Add full resets to AutoResetWrapper. #140 #179

PiperOrigin-RevId: 796491534
Change-Id: I187b492c92c2615ea9583587c9b0badf81322c62

Author: btaba

CHANGELOG.md

@@ -13,6 +13,9 @@ All notable changes to this project will be documented in this file.
 - Remove `mjx_env.init` in favor of `mjx_env.make_data` since `make_data`
   now requires an `MjModel` argument rather than an `mjx.Model` argument.
 - Add device to `mjx_env.make_data`, fixes #174.
+- Update AutoResetWrapper to allow full resets on done. Fixes #179. Also
+  provides a means for doing curriculum learning via
+  `state.info['AutoResetWrapper_done_count']`, see #140.
 
 ## [0.0.5] - 2025-06-23
 

mujoco_playground/_src/wrapper.py

@@ -93,18 +93,22 @@ def wrap_for_brax_training(
     randomization_fn: Optional[
         Callable[[mjx.Model], Tuple[mjx.Model, mjx.Model]]
     ] = None,
+    full_reset: bool = False,
 ) -> Wrapper:
   """Common wrapper pattern for all brax training agents.
 
   Args:
     env: environment to be wrapped
     vision: whether the environment will be vision based
-    num_vision_envs: number of environments the renderer should generate,
-      should equal the number of batched envs
+    num_vision_envs: number of environments the renderer should generate, should
+      equal the number of batched envs
     episode_length: length of episode
     action_repeat: how many repeated actions to take per step
     randomization_fn: randomization function that produces a vectorized model
       and in_axes to vmap over
+    full_reset: whether to call `env.reset` during `env.step` on done rather
+      than resetting to a cached first state. Setting full_reset=True may
+      increase wallclock time because it forces full resets to random states.
 
   Returns:
     An environment that is wrapped with Episode and AutoReset wrappers.  If the
@@ -118,24 +122,66 @@ def wrap_for_brax_training(
   else:
     env = BraxDomainRandomizationVmapWrapper(env, randomization_fn)
   env = brax_training.EpisodeWrapper(env, episode_length, action_repeat)
-  env = BraxAutoResetWrapper(env)
+  env = BraxAutoResetWrapper(env, full_reset=full_reset)
   return env
 
 
 class BraxAutoResetWrapper(Wrapper):
-  """Automatically resets Brax envs that are done."""
+  """Automatically resets Brax envs that are done.
+
+  If `full_reset` is disabled (default):
+    * the environment will reset to a cached first state.
+    * only data and obs are reset, not the environment info.
+
+  If `full_reset` is enabled:
+    * the environment will call env.reset during env.step on done.
+    * `full_reset` will thus incur a penalty in wallclock time depending on the
+      complexity of the reset function.
+    * info is fully reset, except for info under the key
+      `AutoResetWrapper_preserve_info`, which is passed through from the prior
+      step. This can be used for curriculum learning.
+
+  Attributes:
+    env: The wrapped environment.
+    full_reset: Whether to call `env.reset` during `env.step` on done.
+  """
+
+  def __init__(self, env: Any, full_reset: bool = False):
+    super().__init__(env)
+    self._full_reset = full_reset
+    self._info_key = 'AutoResetWrapper'
 
   def reset(self, rng: jax.Array) -> mjx_env.State:
-    state = self.env.reset(rng)
-    state.info['first_state'] = state.data
-    state.info['first_obs'] = state.obs
+    rng_key = jax.vmap(jax.random.split)(rng)
+    rng, key = rng_key[..., 0], rng_key[..., 1]
+    state = self.env.reset(key)
+    state.info[f'{self._info_key}_first_data'] = state.data
+    state.info[f'{self._info_key}_first_obs'] = state.obs
+    state.info[f'{self._info_key}_rng'] = rng
+    state.info[f'{self._info_key}_done_count'] = jp.zeros(
+        key.shape[:-1], dtype=int
+    )
     return state
 
   def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
+    # grab the reset state.
+    reset_state = None
+    rng_key = jax.vmap(jax.random.split)(state.info[f'{self._info_key}_rng'])
+    reset_rng, reset_key = rng_key[..., 0], rng_key[..., 1]
+    if self._full_reset:
+      reset_state = self.reset(reset_key)
+      reset_data = reset_state.data
+      reset_obs = reset_state.obs
+    else:
+      reset_data = state.info[f'{self._info_key}_first_data']
+      reset_obs = state.info[f'{self._info_key}_first_obs']
+
     if 'steps' in state.info:
+      # reset steps to 0 if done.
       steps = state.info['steps']
       steps = jp.where(state.done, jp.zeros_like(steps), steps)
       state.info.update(steps=steps)
+
     state = state.replace(done=jp.zeros_like(state.done))
     state = self.env.step(state, action)
 
@@ -147,11 +193,25 @@ def where_done(x, y):
         done = jp.reshape(done, [x.shape[0]] + [1] * (len(x.shape) - 1))
       return jp.where(done, x, y)
 
-    data = jax.tree.map(
-        where_done, state.info['first_state'], state.data
-    )
-    obs = jax.tree.map(where_done, state.info['first_obs'], state.obs)
-    return state.replace(data=data, obs=obs)
+    data = jax.tree.map(where_done, reset_data, state.data)
+    obs = jax.tree.map(where_done, reset_obs, state.obs)
+
+    next_info = state.info
+    done_count_key = f'{self._info_key}_done_count'
+    if self._full_reset and reset_state:
+      next_info = jax.tree.map(where_done, reset_state.info, state.info)
+      next_info[done_count_key] = state.info[done_count_key]
+
+      if 'steps' in next_info:
+        next_info['steps'] = state.info['steps']
+      preserve_info_key = f'{self._info_key}_preserve_info'
+      if preserve_info_key in next_info:
+        next_info[preserve_info_key] = state.info[preserve_info_key]
+
+    next_info[done_count_key] += state.done.astype(int)
+    next_info[f'{self._info_key}_rng'] = reset_rng
+
+    return state.replace(data=data, obs=obs, info=next_info)
 
 
 class BraxDomainRandomizationVmapWrapper(Wrapper):

mujoco_playground/_src/wrapper_test.py

@@ -13,46 +13,116 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for the wrapper module."""
+
 import functools
 
 from absl.testing import absltest
+from absl.testing import parameterized
+from brax.envs.wrappers import training as brax_training
 import jax
 import jax.numpy as jp
-import numpy as np
-
 from mujoco_playground._src import dm_control_suite
 from mujoco_playground._src import wrapper
+import numpy as np
 
 
-class WrapperTest(absltest.TestCase):
+class WrapperTest(parameterized.TestCase):
 
-  def test_auto_reset_wrapper(self):
+  @parameterized.named_parameters(
+      ('full_reset', True),
+      ('cache_reset', False),
+  )
+  def test_auto_reset_wrapper(self, full_reset):
+    """Tests the AutoResetWrapper."""
     class DoneEnv:
 
       def __init__(self, env):
         self._env = env
 
       def reset(self, key):
-        return self._env.reset(key)
+        state = self._env.reset(key)
+        state.info['AutoResetWrapper_preserve_info'] = 1
+        state.info['other_info'] = 1
+        return state
 
       def step(self, state, action):
         state = self._env.step(state, jp.ones_like(action))
         state = state.replace(done=action[0] > 0)
+        state.info['AutoResetWrapper_preserve_info'] = 2
+        state.info['other_info'] = 2
         return state
 
     env = wrapper.BraxAutoResetWrapper(
-        DoneEnv(dm_control_suite.load('CartpoleBalance'))
+        brax_training.VmapWrapper(
+            DoneEnv(dm_control_suite.load('CartpoleBalance'))
+        ),
+        full_reset=full_reset,
     )
 
     jit_reset = jax.jit(env.reset)
     jit_step = jax.jit(env.step)
-    state = jit_reset(jax.random.PRNGKey(0))
-    first_qpos = state.info['first_state'].qpos
+    state = jit_reset(jax.random.PRNGKey(0)[None])
+    first_qpos = state.data.qpos
 
-    state = jit_step(state, -jp.ones(env._env.action_size))
+    # First step should not be done.
+    state = jit_step(state, -jp.ones(env._env.action_size)[None])
+    np.testing.assert_allclose(state.info['AutoResetWrapper_done_count'], 0)
     self.assertGreater(np.linalg.norm(state.data.qpos - first_qpos), 1e-3)
-    state = jit_step(state, jp.ones(env._env.action_size))
-    np.testing.assert_allclose(state.data.qpos, first_qpos, atol=1e-6)
+    self.assertEqual(state.info['AutoResetWrapper_preserve_info'], 2)
+    self.assertEqual(state.info['other_info'], 2)
+
+    for i in range(1, 3):
+      state = jit_step(state, jp.ones(env._env.action_size)[None])
+      jax.tree.map(lambda x: x.block_until_ready(), state)
+      if full_reset:
+        self.assertTrue((state.data.qpos != first_qpos).all())
+      else:
+        np.testing.assert_allclose(state.data.qpos, first_qpos, atol=1e-6)
+      np.testing.assert_allclose(state.info['AutoResetWrapper_done_count'], i)
+      self.assertEqual(state.info['AutoResetWrapper_preserve_info'], 2)
+      expected_other_info = 1 if full_reset else 2
+      self.assertEqual(state.info['other_info'], expected_other_info)
+
+  @parameterized.named_parameters(
+      ('full_reset', True),
+      ('cache_reset', False),
+  )
+  def test_evalwrapper_with_reset(self, full_reset):
+    """Tests EvalWrapper with reset in the AutoResetWrapper."""
+    episode_length = 10
+    num_envs = 4
+
+    env = dm_control_suite.load('CartpoleBalance')
+    env = wrapper.wrap_for_brax_training(
+        env,
+        episode_length=episode_length,
+        full_reset=full_reset,
+    )
+    env = brax_training.EvalWrapper(env)
+
+    jit_reset = jax.jit(env.reset)
+    jit_step = jax.jit(env.step)
+
+    rng = jax.random.PRNGKey(0)
+    rng = jax.random.split(rng, num_envs)
+    state = jit_reset(rng)
+    first_obs = state.obs
+    action = jp.zeros((num_envs, env.action_size))
+
+    for _ in range(episode_length):
+      state = jit_step(state, action)
+
+    # All episodes should finish at episode_length.
+    avg_episode_length = state.info['eval_metrics'].episode_steps.mean()
+    np.testing.assert_allclose(avg_episode_length, episode_length, atol=1e-6)
+    active_episodes = state.info['eval_metrics'].active_episodes
+    self.assertTrue(np.all(active_episodes == 0))
+
+    np.testing.assert_array_equal(state.info['steps'], 10 * np.ones(num_envs))
+    if full_reset:
+      self.assertTrue((state.obs != first_obs).all())
+    else:
+      np.testing.assert_allclose(state.obs, first_obs, rtol=1e-6)
 
   def test_domain_randomization_wrapper(self):
     def randomization_fn(model, rng):