A differentiable physics engine and multibody dynamics library for control and robot learning.
JaxSim is a differentiable physics engine and multibody dynamics library designed for applications in control and robot learning, implemented with JAX.
Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
google/jax
following a functional programming paradigm.[!WARNING] This project is still experimental, APIs could change between releases without notice.
[!NOTE] JaxSim currently focuses on locomotion applications. Only contacts between bodies and smooth ground surfaces are supported.
The JaxSim API documentation is available at jaxsim.readthedocs.io.
You can install the project using conda
as follows:
conda install jaxsim -c conda-forge
You can enforce GPU support, if needed, by also specifying "jaxlib = * = *cuda*"
.
You can install the project using pypa/pip
, preferably in a virtual environment, as follows:
pip install jaxsim
Check setup.cfg
for the complete list of optional dependencies.
You can obtain a full installation using jaxsim[all]
.
If you need GPU support, follow the official installation instructions of JAX.
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 .
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 showing the potential of JAX in this field.
Brax v2 was later implemented reduced coordinates, following an approach comparable to JaxSim. The development then shifted to MJX, which today provides a JAX-based implementation of the Mujoco APIs.
The main differences between MJX/Brax and JaxSim are as follows:
We welcome contributions from the community. Please read the contributing guide to get started.
@software{ferigo_jaxsim_2022,
author = {Diego Ferigo and Filippo Luca Ferretti and Silvio Traversaro and Daniele Pucci},
title = {{JaxSim}: A Differentiable Physics Engine and Multibody Dynamics Library for Control and Robot Learning},
url = {http://github.com/ami-iit/jaxsim},
year = {2022},
}
Author | Maintainers |
---|---|