Thread on possible JAX-based physics engines

[clement-moulin-frier]

Currently we rely on JAX-MD, but the library is not very well maintained, which could become an issue in the future.

Our main requirement is to be able to simulate a large number of entities in parallel (> 10K), which require efficient search of nearest neighbors.

Some potential alternatives recommended by @eleninisioti .

• MuJoCo XLA (MJX) - MuJoCo Documentation
• adaptive-intelligent-robotics/Kheperax: High-performance JAX-powered simulator for robotic navigation in 2D mazes, optimized for Quality-Diversity algorithm research and benchmarking.
  • Looks like they deal with wall detection
• FLAIROx/Kinetix: Reinforcement learning on general 2D physics environments in JAX. ICLR 2025 Oral.
• MichaelTMatthews/Jax2D: Highly scalable 2D JAX physics engine.
  • But
    • The main reason to use Jax2D over other JAX physics engines such as Brax or MJX is that Jax2D scenes are (largely) dynamically specified. However, Jax2D always has O(n^2) runtime with respect to the number of entities in a scene, since we must always calculate the full collision resolution for every pair of entities. This means it is usually not appropriate for simulating scenes with large numbers (>100) of entities.
    • In short: Jax2D excels at simulating lots of small and diverse scenes in parallel very fast.

This thread can be used to post more potential alternative solutions to JAX-MD