Add donate_argnums to brax PPO. Avoids graph recaptures with MJX-Warp.

PiperOrigin-RevId: 816256474
Change-Id: I2c5573ca8c3ee7a279fbe160bbf30a62cfe7ff88

Author: btaba

brax/training/acting.py

@@ -61,7 +61,6 @@ def generate_unroll(
 ) -> Tuple[State, Transition]:
   """Collect trajectories of given unroll_length."""
 
-  @jax.jit
   def f(carry, unused_t):
     state, current_key = carry
     current_key, next_key = jax.random.split(current_key)
@@ -70,8 +69,9 @@ def f(carry, unused_t):
     )
     return (nstate, next_key), transition
 
+  f_jit = jax.jit(f, donate_argnums=(0,))
   (final_state, _), data = jax.lax.scan(
-      f, (env_state, key), (), length=unroll_length
+      f_jit, (env_state, key), (), length=unroll_length
   )
   return final_state, data
 
@@ -111,9 +111,12 @@ def __init__(
     self._eval_walltime = 0.0
 
     eval_env = envs.training.EvalWrapper(eval_env)
+    self._eval_state_to_donate = jax.jit(eval_env.reset)(
+        jax.random.split(key, num_eval_envs)
+    )
 
     def generate_eval_unroll(
-        policy_params: PolicyParams, key: PRNGKey
+        eval_env_state_donated: State, policy_params: PolicyParams, key: PRNGKey
     ) -> State:
       reset_keys = jax.random.split(key, num_eval_envs)
       eval_first_state = eval_env.reset(reset_keys)
@@ -125,7 +128,9 @@ def generate_eval_unroll(
           unroll_length=episode_length // action_repeat,
       )[0]
 
-    self._generate_eval_unroll = jax.jit(generate_eval_unroll)
+    self._generate_eval_unroll = jax.jit(
+        generate_eval_unroll, donate_argnums=(0,), keep_unused=True
+    )
     self._steps_per_unroll = episode_length * num_eval_envs
 
   def run_evaluation(
@@ -138,7 +143,11 @@ def run_evaluation(
     self._key, unroll_key = jax.random.split(self._key)
 
     t = time.time()
-    eval_state = self._generate_eval_unroll(policy_params, unroll_key)
+    eval_state = self._generate_eval_unroll(
+        self._eval_state_to_donate, policy_params, unroll_key
+    )
+    self._eval_state_to_donate = eval_state
+
     eval_metrics = eval_state.info['eval_metrics']
     eval_metrics.active_episodes.block_until_ready()
     epoch_eval_time = time.time() - t

brax/training/agents/ppo/train.py

@@ -394,16 +394,28 @@ def train(
       randomization_fn,
   )
 
-  if local_devices_to_use > 1 or use_pmap_on_reset:
-    reset_fn = jax.pmap(env.reset, axis_name=_PMAP_AXIS_NAME)
-  else:
-    reset_fn = jax.jit(jax.vmap(env.reset))
+  def reset_fn_donated_env_state(env_state_donated, key_envs):
+    return env.reset(key_envs)
 
   key_envs = jax.random.split(key_env, num_envs // process_count)
   key_envs = jnp.reshape(
       key_envs, (local_devices_to_use, -1) + key_envs.shape[1:]
   )
-  env_state = reset_fn(key_envs)
+  if local_devices_to_use > 1 or use_pmap_on_reset:
+    reset_fn_ = jax.pmap(env.reset, axis_name=_PMAP_AXIS_NAME)
+    env_state = reset_fn_(key_envs)
+    reset_fn = jax.pmap(
+        reset_fn_donated_env_state,
+        axis_name=_PMAP_AXIS_NAME,
+        donate_argnums=(0,),
+    )
+  else:
+    reset_fn_ = jax.jit(jax.vmap(env.reset))
+    env_state = reset_fn_(key_envs)
+    reset_fn = jax.jit(
+        reset_fn_donated_env_state, donate_argnums=(0,), keep_unused=True
+    )
+
   # Discard the batch axes over devices and envs.
   obs_shape = jax.tree_util.tree_map(lambda x: x.shape[2:], env_state.obs)
 
@@ -611,7 +623,14 @@ def training_epoch(
     loss_metrics = jax.tree_util.tree_map(jnp.mean, loss_metrics)
     return training_state, state, loss_metrics
 
-  training_epoch = jax.pmap(training_epoch, axis_name=_PMAP_AXIS_NAME)
+  training_epoch = jax.pmap(
+      training_epoch,
+      axis_name=_PMAP_AXIS_NAME,
+      donate_argnums=(
+          0,
+          1,
+      ),
+  )
 
   # Note that this is NOT a pure jittable method.
   def training_epoch_with_timing(
@@ -755,7 +774,8 @@ def training_epoch_with_timing(
           lambda x, s: jax.random.split(x[0], s), in_axes=(0, None)
       )(key_envs, key_envs.shape[1])
       # TODO(brax-team): move extra reset logic to the AutoResetWrapper.
-      env_state = reset_fn(key_envs) if num_resets_per_eval > 0 else env_state
+      if num_resets_per_eval > 0:
+        env_state = reset_fn((training_state, env_state), key_envs)
 
     if process_id != 0:
       continue

brax/training/agents/ppo/train_test.py

@@ -54,8 +54,8 @@ def testTrain(self, obs_mode):
         normalize_advantage=False,
     )
     self.assertGreater(metrics['eval/episode_reward'], 135)
-    self.assertEqual(fast.reset_count, 2)  # type: ignore
-    self.assertEqual(fast.step_count, 2)  # type: ignore
+    self.assertEqual(fast.reset_count, 4)  # type: ignore
+    self.assertEqual(fast.step_count, 3)  # type: ignore
 
   @parameterized.parameters(
       ('normal', 'scalar'),

brax/training/agents/sac/train.py

@@ -416,7 +416,8 @@ def f(carry, unused):
     )[0]
 
   prefill_replay_buffer = jax.pmap(
-      prefill_replay_buffer, axis_name=_PMAP_AXIS_NAME
+      prefill_replay_buffer, axis_name=_PMAP_AXIS_NAME,
+      donate_argnums=(0, 1, 2)
   )
 
   def training_epoch(
@@ -441,7 +442,9 @@ def f(carry, unused_t):
     metrics = jax.tree_util.tree_map(jnp.mean, metrics)
     return training_state, env_state, buffer_state, metrics
 
-  training_epoch = jax.pmap(training_epoch, axis_name=_PMAP_AXIS_NAME)
+  training_epoch = jax.pmap(
+      training_epoch, axis_name=_PMAP_AXIS_NAME, donate_argnums=(0, 1, 2)
+  )
 
   # Note that this is NOT a pure jittable method.
   def training_epoch_with_timing(

brax/training/agents/sac/train_test.py

@@ -45,11 +45,12 @@ def testTrain(self):
         grad_updates_per_step=64,
         num_evals=3,
         seed=0,
+        eval_env=envs.get_environment('fast'),
     )
     self.assertGreater(metrics['eval/episode_reward'], 140 * 0.995)
-    self.assertEqual(fast.reset_count, 3)  # type: ignore
-    # once for prefill, once for train, once for eval
-    self.assertEqual(fast.step_count, 3)  # type: ignore
+    self.assertEqual(fast.reset_count, 2)  # type: ignore
+    # once for prefill, once for train
+    self.assertEqual(fast.step_count, 2)  # type: ignore
 
   @parameterized.parameters(True, False)
   def testNetworkEncoding(self, normalize_observations):

docs/release-notes/next-release.md

@@ -6,3 +6,4 @@
 * Allow episode metrics during eval to be normalized by the episode length, as long as the metric name ends with "per_step".
 * Add adaptive learning rate to PPO. Desired KL is sensitive to network initialization weights and entropy cost and may require some tuning for your environment.
 * Add loss metrics to the PPO training logger.
+* Add `donate_argnums` to brax PPO to somewhat mitigate repeated graph captures when using MJX-Warp.