a question about re-implementing the agile policy in brax with MJX

Hi Tairan,

Thank you for releasing the code!

I am trying to re-implement the agile PPO policy in `brax` because I need the physics simulation to be differentiable for my application. I would like to know if this is a feasible option at all? Thank you!


Best,
Randy