Andrew Luo
Andrew Luo
Not sure if this is the right place to post this - Re: whether ``` jax.grad ``` delivers useful gradients: I have been playing with gradients over mjx.step (Newton solver;...
I have a depth camera implementation using ```mjx.ray``` from a personal project, that I discuss [here](https://github.com/google-deepmind/mujoco/issues/1485). It renders quite quickly and I've successfully used it for vision-based RL.
I haven't tried differentiating through mjx.step with settings other than from [1182](https://github.com/google-deepmind/mujoco/issues/1182). Do these settings work for you? I found that with the eulerdamp flag disabled, I could stably simulate...
In case this is still helpful, I think mj_multiRay() can now be implemented by using jax.vmap on mjx.ray(). For example, please see this simple [depth camera implementation](https://github.com/Andrew-Luo1/jax_shac/blob/main/vision/2dof_ball.ipynb). Batching across environments,...
In my experience, ocassional nans are to be expected from simulators. I also ran into nans with PPO training. Here's three thoughts: 1) Seeing how it's a floating point error,...