small improvements...

kscalelabs · Oct 16, 2024 · c5cac35 · c5cac35
1 parent 569a3ee
commit c5cac35
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -165,7 +165,8 @@ def _env_step(
                 log_prob = pi.log_prob(action)
 
                 # Updates the environment state using the simulation.
-                rng, *step_rngs = jax.random.split(rng, config.training.num_envs + 1)
+                rng, step_rng = jax.random.split(rng)
+                step_rngs = jax.random.split(step_rng, config.training.num_envs)
                 env_state: State = step_fn(env_state, action, step_rngs)
 
                 # Normalizing observations improves training