From 0a00a978e8e4e9f99dd2dd5e0df67c3c2a003c40 Mon Sep 17 00:00:00 2001
From: Jona Reynaert <jona.reynaert@ugent.be>
Date: Mon, 4 May 2026 00:00:05 +0200
Subject: [PATCH] fix: formatted file

---
 .../trainers/PPOTrainer.py                    | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py
index 6f5a301..9f22ae3 100644
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@@ -240,15 +240,18 @@ def _rollout_jit(
     truncated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_)
 
     (
-        agent_state,
-        episode_stats,
-        next_obs,
-        next_done,
-        key,
-        env_state,
-        terminated_any,
-        truncated_any,
-    ), storage = jax.lax.scan(
+        (
+            agent_state,
+            episode_stats,
+            next_obs,
+            next_done,
+            key,
+            env_state,
+            terminated_any,
+            truncated_any,
+        ),
+        storage,
+    ) = jax.lax.scan(
         partial(
             _step_once,
             sensor=sensor,