Skip to main content

A differentiable physics engine and multibody dynamics library for control and robot learning.

Project description

JaxSim

JaxSim is a differentiable physics engine built with JAX, tailored for co-design and robotic learning applications.



Features

  • Physically consistent differentiability w.r.t. hardware parameters.
  • Closed chain dynamics support.
  • Reduced-coordinate physics engine for fixed-base and floating-base robots.
  • Fully Python-based, leveraging JAX following a functional programming paradigm.
  • Seamless execution on CPUs, GPUs, and TPUs.
  • Supports JIT compilation and automatic vectorization for high performance.
  • Compatible with SDF models and URDF (via sdformat conversion).

[!WARNING] This project is still experimental. APIs may change between releases without notice.

[!NOTE] JaxSim currently focuses on locomotion applications. Only contacts between bodies and smooth ground surfaces are supported.

How to use it

import pathlib

import icub_models
import jax.numpy as jnp

import jaxsim.api as js

# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")

joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
          'r_ankle_roll')

# Build and reduce the model
model_description = pathlib.Path(model_path)

full_model = js.model.JaxSimModel.build_from_model_description(
    model_description=model_description,
)

model = js.model.reduce(model=full_model, considered_joints=joints)

# Get the number of degrees of freedom
ndof = model.dofs()

# Initialize data and simulation
# Note that the default data representation is mixed velocity representation
data = js.data.JaxSimModelData.build(
    model=model, base_position=jnp.array([0.0, 0.0, 1.0])
)

T = jnp.arange(start=0, stop=1.0, step=model.time_step)

tau = jnp.zeros(ndof)

# Simulate
for _ in T:
    data = js.model.step(
        model=model, data=data, link_forces=None, joint_force_references=tau
    )

Check the example folder for additional use cases!

Installation

With conda

You can install the project using conda as follows:

conda install jaxsim -c conda-forge

GPU support for JAX will be automatically installed if a compatible GPU is detected.

With pixi

Note

The minimum version of pixi required is 0.39.0.

Since the pixi.lock file is stored using Git LFS, make sure you have Git LFS installed and properly configured on your system before installation. After cloning the repository, run:

git lfs install && git lfs pull

This ensures all LFS-tracked files are properly downloaded before you proceed with the installation.

You can add the jaxsim dependency in your pixi project as follows:

pixi add jaxsim

If you are on Linux and you want to use a cuda-powered version of jax, remember to add the appropriate line in the system-requirements table, i.e. adding

[system-requirements]
cuda = "13"

if you are using a pixi.toml file or

[tool.pixi.system-requirements]
cuda = "13"

if you are using a pyproject.toml file.

With pip

You can install the project using pypa/pip, preferably in a virtual environment, as follows:

pip install jaxsim

Check pyproject.toml for the complete list of optional dependencies. You can obtain a full installation using jaxsim[all].

If you need URDF support, follow the official instructions to install Gazebo Sim on your operating system, making sure to obtain sdformat ≥ 13.0 and gz-tools ≥ 2.0.

You don't need to install the entire Gazebo Sim suite. For example, on Ubuntu, it is sufficient to install the libsdformat* and gz-tools2 packages.

If you need GPU support, follow the official installation instructions of JAX.

Contributors installation (with conda)

If you want to contribute to the project, we recommend creating the following jaxsim conda environment first:

conda env create -f environment.yml

Then, activate the environment and install the project in editable mode:

conda activate jaxsim
pip install --no-deps -e .
Contributors installation (with pixi)

Note

The minimum version of pixi required is 0.39.0.

Since the pixi.lock file is stored using Git LFS, make sure you have Git LFS installed and properly configured on your system before installation. After cloning the repository, run:

git lfs install && git lfs pull

This ensures all LFS-tracked files are properly downloaded before you proceed with the installation.

You can install the default dependencies of the project using pixi as follows:

pixi install

See pixi task list for a list of available tasks.

Documentation

The JaxSim API documentation is available at jaxsim.readthedocs.io.

Additional features

Jaxsim can also be used as a multi-body dynamics library! With full support for automatic differentiation of RBDAs (forwards and reverse mode) and automatic differentiation against both kinematic and dynamic parameters.

Using JaxSim as a multibody dynamics library

import pathlib

import icub_models
import jax.numpy as jnp

import jaxsim.api as js

# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")

joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
          'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
          'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
          'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
          'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
          'r_ankle_roll')

# Build and reduce the model
model_description = pathlib.Path(model_path)

full_model = js.model.JaxSimModel.build_from_model_description(
    model_description=model_description,
)

model = js.model.reduce(model=full_model, considered_joints=joints)

# Initialize model data
data = js.data.JaxSimModelData.build(
    model=model,
    base_position=jnp.array([0.0, 0.0, 1.0]),
)

# Frame and dynamics computations
frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")

# Frame transformation
W_H_F = js.frame.transform(
    model=model, data=data, frame_index=frame_index
)

# Frame Jacobian
W_J_F = js.frame.jacobian(
    model=model, data=data, frame_index=frame_index
)

# Dynamics properties
M = js.model.free_floating_mass_matrix(model=model, data=data)      # Mass matrix
h = js.model.free_floating_bias_forces(model=model, data=data)      # Bias forces
g = js.model.free_floating_gravity_forces(model=model, data=data)   # Gravity forces
C = js.model.free_floating_coriolis_matrix(model=model, data=data)  # Coriolis matrix

# Print dynamics results
print(f"{M.shape=} \n{h.shape=} \n{g.shape=} \n{C.shape=}")

Credits

The RBDAs are based on the theory of the Rigid Body Dynamics Algorithms book by Roy Featherstone. The algorithms and some simulation features were inspired by its accompanying code.

The development of JaxSim started in late 2021, inspired by early versions of google/brax. At that time, Brax was implemented in maximal coordinates, and we wanted a physics engine in reduced coordinates. We are grateful to the Brax team for their work and for showing the potential of JAX in this field.

Brax v2 was later implemented with reduced coordinates, following an approach comparable to JaxSim. The development then shifted to MJX, which provides a JAX-based implementation of the Mujoco APIs.

The main differences between MJX/Brax and JaxSim are as follows:

  • JaxSim supports out-of-the-box all SDF models with Pose Frame Semantics.
  • JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.

Contributing

We welcome contributions from the community. Please read the contributing guide to get started.

Citing

If you use JaxSim in your work, please cite the following paper:

@article{ferretti_contact_aware_2026,
  author       = {Filippo Luca Ferretti and Diego Ferigo and Alessandro Croci and Carlotta Sartore and Omar G. Younis and Silvio Traversaro and Daniele Pucci},
  title        = {Contact-Aware Morphology Optimization via Physically Consistent Differentiable Simulation},
  journal      = {IEEE Robotics and Automation Letters},
  year         = {2026},
  doi          = {10.1109/LRA.2026.3678125}
}

People

Authors Maintainers

License

BSD3

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxsim-0.8.3.dev56.tar.gz (203.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

jaxsim-0.8.3.dev56-py3-none-any.whl (169.0 kB view details)

Uploaded Python 3

File details

Details for the file jaxsim-0.8.3.dev56.tar.gz.

File metadata

  • Download URL: jaxsim-0.8.3.dev56.tar.gz
  • Upload date:
  • Size: 203.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for jaxsim-0.8.3.dev56.tar.gz
Algorithm Hash digest
SHA256 4d2a609374a650d5dbad6dd0f88c1c42e6bd9ea3e0d02c19b72ca2f92a31eb3d
MD5 12b6419ae71695a12cd92392705f1268
BLAKE2b-256 27dbc3fa553890b4cea63e8bea05b7194a09f464fce86ad06c4dbbef274202f0

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxsim-0.8.3.dev56.tar.gz:

Publisher: ci_cd.yml on gbionics/jaxsim

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file jaxsim-0.8.3.dev56-py3-none-any.whl.

File metadata

  • Download URL: jaxsim-0.8.3.dev56-py3-none-any.whl
  • Upload date:
  • Size: 169.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for jaxsim-0.8.3.dev56-py3-none-any.whl
Algorithm Hash digest
SHA256 faf1ce3426f3582e9ba45608444201853040b723a8853c69887fef7db9131d2e
MD5 9291e798c9241dd78c1f040458a86a58
BLAKE2b-256 c5dd99f84c0566cf92e2dbfd2b2d1961e9eddfac4b33bd0e2bc5266f0251be92

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxsim-0.8.3.dev56-py3-none-any.whl:

Publisher: ci_cd.yml on gbionics/jaxsim

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page