Ctrl+K

JAXsim

JAXsim#

A scalable physics engine and multibody dynamics library implemented with JAX. With JIT batteries đŸ”‹

Note

This simulator currently focuses on locomotion applications. Only contacts with ground are supported.

Features#

Performance

Physics engine in reduced coordinates implemented with JAX. Compatibility with JIT compilation for increased performance and transparent support to execute logic on CPUs, GPUs, and TPUs. Parallel multi-body simulations on hardware accelerators for significantly increased throughput

Model Parsing

Support for SDF models (and, upon conversion, URDF models). Revolute, prismatic, and fixed joints supported.

Automatic Differentiation

Support for automatic differentiation of rigid body dynamics algorithms (RBDAs) for model-based robotics research. Soft contacts model supporting full friction cone and sticking / slipping transition.

Complex Dynamics

JAXsim provides a variety of integrators for the simulation of multibody dynamics, including RK4, Heun, Euler, and more. Support of multiple velocities representations.

Examples#

Explore and learn how to use the library through practical demonstrations available in the examples folder.

Credits#

The physics module of JAXsim is based on the theory of the Rigid Body Dynamics Algorithms book by Roy Featherstone. We structured part of our logic following its accompanying code. The physics engine is developed entirely in Python using JAX.

The inspiration for developing JAXsim originally stemmed from early versions of google/brax. Here below we summarize the differences between the projects:

  • JAXsim simulates multibody dynamics in reduced coordinates, while brax v1 uses maximal coordinates.

  • The new v2 APIs of brax (and the new MJX) were then implemented in reduced coordinates, following an approach comparable to JAXsim, with major differences in contact handling.

  • The rigid-body algorithms used in JAXsim allow to efficiently compute quantities based on the Euler-Poincarè formulation of the equations of motion, necessary for model-based robotics research.

  • JAXsim supports SDF (and, indirectly, URDF) models, assuming the model is described with the recent Pose Frame Semantics.

  • Contrarily to brax, JAXsim only supports collision detection between bodies and a compliant ground surface.

  • The RBDAs of JAXsim support automatic differentiation, but this functionality has not been thoroughly tested.

People#

Author and Maintainer#

Diego Ferigo

Maintainer#

Filippo Luca Ferretti

License#

BSD3

Contents