checkpoint

tests/test_training.py

@@ -0,0 +1,53 @@
+import jax.numpy as jnp
+import jax.random as jrandom
+
+from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
+from general_bots_training.ppo import compute_advantages_and_returns
+
+
+def test_returns_use_raw_advantages_before_policy_normalization():
+    rewards = jnp.array([[1.0, 3.0]])
+    values = jnp.zeros_like(rewards)
+    next_value = jnp.zeros(2)
+    dones = jnp.ones_like(rewards, dtype=bool)
+
+    advantages, returns = compute_advantages_and_returns(rewards, values, next_value, dones)
+
+    assert jnp.allclose(advantages, jnp.array([[-1.0, 1.0]]))
+    assert jnp.allclose(returns, rewards)
+
+
+def test_pass_is_one_global_action_and_round_trips():
+    network = PolicyValueNetwork(jrandom.PRNGKey(0))
+    obs = jnp.zeros((14, 4, 4))
+    mask = jnp.zeros((4, 4, 4), dtype=bool)
+
+    action, _, sampled_logprob, entropy = network(obs, mask, jrandom.PRNGKey(1))
+    _, _, evaluated_logprob, _ = network(obs, mask, jrandom.PRNGKey(2), action)
+
+    assert action.tolist() == [1, 0, 0, 0, 0]
+    assert jnp.allclose(sampled_logprob, evaluated_logprob)
+    assert jnp.isclose(entropy, 0.0)
+
+
+def test_observation_encoder_includes_normalized_timestep():
+    class Observation:
+        armies = jnp.zeros((2, 3), dtype=jnp.int32)
+        generals = jnp.zeros((2, 3), dtype=bool)
+        castles = jnp.zeros((2, 3), dtype=bool)
+        mountains = jnp.zeros((2, 3), dtype=bool)
+        neutral_cells = jnp.ones((2, 3), dtype=bool)
+        owned_cells = jnp.zeros((2, 3), dtype=bool)
+        opponent_cells = jnp.zeros((2, 3), dtype=bool)
+        fog_cells = jnp.zeros((2, 3), dtype=bool)
+        structures_in_fog = jnp.zeros((2, 3), dtype=bool)
+        owned_land_count = jnp.array(0)
+        owned_army_count = jnp.array(0)
+        opponent_land_count = jnp.array(0)
+        opponent_army_count = jnp.array(0)
+        timestep = jnp.array(1200)
+
+    encoded = obs_to_tensor(Observation())
+
+    assert encoded.shape == (14, 2, 3)
+    assert jnp.allclose(encoded[-1], 1.0)