Reset finished JAX PPO environments

src/lost_cities_jax/ppo.py

@@ -238,7 +238,7 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
     def train_iteration(
         state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array
     ):
-        rng, rollout_key, update_key = jax.random.split(rng, 3)
+        rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4)
         env_state, transitions, rollout_metrics = rollout_fn(
             state, env_state, rollout_key, shaping_coef
         )
@@ -250,6 +250,7 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
             cfg.ppo.gae_lambda,
         )
         state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
+        env_state = reset_done_envs(env_state, reset_key, cfg.ppo.batch_games)
         return state, env_state, rng, {**rollout_metrics, **update_metrics}
 
     return train_iteration
@@ -374,6 +375,14 @@ def random_rollout(cfg: JaxPPOConfig) -> dict:
     return _metrics_to_row(metrics, update=0, shaping_coef=0.0, cfg=cfg)
 
 
+def reset_done_envs(env_state: State, rng: jax.Array, batch_games: int) -> State:
+    fresh = jax.vmap(reset)(jax.random.split(rng, batch_games))
+    done = env_state.done
+    return jax.tree_util.tree_map(
+        lambda old, new: jnp.where(_broadcast_done(done, old), new, old), env_state, fresh
+    )
+
+
 def ppo_update(
     state: TrainState,
     transitions: Transition,
@@ -644,6 +653,11 @@ def masked_mean(values: jax.Array, weights: jax.Array) -> jax.Array:
     return jnp.sum(values * weights) / jnp.maximum(jnp.sum(weights), 1.0)
 
 
+def _broadcast_done(done: jax.Array, leaf: jax.Array) -> jax.Array:
+    shape = done.shape + (1,) * (leaf.ndim - 1)
+    return done.reshape(shape)
+
+
 def _flatten_transitions(transitions: Transition) -> Transition:
     return jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)