Dial-MPC implementation in Genesis #28
Description
I’m currently trying to reproduce the paper Dial-MPC in the Genesis simulator as part of my academic work. However, I’m running into a few issues:
1. Integration with JAX: Dial-MPC is implemented using JAX, while Genesis is not natively compatible with JAX. I’m not sure how to make the two interact effectively, especially for running parallel rollouts or differentiable computations.
2. Managing State in Genesis: The way Genesis handles simulation state (observations, resets, stepping) seems quite different from other simulators like Brax. I’m struggling to figure out how to properly clone or restore the environment state during MPC rollouts, which is crucial for sampling multiple trajectories from the same initial state.
If anyone has experience with MPC in Genesis, I’d really appreciate any tips or example code.
Thanks in advance!