There’s been loads of work in this area (differentiable simulators) in JAX. I think Google/Deepmind are porting mujoco to JAX (https://mujoco.readthedocs.io/en/

There’s been loads of work in this area (differentiable simulators) in JAX. I think Google/Deepmind are porting mujoco to JAX (https://mujoco.readthedocs.io/en/stable/mjx.html https://mujoco.readthedocs.io/en/stable/mjx.html), and for fluid dynamics or deformable bodies there are research projects like https://github.com/google/jax-cfd https://github.com/google/jax-cfd and https://daxbench.github.io/ https://daxbench.github.io/.