Reset function in BraxAutoResetWrapper

[bkpcoding]

Hello, the auto-reset wrapper used in the mujoco playground seems to use the same initial state for every reset. I was interested if there was a specific reason for this? Because in some of my experiments with FastTD3 as shown below resetting the environments to random states performs much better than resetting to initial state. Thank you!!

[btaba]

Hi @bkpcoding , we did this mainly for performance (wallclock) reasons. How do you reset to random states? With a host-callback or with a where? We do the host-callback way here

[bkpcoding]

Hmm, I get similar graph in wallclock time as well (maybe 5mins slow). So I believe it is a good trade-off for higher performance?

I use something like this:

  def random_reset(self, rng: jax.Array, state: mjx_env.State, should_reset: bool) -> mjx_env.State:

    def apply_reset(args):
      rng, state = args
      new_state = self.reset(rng)
      info = {
          **new_state.info,
          'episode_done': state.info.get('episode_done', jp.array(False)),
          'episode_metrics': state.info.get('episode_metrics', {}),
          'first_obs': state.info.get('first_obs', {
              'state': state.obs,
          }),
          'first_state': state.info.get('first_state', state.obs),
          'raw_obs': state.info.get('raw_obs', {
              'state': state.obs,
          }),
          'steps': state.info.get('steps', jp.array(0)),
          'truncation': state.info.get('truncation', jp.array(False)),
      }
      # make the state the same structure as we have the other stage
      return new_state.replace(info=info)
    
    def no_reset(args):
      _, state = args
      return state
    return jax.lax.cond(should_reset, apply_reset, no_reset, (rng, state))

[btaba]

Ok interesting! I suppose it makes sense to allow the option to call reset on every step. FYI jax.lax.cond turns into a where in a vmap, and it makes it a bit more clear that both branches run on every step.

https://docs.jax.dev/en/latest/_autosummary/jax.lax.cond.html#:~:text=In%20contrast%20with,JAX%E2%80%99s%20tracing%20model).

[btaba]

FYI #179