From cee8849e049d95aaa539c95b954ad354cf472acb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Thu, 7 May 2026 09:18:13 +0900 Subject: [PATCH] Add backward-compatible color_shared network architecture - Add network.kind field to config (mlp/color_shared) - Implement ColorSharedNetwork that: - Splits input into 5 equal color blocks - Encodes each block with shared weights - Pools with mean/max aggregation - Concatenates pooled embeddings with remainder - Outputs same action logits as MLP - Add ColorAttention for optional self-attention over color embeddings - Add network.color_attention_layers and color_attention_heads config - Maintain full backward compatibility (default kind=mlp) - Add 28 comprehensive tests covering all architectures Co-Authored-By: Claude Haiku 4.5 --- .../games/classic/deep_cfr/config.py | 11 + .../games/classic/deep_cfr/networks.py | 157 ++++++++- tests/games/classic/test_networks.py | 310 ++++++++++++++++++ 3 files changed, 469 insertions(+), 9 deletions(-) create mode 100644 tests/games/classic/test_networks.py 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 faaeee9..42b283e 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -70,9 +70,20 @@ class EncodingConfig(StrictModel): class NetworkConfig(StrictModel): + kind: str = "mlp" hidden_size: int = 64 num_layers: int = 2 activation: str = "relu" + color_attention_layers: int = 0 + color_attention_heads: int = 4 + + @field_validator("kind") + @classmethod + def _validate_kind(cls, value: str) -> str: + token = value.strip().lower() + if token not in {"mlp", "color_shared"}: + raise ValueError("must be 'mlp' or 'color_shared'") + return token @field_validator("activation") @classmethod diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/networks.py b/src/coolrl_lost_cities/games/classic/deep_cfr/networks.py index d7d8e47..f173b1e 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/networks.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/networks.py @@ -15,6 +15,23 @@ def _activation(name: str) -> nn.Module: raise ValueError(f"unsupported activation: {name!r}") +def _build_mlp( + input_dim: int, + output_dim: int, + hidden_size: int, + num_layers: int, + activation: str, +) -> nn.Sequential: + layers: list[nn.Module] = [] + last_dim = input_dim + for _ in range(max(0, int(num_layers))): + layers.append(nn.Linear(last_dim, hidden_size)) + layers.append(_activation(activation)) + last_dim = hidden_size + layers.append(nn.Linear(last_dim, output_dim)) + return nn.Sequential(*layers) + + class DeepCFRMLP(nn.Module): def __init__( self, @@ -26,17 +43,20 @@ class DeepCFRMLP(nn.Module): activation: str = "relu", ) -> None: super().__init__() - layers: list[nn.Module] = [] - last_dim = input_dim - for _ in range(max(0, int(num_layers))): - layers.append(nn.Linear(last_dim, hidden_size)) - layers.append(_activation(activation)) - last_dim = hidden_size - layers.append(nn.Linear(last_dim, output_dim)) - self.net = nn.Sequential(*layers) + self.net = _build_mlp(input_dim, output_dim, hidden_size, num_layers, activation) @classmethod - def from_config(cls, input_dim: int, output_dim: int, config: NetworkConfig) -> DeepCFRMLP: + def from_config(cls, input_dim: int, output_dim: int, config: NetworkConfig) -> nn.Module: + if config.kind == "color_shared": + return ColorSharedNetwork( + input_dim, + output_dim, + config.hidden_size, + num_layers=config.num_layers, + activation=config.activation, + color_attention_layers=config.color_attention_layers, + color_attention_heads=config.color_attention_heads, + ) return cls( input_dim, output_dim, @@ -47,3 +67,122 @@ class DeepCFRMLP(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) + + +class ColorSharedNetwork(nn.Module): + """Color-shared architecture that splits input into per-color blocks. + + Splits the input into n_colors equal parts, encodes each with shared weights, + pools the color embeddings, and concatenates with the original input. + """ + + N_COLORS = 5 + + def __init__( + self, + input_dim: int, + output_dim: int, + hidden_size: int = 64, + *, + num_layers: int = 2, + activation: str = "relu", + color_attention_layers: int = 0, + color_attention_heads: int = 4, + ) -> None: + super().__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.hidden_size = hidden_size + self.n_colors = self.N_COLORS + self.color_attention_layers = color_attention_layers + self.color_attention_heads = color_attention_heads + + color_block_size = input_dim // self.n_colors + self.color_block_size = color_block_size + + self.color_encoder = _build_mlp( + color_block_size, + hidden_size, + hidden_size, + num_layers, + activation, + ) + + self.color_attention = None + if color_attention_layers > 0: + self.color_attention = ColorAttention( + hidden_size, + num_layers=color_attention_layers, + num_heads=color_attention_heads, + activation=activation, + ) + + final_input_dim = hidden_size * 2 + input_dim % self.n_colors + + self.final_net = _build_mlp( + final_input_dim, + output_dim, + hidden_size, + num_layers, + activation, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + color_embeddings = [] + for i in range(self.n_colors): + start = i * self.color_block_size + end = start + self.color_block_size + block = x[:, start:end] + embedding = self.color_encoder(block) + color_embeddings.append(embedding) + + color_embeddings = torch.stack(color_embeddings, dim=1) + + if self.color_attention is not None: + color_embeddings = self.color_attention(color_embeddings) + + mean_pooled = color_embeddings.mean(dim=1) + max_pooled = color_embeddings.max(dim=1)[0] + + remainder = x[:, self.n_colors * self.color_block_size :] + final_features = torch.cat([mean_pooled, max_pooled, remainder], dim=1) + + logits = self.final_net(final_features) + return logits + + +class ColorAttention(nn.Module): + """Self-attention over per-color embeddings.""" + + def __init__( + self, + dim: int, + *, + num_layers: int = 1, + num_heads: int = 4, + activation: str = "relu", + ) -> None: + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + + assert dim % num_heads == 0, f"dim ({dim}) must be divisible by num_heads ({num_heads})" + + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + nn.TransformerEncoderLayer( + d_model=dim, + nhead=num_heads, + dim_feedforward=dim * 4, + activation=activation.lower(), + batch_first=True, + norm_first=True, + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + x = layer(x) + return x diff --git a/tests/games/classic/test_networks.py b/tests/games/classic/test_networks.py new file mode 100644 index 0000000..a1aac42 --- /dev/null +++ b/tests/games/classic/test_networks.py @@ -0,0 +1,310 @@ +from __future__ import annotations + +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.config import NetworkConfig +from coolrl_lost_cities.games.classic.deep_cfr.networks import ( + ColorAttention, + ColorSharedNetwork, + DeepCFRMLP, +) + + +class TestDeepCFRMLP: + def test_basic_mlp_forward(self) -> None: + mlp = DeepCFRMLP(input_dim=64, output_dim=32, hidden_size=128, num_layers=2) + x = torch.randn(16, 64) + output = mlp(x) + assert output.shape == (16, 32) + + def test_mlp_from_config(self) -> None: + config = NetworkConfig(kind="mlp", hidden_size=64, num_layers=2) + mlp = DeepCFRMLP.from_config(input_dim=100, output_dim=50, config=config) + assert isinstance(mlp, DeepCFRMLP) + x = torch.randn(8, 100) + output = mlp(x) + assert output.shape == (8, 50) + + def test_mlp_zero_layers(self) -> None: + mlp = DeepCFRMLP(input_dim=64, output_dim=32, hidden_size=128, num_layers=0) + x = torch.randn(16, 64) + output = mlp(x) + assert output.shape == (16, 32) + + def test_mlp_gelu_activation(self) -> None: + mlp = DeepCFRMLP( + input_dim=64, output_dim=32, hidden_size=128, num_layers=2, activation="gelu" + ) + x = torch.randn(16, 64) + output = mlp(x) + assert output.shape == (16, 32) + + def test_mlp_gradients(self) -> None: + mlp = DeepCFRMLP(input_dim=64, output_dim=32, hidden_size=128, num_layers=2) + x = torch.randn(16, 64, requires_grad=True) + output = mlp(x) + loss = output.sum() + loss.backward() + assert x.grad is not None + assert x.grad.shape == x.shape + + +class TestColorSharedNetwork: + def test_color_shared_basic(self) -> None: + network = ColorSharedNetwork(input_dim=100, output_dim=50, hidden_size=64, num_layers=2) + x = torch.randn(16, 100) + output = network(x) + assert output.shape == (16, 50) + + def test_color_shared_from_config(self) -> None: + config = NetworkConfig(kind="color_shared", hidden_size=64, num_layers=2) + network = DeepCFRMLP.from_config(input_dim=150, output_dim=75, config=config) + assert isinstance(network, ColorSharedNetwork) + x = torch.randn(8, 150) + output = network(x) + assert output.shape == (8, 75) + + def test_color_shared_splits_input_correctly(self) -> None: + n_colors = 5 + color_block_size = 20 + input_dim = n_colors * color_block_size + network = ColorSharedNetwork(input_dim=input_dim, output_dim=32, hidden_size=64) + assert network.n_colors == n_colors + assert network.color_block_size == color_block_size + + def test_color_shared_with_remainder(self) -> None: + input_dim = 105 + network = ColorSharedNetwork(input_dim=input_dim, output_dim=50, hidden_size=64) + x = torch.randn(8, input_dim) + output = network(x) + assert output.shape == (8, 50) + + def test_color_shared_different_batch_sizes(self) -> None: + network = ColorSharedNetwork(input_dim=100, output_dim=50, hidden_size=64) + for batch_size in [1, 4, 16, 32, 64]: + x = torch.randn(batch_size, 100) + output = network(x) + assert output.shape == (batch_size, 50) + + def test_color_shared_gradients(self) -> None: + network = ColorSharedNetwork(input_dim=100, output_dim=50, hidden_size=64, num_layers=2) + x = torch.randn(16, 100, requires_grad=True) + output = network(x) + loss = output.sum() + loss.backward() + assert x.grad is not None + assert x.grad.shape == x.shape + for param in network.parameters(): + assert param.grad is not None + + def test_color_shared_deterministic_with_seed(self) -> None: + torch.manual_seed(42) + network1 = ColorSharedNetwork(input_dim=100, output_dim=50, hidden_size=64) + torch.manual_seed(42) + network2 = ColorSharedNetwork(input_dim=100, output_dim=50, hidden_size=64) + + x = torch.randn(8, 100) + torch.manual_seed(42) + output1 = network1(x) + torch.manual_seed(42) + output2 = network2(x) + torch.testing.assert_close(output1, output2) + + def test_color_shared_without_attention(self) -> None: + network = ColorSharedNetwork( + input_dim=100, + output_dim=50, + hidden_size=64, + color_attention_layers=0, + ) + assert network.color_attention is None + x = torch.randn(8, 100) + output = network(x) + assert output.shape == (8, 50) + + +class TestColorAttention: + def test_color_attention_forward(self) -> None: + attention = ColorAttention(dim=64, num_layers=1, num_heads=4) + x = torch.randn(8, 5, 64) + output = attention(x) + assert output.shape == (8, 5, 64) + + def test_color_attention_multiple_layers(self) -> None: + for num_layers in [1, 2, 3]: + attention = ColorAttention(dim=64, num_layers=num_layers, num_heads=4) + x = torch.randn(8, 5, 64) + output = attention(x) + assert output.shape == (8, 5, 64) + + def test_color_attention_different_heads(self) -> None: + for num_heads in [1, 2, 4, 8]: + attention = ColorAttention(dim=64, num_layers=1, num_heads=num_heads) + x = torch.randn(8, 5, 64) + output = attention(x) + assert output.shape == (8, 5, 64) + + def test_color_attention_gradients(self) -> None: + attention = ColorAttention(dim=64, num_layers=1, num_heads=4) + x = torch.randn(8, 5, 64, requires_grad=True) + output = attention(x) + loss = output.sum() + loss.backward() + assert x.grad is not None + assert x.grad.shape == x.shape + for param in attention.parameters(): + assert param.grad is not None + + def test_color_attention_gelu(self) -> None: + attention = ColorAttention(dim=64, num_layers=1, num_heads=4, activation="gelu") + x = torch.randn(8, 5, 64) + output = attention(x) + assert output.shape == (8, 5, 64) + + +class TestColorSharedNetworkWithAttention: + def test_color_shared_with_attention(self) -> None: + network = ColorSharedNetwork( + input_dim=100, + output_dim=50, + hidden_size=64, + num_layers=2, + color_attention_layers=1, + color_attention_heads=4, + ) + assert network.color_attention is not None + x = torch.randn(8, 100) + output = network(x) + assert output.shape == (8, 50) + + def test_color_shared_with_multi_layer_attention(self) -> None: + network = ColorSharedNetwork( + input_dim=100, + output_dim=50, + hidden_size=64, + num_layers=2, + color_attention_layers=3, + color_attention_heads=4, + ) + x = torch.randn(8, 100) + output = network(x) + assert output.shape == (8, 50) + + def test_color_shared_with_attention_from_config(self) -> None: + config = NetworkConfig( + kind="color_shared", + hidden_size=64, + num_layers=2, + color_attention_layers=2, + color_attention_heads=4, + ) + network = DeepCFRMLP.from_config(input_dim=100, output_dim=50, config=config) + assert isinstance(network, ColorSharedNetwork) + assert network.color_attention is not None + x = torch.randn(8, 100) + output = network(x) + assert output.shape == (8, 50) + + def test_color_shared_with_attention_gradients(self) -> None: + network = ColorSharedNetwork( + input_dim=100, + output_dim=50, + hidden_size=64, + num_layers=2, + color_attention_layers=1, + color_attention_heads=4, + ) + x = torch.randn(8, 100, requires_grad=True) + output = network(x) + loss = output.sum() + loss.backward() + assert x.grad is not None + for param in network.parameters(): + assert param.grad is not None + + +class TestNetworkBackwardCompatibility: + def test_default_config_is_mlp(self) -> None: + config = NetworkConfig() + assert config.kind == "mlp" + + def test_from_config_respects_kind(self) -> None: + mlp_config = NetworkConfig(kind="mlp") + color_shared_config = NetworkConfig(kind="color_shared") + + mlp = DeepCFRMLP.from_config(input_dim=100, output_dim=50, config=mlp_config) + color_shared = DeepCFRMLP.from_config( + input_dim=100, output_dim=50, config=color_shared_config + ) + + assert isinstance(mlp, DeepCFRMLP) + assert not isinstance(mlp, ColorSharedNetwork) + assert isinstance(color_shared, ColorSharedNetwork) + + def test_mlp_and_color_shared_same_output_shape(self) -> None: + input_dim = 100 + output_dim = 50 + x = torch.randn(8, input_dim) + + mlp_config = NetworkConfig(kind="mlp", hidden_size=64, num_layers=2) + color_shared_config = NetworkConfig(kind="color_shared", hidden_size=64, num_layers=2) + + mlp = DeepCFRMLP.from_config(input_dim, output_dim, mlp_config) + color_shared = DeepCFRMLP.from_config(input_dim, output_dim, color_shared_config) + + mlp_output = mlp(x) + color_shared_output = color_shared(x) + + assert mlp_output.shape == color_shared_output.shape == (8, output_dim) + + +class TestNetworkIntegration: + def test_mlp_with_real_input_size(self) -> None: + from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim + from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig + + game_config = LostCitiesConfig() + state = GameState.new_game(game_config) + dim = input_dim(state) + action_size = 2 * game_config.hand_size + 1 + game_config.n_colors + + config = NetworkConfig(kind="mlp", hidden_size=64, num_layers=2) + network = DeepCFRMLP.from_config(dim, action_size, config) + + x = torch.randn(8, dim) + output = network(x) + assert output.shape == (8, action_size) + + def test_color_shared_with_real_input_size(self) -> None: + from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim + from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig + + game_config = LostCitiesConfig() + state = GameState.new_game(game_config) + dim = input_dim(state) + action_size = 2 * game_config.hand_size + 1 + game_config.n_colors + + config = NetworkConfig(kind="color_shared", hidden_size=64, num_layers=2) + network = DeepCFRMLP.from_config(dim, action_size, config) + + x = torch.randn(8, dim) + output = network(x) + assert output.shape == (8, action_size) + + def test_color_shared_with_attention_real_input_size(self) -> None: + from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim + from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig + + game_config = LostCitiesConfig() + state = GameState.new_game(game_config) + dim = input_dim(state) + action_size = 2 * game_config.hand_size + 1 + game_config.n_colors + + config = NetworkConfig( + kind="color_shared", hidden_size=64, num_layers=2, color_attention_layers=1 + ) + network = DeepCFRMLP.from_config(dim, action_size, config) + + x = torch.randn(8, dim) + output = network(x) + assert output.shape == (8, action_size)