Is your feature request related to a problem? Please describe.
There are many applications where we would like to track gradients throughout the simulator. One example could be to optimize some of the simulation parameters such as the elastic modulus, the joint stiffness etc. to fit the recorded experimental data as closely as possible.
Another limitation of the current simulator implementation is that it cannot be run massively in parallel. Because of the Python multiprocessing limitations, it is limited to the number of cores of your CPU (to the best of my knowledge). This constrains the performance of Reinforcement Learning (RL) algorithms for example, as the number of parallel simulations is limited to just a few.
In contrast, there are examples now of simulators running on the GPU, such as Nvidia Isaac or Brax to enable massive parallelism of simulations.
Describe the solution you'd like
While I am not an expert in JAX, an implementation of the simulator in Jax could be a good way to achieve multiple improvements in one go:
- Differentiability
- Running the simulator on GPUs / TPUs3 / etc.
- Running the simulator massively parallel on the GPU.
- JIT compilation of the entire simulator
Describe alternatives you've considered
I saw on your documentation that you are currently working on a C++ implementation of the simulator. I wonder if the majority of the expected speed-ups compared to the current Numba / Python implementation could also be achieved with JAX JIT compilation? Alternatively, the gradient could also be manually tracked in C++, but that would be quite a massive undertaking I guess...
Is your feature request related to a problem? Please describe.
There are many applications where we would like to track gradients throughout the simulator. One example could be to optimize some of the simulation parameters such as the elastic modulus, the joint stiffness etc. to fit the recorded experimental data as closely as possible.
Another limitation of the current simulator implementation is that it cannot be run massively in parallel. Because of the Python multiprocessing limitations, it is limited to the number of cores of your CPU (to the best of my knowledge). This constrains the performance of Reinforcement Learning (RL) algorithms for example, as the number of parallel simulations is limited to just a few.
In contrast, there are examples now of simulators running on the GPU, such as Nvidia Isaac or Brax to enable massive parallelism of simulations.
Describe the solution you'd like
While I am not an expert in JAX, an implementation of the simulator in Jax could be a good way to achieve multiple improvements in one go:
Describe alternatives you've considered
I saw on your documentation that you are currently working on a C++ implementation of the simulator. I wonder if the majority of the expected speed-ups compared to the current Numba / Python implementation could also be achieved with JAX JIT compilation? Alternatively, the gradient could also be manually tracked in C++, but that would be quite a massive undertaking I guess...