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 <noreply@anthropic.com>

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

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] = []
+        self.net = _build_mlp(input_dim, output_dim, hidden_size, num_layers, activation)
-        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)
 
     @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

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)