Deep CFR LCFR DCFR loss weighting 추가
This commit is contained in:
@@ -88,6 +88,8 @@ def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]:
|
||||
overrides.setdefault("evaluation", {})["games"] = args.eval_games
|
||||
if args.regret_fallback is not None:
|
||||
overrides.setdefault("regret_matching", {})["all_negative_fallback"] = args.regret_fallback
|
||||
if args.training_weighting is not None:
|
||||
overrides.setdefault("training_weighting", {})["mode"] = args.training_weighting
|
||||
if args.no_save:
|
||||
checkpoint_overrides = overrides.setdefault("checkpoint", {})
|
||||
checkpoint_overrides["save_latest"] = False
|
||||
@@ -225,6 +227,11 @@ def main(argv: list[str] | None = None) -> None:
|
||||
choices=("uniform", "argmax_tiebreak"),
|
||||
help="Override regret_matching.all_negative_fallback.",
|
||||
)
|
||||
train.add_argument(
|
||||
"--training-weighting",
|
||||
choices=("none", "lcfr", "dcfr"),
|
||||
help="Override training_weighting.mode.",
|
||||
)
|
||||
train.add_argument("--seed", type=int)
|
||||
train.add_argument("--no-save", action="store_true")
|
||||
train.add_argument("--save-latest-only", action="store_true")
|
||||
|
||||
@@ -175,6 +175,22 @@ class RegretMatchingConfig(StrictModel):
|
||||
return token
|
||||
|
||||
|
||||
class TrainingWeightingConfig(StrictModel):
|
||||
mode: str = "none"
|
||||
lcfr_alpha: float = 1.0
|
||||
dcfr_alpha: float = 1.5
|
||||
dcfr_beta: float = 0.0
|
||||
dcfr_gamma: float = 2.0
|
||||
|
||||
@field_validator("mode")
|
||||
@classmethod
|
||||
def _validate_mode(cls, value: str) -> str:
|
||||
token = value.strip().lower()
|
||||
if token not in {"none", "lcfr", "dcfr"}:
|
||||
raise ValueError("must be 'none', 'lcfr', or 'dcfr'")
|
||||
return token
|
||||
|
||||
|
||||
class SelfPlayLeagueConfig(StrictModel):
|
||||
snapshot_every: int = 1
|
||||
max_snapshots: int = 20
|
||||
@@ -281,6 +297,7 @@ class DeepCFRConfig(StrictModel):
|
||||
network: NetworkConfig = Field(default_factory=NetworkConfig)
|
||||
traversal: TraversalConfig = Field(default_factory=TraversalConfig)
|
||||
regret_matching: RegretMatchingConfig = Field(default_factory=RegretMatchingConfig)
|
||||
training_weighting: TrainingWeightingConfig = Field(default_factory=TrainingWeightingConfig)
|
||||
self_play: SelfPlayLeagueConfig = Field(default_factory=SelfPlayLeagueConfig)
|
||||
optimization: OptimizationConfig = Field(default_factory=OptimizationConfig)
|
||||
memory: MemoryConfig = Field(default_factory=MemoryConfig)
|
||||
|
||||
@@ -691,7 +691,7 @@ class DeepCFRTrainer:
|
||||
def _batch_tensors(
|
||||
self,
|
||||
batch: list[TrainingSample],
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
started = time.perf_counter()
|
||||
x = torch.as_tensor(
|
||||
np.stack([sample.info_state for sample in batch]),
|
||||
@@ -708,12 +708,24 @@ class DeepCFRTrainer:
|
||||
dtype=torch.bool,
|
||||
device=self.device,
|
||||
)
|
||||
iterations = torch.as_tensor(
|
||||
[sample.iteration for sample in batch],
|
||||
dtype=torch.float32,
|
||||
device=self.device,
|
||||
)
|
||||
self._runtime_metrics["batch_tensor_seconds"] = (
|
||||
float(self._runtime_metrics.get("batch_tensor_seconds", 0.0))
|
||||
+ time.perf_counter()
|
||||
- started
|
||||
)
|
||||
return x, y, legal
|
||||
return x, y, legal, iterations
|
||||
|
||||
def _iteration_weights(self, iterations: torch.Tensor, exponent: float) -> torch.Tensor:
|
||||
if exponent == 0.0:
|
||||
return torch.ones_like(iterations, dtype=torch.float32, device=self.device)
|
||||
current = max(1, int(self.iteration))
|
||||
relative = (iterations / float(current)).clamp(min=0.0, max=1.0)
|
||||
return relative.pow(float(exponent))
|
||||
|
||||
def _train_advantage(
|
||||
self,
|
||||
@@ -734,10 +746,33 @@ class DeepCFRTrainer:
|
||||
+ time.perf_counter()
|
||||
- sample_started
|
||||
)
|
||||
x, y, legal = self._batch_tensors(batch)
|
||||
x, y, legal, sample_iterations = self._batch_tensors(batch)
|
||||
pred = network(x)
|
||||
diff = (pred - y).masked_fill(~legal, 0.0)
|
||||
if self.config.training_weighting.mode == "none":
|
||||
loss = diff.square().sum() / legal.sum().clamp_min(1)
|
||||
elif self.config.training_weighting.mode == "lcfr":
|
||||
sample_weights = self._iteration_weights(
|
||||
sample_iterations, self.config.training_weighting.lcfr_alpha
|
||||
)
|
||||
action_weights = sample_weights[:, None] * legal.float()
|
||||
loss = (diff.square() * action_weights).sum() / action_weights.sum().clamp_min(
|
||||
1.0e-12
|
||||
)
|
||||
else:
|
||||
positive_weights = self._iteration_weights(
|
||||
sample_iterations, self.config.training_weighting.dcfr_alpha
|
||||
)
|
||||
negative_weights = self._iteration_weights(
|
||||
sample_iterations, self.config.training_weighting.dcfr_beta
|
||||
)
|
||||
target_weights = torch.where(
|
||||
y >= 0.0, positive_weights[:, None], negative_weights[:, None]
|
||||
)
|
||||
action_weights = target_weights * legal.float()
|
||||
loss = (diff.square() * action_weights).sum() / action_weights.sum().clamp_min(
|
||||
1.0e-12
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
if self.config.optimization.grad_clip > 0.0:
|
||||
@@ -765,10 +800,26 @@ class DeepCFRTrainer:
|
||||
+ time.perf_counter()
|
||||
- sample_started
|
||||
)
|
||||
x, y, legal = self._batch_tensors(batch)
|
||||
x, y, legal, sample_iterations = self._batch_tensors(batch)
|
||||
logits = network(x).masked_fill(~legal, torch.finfo(torch.float32).min)
|
||||
log_probs = nn.functional.log_softmax(logits, dim=-1).masked_fill(~legal, 0.0)
|
||||
loss = -(y * log_probs).sum(dim=-1).mean()
|
||||
per_sample_loss = -(y * log_probs).sum(dim=-1)
|
||||
if self.config.training_weighting.mode == "none":
|
||||
loss = per_sample_loss.mean()
|
||||
elif self.config.training_weighting.mode == "lcfr":
|
||||
sample_weights = self._iteration_weights(
|
||||
sample_iterations, self.config.training_weighting.lcfr_alpha
|
||||
)
|
||||
loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum().clamp_min(
|
||||
1.0e-12
|
||||
)
|
||||
else:
|
||||
sample_weights = self._iteration_weights(
|
||||
sample_iterations, self.config.training_weighting.dcfr_gamma
|
||||
)
|
||||
loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum().clamp_min(
|
||||
1.0e-12
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
if self.config.optimization.grad_clip > 0.0:
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim
|
||||
from coolrl_lost_cities.games.classic.deep_cfr.traversal import CythonDeepCFRTraverser
|
||||
from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
|
||||
@@ -65,6 +66,7 @@ def test_deep_cfr_loads_mapped_legacy_reproduction_config() -> None:
|
||||
assert config.evaluation.device == "trainer"
|
||||
assert config.evaluation.resolved_num_workers() == 4
|
||||
assert config.regret_matching.all_negative_fallback == "uniform"
|
||||
assert config.training_weighting.mode == "none"
|
||||
assert config.checkpoint.save_iteration_interval == 10
|
||||
assert (
|
||||
config.checkpoint.directory == "runs/deep_cfr/deep_cfr_selfplay_full_depth_slot_playability"
|
||||
@@ -86,6 +88,7 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None:
|
||||
"eval_every": None,
|
||||
"eval_games": None,
|
||||
"regret_fallback": "argmax_tiebreak",
|
||||
"training_weighting": "lcfr",
|
||||
"no_save": True,
|
||||
"save_latest_only": False,
|
||||
"save_iteration_interval": None,
|
||||
@@ -103,6 +106,7 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None:
|
||||
assert overridden.traversal.resolved_traversals_per_player() == 1
|
||||
assert overridden.traversal.resolved_num_workers() == 0
|
||||
assert overridden.regret_matching.all_negative_fallback == "argmax_tiebreak"
|
||||
assert overridden.training_weighting.mode == "lcfr"
|
||||
assert overridden.checkpoint.save_every_iteration is False
|
||||
assert overridden.checkpoint.save_latest is False
|
||||
|
||||
@@ -122,6 +126,7 @@ def test_deep_cfr_train_cli_checkpoint_save_overrides() -> None:
|
||||
"eval_every": None,
|
||||
"eval_games": None,
|
||||
"regret_fallback": None,
|
||||
"training_weighting": None,
|
||||
"no_save": False,
|
||||
"save_latest_only": True,
|
||||
"save_iteration_interval": 1,
|
||||
@@ -137,6 +142,26 @@ def test_deep_cfr_train_cli_checkpoint_save_overrides() -> None:
|
||||
assert overridden.checkpoint.save_iteration_interval == 1
|
||||
|
||||
|
||||
def test_deep_cfr_iteration_weights_use_sample_age() -> None:
|
||||
trainer = DeepCFRTrainer(
|
||||
_deep_cfr_config(
|
||||
{
|
||||
"run": {"iterations": 1, "seed": 12},
|
||||
"network": {"hidden_size": 16},
|
||||
"checkpoint": {"save_every_iteration": False},
|
||||
"training_weighting": {"mode": "lcfr", "lcfr_alpha": 1.0},
|
||||
}
|
||||
),
|
||||
LostCitiesConfig(seed=12),
|
||||
)
|
||||
trainer.iteration = 10
|
||||
|
||||
weights = trainer._iteration_weights(torch.tensor([1.0, 5.0, 10.0], device=trainer.device), 1.0)
|
||||
|
||||
assert np.allclose(weights.detach().cpu().numpy(), np.asarray([0.1, 0.5, 1.0]))
|
||||
assert trainer.config.training_weighting.mode == "lcfr"
|
||||
|
||||
|
||||
def test_deep_cfr_batched_evaluation_matches_batch_size_one() -> None:
|
||||
config = _deep_cfr_config(
|
||||
{
|
||||
@@ -293,6 +318,37 @@ def test_deep_cfr_trainer_smoke_run() -> None:
|
||||
assert metrics[0].strategy_loss >= 0.0
|
||||
|
||||
|
||||
def test_deep_cfr_trainer_supports_lcfr_and_dcfr_loss_weighting() -> None:
|
||||
for mode in ("lcfr", "dcfr"):
|
||||
trainer = DeepCFRTrainer(
|
||||
_deep_cfr_config(
|
||||
{
|
||||
"run": {"iterations": 1, "seed": 24},
|
||||
"network": {"hidden_size": 16},
|
||||
"traversal": {
|
||||
"traversals_per_iteration": 1,
|
||||
"max_depth": 2,
|
||||
"max_nodes": 32,
|
||||
},
|
||||
"optimization": {
|
||||
"advantage_train_steps": 1,
|
||||
"strategy_train_steps": 1,
|
||||
"batch_size": 2,
|
||||
},
|
||||
"training_weighting": {"mode": mode},
|
||||
"checkpoint": {"save_every_iteration": False},
|
||||
}
|
||||
),
|
||||
LostCitiesConfig(seed=24),
|
||||
)
|
||||
|
||||
metrics = trainer.train()
|
||||
|
||||
assert len(metrics) == 1
|
||||
assert metrics[0].advantage_loss >= 0.0
|
||||
assert metrics[0].strategy_loss >= 0.0
|
||||
|
||||
|
||||
def test_deep_cfr_cython_traverser_restores_state_and_collects_samples() -> None:
|
||||
trainer = DeepCFRTrainer(
|
||||
_deep_cfr_config(
|
||||
|
||||
Reference in New Issue
Block a user