Potential JAX-based physics for real-time simulation in Slicer?

Support
1

Hello Slicer community,

I’m a developer interested in physics simulation. My project, brax viewer, visualizes JAX-based physics engines in real time, which is a new open-source package developed as part of Google Summer of Code. Briefly, JAX is an extremely fast, differentiable computation package that runs on GPU.

I am looking for a way for my viewer to be used with Slicer. So I share it with the community and have been browsing some relevant discussions. In threads like the one on biomechanics simulation for pectus excavatum, Blender can animate deformation, and tools like OpenSIM or FEBio can simulate but require patient-specific material properties for precision.

Could JAX engines be used to reverse-estimate these material properties in real-time?

There are two relevant physics simulation engines suited for this problem:

  • MJX for rigid body dynamics
  • JAX-FEM for Finite Element Method

JAX has two advantages:

  • Speed: Because JAX runs natively on the GPU, it’s fast enough for real-time simulation. It can also host many environments in parallel, which is ideal for robot trajectory exploration.
  • Differentiability: The entire simulation is differentiable, making it powerful for optimization problems like material property estimation and learning-based robot training.

Based on this discussion, Slicer’s WebServer and OpenIGTLink seem to be good interfaces.

Thank you for your time and any thoughts you have! I am open to any suggestions.

2

@Bruce_Hu, this sounds interesting, thanks for sharing. We have a 3D Slicer extension integrating the SOFA framework ( GitHub - Slicer/SlicerSOFA: 3D Slicer extension to enable simulations using the SOFA framework ). A common pattern to integrate simulation libraries in 3D Sicer is using numpy.array objects as a common language to transfer data between the foreign library and 3D Slicer. This is what is used in SlicerSOFA and more recently @pieper has succesfully demoed integration with NVIDIA Warp.

@Bruce_Hu, what do you think is the advantage of using Brax viewer with streamed slicer data vs. integrating JAX directly in Slicer? 3D Slicer has already a wealth of functionality for visualization.

3

Welcome @Bruce_Hu - it sounds like you are into some interesting things and I hope you’ll find Slicer useful for your work. Now is an exciting time for simulation and medical imaging and as @RafaelPalomar says, we are actively exploring a lot of options.

You may already know of this, but the Newton Physics project also looks interesting in this space.

In addition to the WebServer and OpenIGTLink, you may want to explore RPyC and shared memory as in the client and server examples here. There’s some description of our electromagnetic simulation examples and visualizations here. Note that we preferred the client-server model for SimNIBS because, like JAX-FEM, it is GPL. But it’s also an option for decoupling simulation from UI processes in general.

As a general rule we also like the option we like being able to have the option of installing all the packages in the Slicer python interpreter for portability, although that’s not always possible. It does seem to work well for Warp in my experience.

Yes, this is a great idea and definitely something that would be valuable. We often have 4D datasets and other data that could be used for this purpose.

4

Thanks, @RafaelPalomar. You’re right, the goal is definitely to integrate JAX directly and use Slicer’s own visualization, not build a separate one. Using numpy as the common language is ideal, since jax.numpy mirrors numpy’s API.

Two reasons to integrate JAX are:

  1. Integrated Reinforcement Learning Pipeline
    Robots can be trained directly within the Slicer environment for autonomous systems.
    JAX can be used for simulation - MJX, and reinforcement learning trainer - Brax. (In soft dynamics simulation, SOFA, and NVIDIA Warp seems more suitable than MJX.)
  2. Accelerated Scientific Compute Backend
    JAX is highly optimized for high-throughput computation and cleanly supports different hardware (CPU/GPUs/TPUs). However, its I/O, compilation, and asynchronous architecture are unique, so it requires careful design for integration with the Slicer server (non-blocking, real-time visualization).
5

@Bruce_Hu, thanks for the reply. It is great to see the number of alternatives and the potential to make them work in 3D Slicer. If you have the time and interest I would suggest you join one of the developer meetings to have a discussion on this. These meetings are open and happen on Tuesdays at 10 am EST; usually they are announced at Weekly meetings - 3D Slicer Community the same day or the day before.

One idea could be to have an extension offering an API and tools to interact with simulations in Slicer through different back-ends.

Thanks for the discussion!

1 Like
6

Thank you for the reply, @pieper!

I think the server-client model works well with JAX’s asynchronous execution, and my Brax Viewer follows this design too.

Thanks for the suggestion to Newton Physics and RPyC. It’s nice to see Newton’s support for MuJoCo, as it is widely used in reinforcement learning frameworks like Brax for training algorithms. I will definitely look into these.

I understand that a direct installation is ideal for portability. JAX’s installation is automated by pip install jax[device], where device can be CUDA/TPU. So there could be a direct integration method.

And regarding licenses, is Apache 2.0 good for direct integration? JAX, Brax and my Brax Viewer are all Apache 2.0.

7

Yes, we definitely prefer Apache 2 (see, for example, comments here on why we avoid GPL code). It seems to be a common thought, since Warp and Newton are recent Apache examples too.

As Rafael suggests, discussing these topics at a developer call would be a chance to get us all better up to speed on the options these packages provide and to think about how we can apply them in various medical image computing contexts.

1 Like