coolguyandClaude Haiku 4.5
cee8849e04
- 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>
311 lines
11 KiB
Python
311 lines
11 KiB
Python
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)
|