fix(ppo training): forgot extra change from dev

2026-05-07 08:42:35 +02:00 · 2026-05-07 08:42:35 +02:00 · 23fcee2110
commit 23fcee2110
parent 090e774ad7
1 changed files with 8 additions and 5 deletions
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@ -137,8 +137,11 @@ def _step_once(
    key, reset_key = jax.random.split(key)
    reset_rngs = jax.random.split(reset_key, num_envs)

-    episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
-        episode_stats, env_state, flat_clipped_action, reset_rngs,
+    episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn(
+        episode_stats,
+        env_state,
+        flat_clipped_action,
+        reset_rngs,
    )

    terminated_any = terminated_any | terminated