Recommended way to actively set system's state

[yardenas]

In many sim2real RL applications, it makes sense to leverage our ability to "teleport" between states, without really following the dynamics. For example, in the cartpole environment, instead of exploring and learning to swing up the pendulum to initially get rewards, we can just sample the upright position (how do we know that the upright position is interesting is a matter for a different discussion though).

What's the recommended way to achieve that in Brax?
Going over previous tickets, seems like there are two approaches to achieve that:

• Overriding the reset function, e.g., Manually Set Initial Pose of Agent in the Reset Function #457
• Directly overriding the system's state: Resetting to a particular state #249

What would be the recommended way?
This PR indicates that using reset can significantly slow down the simulation, while I'm not sure if this approach would give a valid state for the simulator (for instance contact forces etc.)

Any ideas anyone?

[erikfrey]

You're welcome to rely on the Reset function to create a pool of states that are sampled from some distribution, or override the state yourself during Step - either will work and it depends on your use case.

Brax is designed so resetting has virtually no impact on simulation speed. My default suggestion is to just add your "start upright" states into reset and everything should just work:

def reset(rng: jax.Array) -> State:
  # set the pendulum upright 20% of the time
  up = jax.random.uniform(rng) < 0.2
  pendulum_ypos = jp.where(up, 1.0, 0.0)
  ...

[yardenas]

Thank you very much @erikfrey! 🙌

, or override the state yourself during Step - either will work and it depends on your use case.

Hope not a too trivial question -- what would be the correct way to do it via Step?

I'm also curious what would be the optimal way to make this distribution dynamic? For example, in an active learning setting, where we'd want to sample only novel/unseen states