Deep CFR LCFR DCFR loss weighting 추가

This commit is contained in:
2026-05-07 06:30:48 +09:00
parent e1a5185d4d
commit 4712a9af89
4 changed files with 137 additions and 6 deletions
@@ -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(