diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py index 3a9e01f..527a37a 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py @@ -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") diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py index 957e9f2..faaeee9 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -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) diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 9c32020..9c650da 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -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) - loss = diff.square().sum() / legal.sum().clamp_min(1) + 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: diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index aec5f64..ea42b81 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -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(