53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
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)
|