Merge pull request #645 from Michael-Beukman:patch-1

PiperOrigin-RevId: 826006956
Change-Id: I512d2f19a9b7a69c58209fbbe2c5cd0a124c59e2

Author: copybara-github

brax/training/agents/ppo/train.py

@@ -775,7 +775,7 @@ def training_epoch_with_timing(
       )(key_envs, key_envs.shape[1])
       # TODO(brax-team): move extra reset logic to the AutoResetWrapper.
       if num_resets_per_eval > 0:
-        env_state = reset_fn((training_state, env_state), key_envs)
+        env_state = reset_fn(env_state, key_envs)
 
     if process_id != 0:
       continue