https://github.com/ami-iit/jaxsim

A differentiable physics engine and multibody dynamics library for control and robot learning.
https://github.com/ami-iit/jaxsim

aba ad automatic-differentiation crba dynamics featherstone jacobian jax jit kinematics ode physics physics-engine rigid-body-dynamics rnea robotics robotics-control robotics-simulation simulation simulation-modeling

Last synced: about 1 month ago
JSON representation

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

Host: GitHub
URL: https://github.com/ami-iit/jaxsim
Owner: ami-iit
License: bsd-3-clause
Created: 2022-01-31T16:44:28.000Z (almost 4 years ago)
Default Branch: main
Last Pushed: 2025-11-14T15:23:11.000Z (about 1 month ago)
Last Synced: 2025-11-14T15:26:54.220Z (about 1 month ago)
Topics: aba, ad, automatic-differentiation, crba, dynamics, featherstone, jacobian, jax, jit, kinematics, ode, physics, physics-engine, rigid-body-dynamics, rnea, robotics, robotics-control, robotics-simulation, simulation, simulation-modeling
Language: Python
Homepage: https://jaxsim.readthedocs.io/
Size: 7.34 MB
Stars: 162
Watchers: 16
Forks: 18
Open Issues: 22
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Codeowners: .github/CODEOWNERS

Awesome Lists containing this project

README

          # 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][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][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

```python

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!

[jax]: https://github.com/google/jax/

[sdformat]: https://github.com/gazebosim/sdformat

[notation]: https://research.tue.nl/en/publications/multibody-dynamics-notation-version-2

[passive_viewer_mujoco]: https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer

## Installation

With conda

You can install the project using [`conda`][conda] as follows:

```bash

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](https://github.com/git-lfs/git-lfs/blob/main/INSTALLING.md) installed and properly configured on your system before installation. After cloning the repository, run:

```bash

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`][pixi] project as follows:

```bash

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`](https://pixi.sh/latest/reference/pixi_manifest/#the-system-requirements-table) table, i.e. adding

~~~toml

[system-requirements]

cuda = "12"

~~~

if you are using a `pixi.toml` file or

~~~toml

[tool.pixi.system-requirements]

cuda = "12"

~~~

if you are using a `pyproject.toml` file.

With pip

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

```bash

pip install jaxsim

```

Check [`pyproject.toml`](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](https://gazebosim.org/docs) 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][jax_gpu] of JAX.

Contributors installation (with conda)

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

```bash

conda env create -f environment.yml

```

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

```bash

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](https://github.com/git-lfs/git-lfs/blob/main/INSTALLING.md) installed and properly configured on your system before installation. After cloning the repository, run:

```bash

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`][pixi] as follows:

```bash

pixi install

```

See `pixi task list` for a list of available tasks.

[conda]: https://anaconda.org/

[pip]: https://github.com/pypa/pip/

[pixi]: https://pixi.sh/

[venv]: https://docs.python.org/3/tutorial/venv.html

[jax_gpu]: https://github.com/google/jax/#installation

## Documentation

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

[readthedocs]: https://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

```python

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][RBDA]

book by Roy Featherstone.

The algorithms and some simulation features were inspired by its accompanying [code][spatial_v2].

[RBDA]: https://link.springer.com/book/10.1007/978-1-4899-7560-7

[spatial_v2]: http://royfeatherstone.org/spatial/index.html#spatial-software

The development of JaxSim started in late 2021, inspired by early versions of [`google/brax`][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][jax] in this field.

Brax v2 was later implemented with reduced coordinates, following an approach comparable to JaxSim.

The development then shifted to [MJX][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][PFS].

- JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.

[brax]: https://github.com/google/brax

[mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html

[PFS]: http://sdformat.org/tutorials?tut=pose_frame_semantics

## Contributing

We welcome contributions from the community.

Please read the [contributing guide](./CONTRIBUTING.md) to get started.

## Citing

If you use JaxSim in your work, please cite our upcoming paper:

```bibtex

@software{ferretti_accelerated_optimization_2025,

  author       = {Filippo Luca Ferretti and Diego Ferigo and Carlotta Sartore and Alessandro Croci and Omar G. Younis and Silvio Traversaro and Daniele Pucci},

  title        = {Hardware-Accelerated Morphology Optimization via Physically Consistent Differentiable Simulation},

  year         = {2025},

  url          = {https://github.com/ami-iit/jaxsim}

}

```

## People

| Authors | Maintainers |

|:------:|:-----------:|

| [][df] [][ff] | [][ff] [][ac] |

[df]: https://github.com/diegoferigo

[ff]: https://github.com/flferretti

[ac]: https://github.com/xela-95

## License

[BSD3](https://choosealicense.com/licenses/bsd-3-clause/)

ecosyste.ms

Data

Tools

Indexes

Applications

Experiments

Awesome

https://github.com/ami-iit/jaxsim

Awesome Lists containing this project

README