Fix reset function call in PPO training

Author: Michael-Beukman

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