How to deal with an impure reset function

[Theo-Cheynel]

Hi,

I want to train an agent to imitate reference motions from a motion capture dataset. I wrote a custom environment, which works well when the reference motion capture clip is "hardcoded" (when it is the same throughout all environments).
However, I would like to make the reference clip vary across envs, in other words,I want the env.reset to pick a random motion clip (at the moment, motion clips are obtained from another class' __getitem__ method).

The thing is, that would make the reset function an impure function, because its outputs would differ everytime, and JAX jitting only supports pure functions. Is there a workaround you can think of ?

At the very least, is it possible to tell ppo.train not to jit the reset function, while still jitting the step function ?

Thanks for your help

[btaba]

Hi @Theo-Cheynel , are you able to do something like this in the env.reset:

mocap = jp.ones((10, 3)) * jp.arange(0, 10)[:, None]  # load this once
rng = jax.random.PRNGKey(0)

def reset(rng):
  rng, key = jax.random.split(rng, 2)
  return rng, jax.random.choice(key, mocap, (1,))

rng, val = jax.jit(reset)(rng)

[btaba]

transferred to discussion due to inactivity

[Theo-Cheynel]

Thanks a lot for answering ! Sorry it took so long, I had to fix several problems with my BRAX environment (one of them was because I stored data in the state.info, updating it at each step, and didn't notice that the AutoResetWrapper wasn't replacing the info, only the qp and obs. In hindsight it made sense, but boy was that tough to debug !).

So now I'm finally facing the issue that I raised in this discussion 😅 So far it's simply overfitting on a single motion sequence, since I didn't deal with the impure function yet. Here's a visual from the very first functional training ! I'm hyped 🚀

motion.imitation.mp4

If I do as you proposed, won't that mean that the entire mocap data is stored by JAX on the GPU, as many times as the number of environments that run in parallel ? I'll try it but I feel like if I have 100 MB of mocap data, il will load 100*128=12.8 GB of RAM. Am I right in assuming so ? And if yes, do you know of a way for JAX to share data across the parallel envs so that it isn't replicated ?

Thank you so much ❤️

[btaba]

Awesome demo! I don't think the full dataset would be replicated per environment, but if memory becomes an issue, you can always make a host callback to reset states (it doesn't all need to happen in one jitted call)