# JaxSim
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.
## Features
- Physics engine in reduced coordinates supporting fixed-base and floating-base robots.
- Multibody dynamics library providing all the necessary components for developing model-based control algorithms.
- Completely developed in Python with [`google/jax`][jax] following a functional programming paradigm.
- Transparent support for running on CPUs, GPUs, and TPUs.
- Full support for JIT compilation for increased performance.
- Full support for automatic vectorization for massive parallelization of open-loop and closed-loop architectures.
- Support for SDF models and, upon conversion with [sdformat][sdformat], URDF models.
- Visualization based on the [passive viewer][passive_viewer_mujoco] of Mujoco.
### JaxSim as a simulator
- Wide range of fixed-step explicit Runge-Kutta integrators.
- Support for variable-step integrators implemented as embedded Runge-Kutta schemes.
- Improved stability by optionally integrating the base orientation on the $\text{SO}(3)$ manifold.
- Soft contacts model supporting full friction cone and sticking-slipping transition.
- Collision detection between points rigidly attached to bodies and uneven ground surfaces.
### JaxSim as a multibody dynamics library
- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
- Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
- Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
- Exposes all the necessary quantities to develop controllers in centroidal coordinates.
### JaxSim for robot learning
- Being developed with JAX, all the RBDAs support automatic differentiation both in forward and reverse modes.
- Support for automatically differentiating against kinematics and dynamics parameters.
- All fixed-step integrators are forward and reverse differentiable.
- All variable-step integrators are forward differentiable.
- Ideal for sampling synthetic data for reinforcement learning (RL).
- Ideal for designing physics-informed neural networks (PINNs) with loss functions requiring model-based quantities.
- Ideal for combining model-based control with learning-based components.
[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
> [!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.
## Documentation
The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].
[readthedocs]: https://jaxsim.readthedocs.io/
## Installation
<details>
<summary>With conda</summary>
You can install the project using [`conda`][conda] as follows:
```bash
conda install jaxsim -c conda-forge
```
You can enforce GPU support, if needed, by also specifying `"jaxlib = * = *cuda*"`.
</details>
<details>
<summary>With pip</summary>
You can install the project using [`pypa/pip`][pip], preferably in a [virtual environment][venv], as follows:
```bash
pip install jaxsim
```
Check [`setup.cfg`](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][jax_gpu] of JAX.
</details>
<details>
<summary>Contributors installation</summary>
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 .
```
</details>
[conda]: https://anaconda.org/
[pip]: https://github.com/pypa/pip/
[venv]: https://docs.python.org/3/tutorial/venv.html
[jax_gpu]: https://github.com/google/jax/#installation
## 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 showing the potential of [JAX][jax] in this field.
Brax v2 was later implemented reduced coordinates, following an approach comparable to JaxSim.
The development then shifted to [MJX][mjx], which today 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.
Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous
function of the state $(\mathbf{q}, \boldsymbol{\nu})$, it doesn't require running any optimization algorithm
when stepping the simulation forward.
- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators.
[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
```bibtex
@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},
}
```
## People
| Author | Maintainers |
|:------:|:-----------:|
| [<img src="https://avatars.githubusercontent.com/u/469199?v=4" width="40">][df] | [<img src="https://avatars.githubusercontent.com/u/102977828?v=4" width="40">][ff] [<img src="https://avatars.githubusercontent.com/u/469199?v=4" width="40">][df] |
[df]: https://github.com/diegoferigo
[ff]: https://github.com/flferretti
## License
[BSD3](https://choosealicense.com/licenses/bsd-3-clause/)
Raw data
{
"_id": null,
"home_page": "https://github.com/ami-iit/jaxsim",
"name": "jaxsim",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.11",
"maintainer_email": null,
"keywords": "physics, physics engine, jax, rigid body dynamics, featherstone, reinforcement learning, robot, robotics, sdf, urdf",
"author": "Diego Ferigo",
"author_email": "diego.ferigo@iit.it",
"download_url": "https://files.pythonhosted.org/packages/83/31/eda33be294f66794c5103eec144f21667cd87a1319a8f0d01bf60fb0b5a0/jaxsim-0.2.0.tar.gz",
"platform": "any",
"description": "# JaxSim\n\nJaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.\n\nIts design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence. \n\n## Features\n\n- Physics engine in reduced coordinates supporting fixed-base and floating-base robots.\n- Multibody dynamics library providing all the necessary components for developing model-based control algorithms.\n- Completely developed in Python with [`google/jax`][jax] following a functional programming paradigm.\n- Transparent support for running on CPUs, GPUs, and TPUs.\n- Full support for JIT compilation for increased performance.\n- Full support for automatic vectorization for massive parallelization of open-loop and closed-loop architectures.\n- Support for SDF models and, upon conversion with [sdformat][sdformat], URDF models.\n- Visualization based on the [passive viewer][passive_viewer_mujoco] of Mujoco.\n\n### JaxSim as a simulator\n\n- Wide range of fixed-step explicit Runge-Kutta integrators.\n- Support for variable-step integrators implemented as embedded Runge-Kutta schemes.\n- Improved stability by optionally integrating the base orientation on the $\\text{SO}(3)$ manifold.\n- Soft contacts model supporting full friction cone and sticking-slipping transition.\n- Collision detection between points rigidly attached to bodies and uneven ground surfaces.\n\n### JaxSim as a multibody dynamics library\n\n- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians. \n- Provides all the quantities included in the Euler-Poincar\u00e8 formulation of the equations of motion.\n- Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].\n- Exposes all the necessary quantities to develop controllers in centroidal coordinates.\n\n### JaxSim for robot learning\n\n- Being developed with JAX, all the RBDAs support automatic differentiation both in forward and reverse modes.\n- Support for automatically differentiating against kinematics and dynamics parameters.\n- All fixed-step integrators are forward and reverse differentiable.\n- All variable-step integrators are forward differentiable.\n- Ideal for sampling synthetic data for reinforcement learning (RL).\n- Ideal for designing physics-informed neural networks (PINNs) with loss functions requiring model-based quantities.\n- Ideal for combining model-based control with learning-based components.\n\n[jax]: https://github.com/google/jax/\n[sdformat]: https://github.com/gazebosim/sdformat\n[notation]: https://research.tue.nl/en/publications/multibody-dynamics-notation-version-2\n[passive_viewer_mujoco]: https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer\n\n> [!WARNING]\n> This project is still experimental, APIs could change between releases without notice.\n\n> [!NOTE]\n> JaxSim currently focuses on locomotion applications.\n> Only contacts between bodies and smooth ground surfaces are supported.\n\n## Documentation\n\nThe JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].\n\n[readthedocs]: https://jaxsim.readthedocs.io/\n\n## Installation\n\n<details>\n<summary>With conda</summary>\n\nYou can install the project using [`conda`][conda] as follows:\n\n```bash\nconda install jaxsim -c conda-forge\n```\n\nYou can enforce GPU support, if needed, by also specifying `\"jaxlib = * = *cuda*\"`.\n\n</details>\n\n<details>\n<summary>With pip</summary>\n\nYou can install the project using [`pypa/pip`][pip], preferably in a [virtual environment][venv], as follows:\n\n```bash\npip install jaxsim\n```\n\nCheck [`setup.cfg`](setup.cfg) for the complete list of optional dependencies.\nYou can obtain a full installation using `jaxsim[all]`.\n\nIf you need GPU support, follow the official [installation instructions][jax_gpu] of JAX.\n\n</details>\n\n<details>\n<summary>Contributors installation</summary>\n\nIf you want to contribute to the project, we recommend creating the following `jaxsim` conda environment first:\n\n```bash\nconda env create -f environment.yml\n```\n\nThen, activate the environment and install the project in editable mode:\n\n```bash\nconda activate jaxsim\npip install --no-deps -e .\n```\n\n</details>\n\n[conda]: https://anaconda.org/\n[pip]: https://github.com/pypa/pip/\n[venv]: https://docs.python.org/3/tutorial/venv.html\n[jax_gpu]: https://github.com/google/jax/#installation\n\n## Credits\n\nThe RBDAs are based on the theory of the [Rigid Body Dynamics Algorithms][RBDA]\nbook by Roy Featherstone.\nThe algorithms and some simulation features were inspired by its accompanying [code][spatial_v2].\n\n[RBDA]: https://link.springer.com/book/10.1007/978-1-4899-7560-7\n[spatial_v2]: http://royfeatherstone.org/spatial/index.html#spatial-software\n\nThe development of JaxSim started in late 2021, inspired by early versions of [`google/brax`][brax].\nAt that time, Brax was implemented in maximal coordinates, and we wanted a physics engine in reduced coordinates.\nWe are grateful to the Brax team for their work and showing the potential of [JAX][jax] in this field.\n\nBrax v2 was later implemented reduced coordinates, following an approach comparable to JaxSim.\nThe development then shifted to [MJX][mjx], which today provides a JAX-based implementation of the Mujoco APIs.\n\nThe main differences between MJX/Brax and JaxSim are as follows:\n\n- JaxSim supports out-of-the-box all SDF models with [Pose Frame Semantics][PFS].\n- JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.\n Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous \n function of the state $(\\mathbf{q}, \\boldsymbol{\\nu})$, it doesn't require running any optimization algorithm\n when stepping the simulation forward.\n- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators. \n\n[brax]: https://github.com/google/brax\n[mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html\n[PFS]: http://sdformat.org/tutorials?tut=pose_frame_semantics\n\n## Contributing\n\nWe welcome contributions from the community.\nPlease read the [contributing guide](./CONTRIBUTING.md) to get started.\n\n## Citing\n\n```bibtex\n@software{ferigo_jaxsim_2022,\n author = {Diego Ferigo and Filippo Luca Ferretti and Silvio Traversaro and Daniele Pucci},\n title = {{JaxSim}: A Differentiable Physics Engine and Multibody Dynamics Library for Control and Robot Learning},\n url = {http://github.com/ami-iit/jaxsim},\n year = {2022},\n}\n```\n\n## People\n\n| Author | Maintainers |\n|:------:|:-----------:|\n| [<img src=\"https://avatars.githubusercontent.com/u/469199?v=4\" width=\"40\">][df] | [<img src=\"https://avatars.githubusercontent.com/u/102977828?v=4\" width=\"40\">][ff] [<img src=\"https://avatars.githubusercontent.com/u/469199?v=4\" width=\"40\">][df] |\n\n[df]: https://github.com/diegoferigo\n[ff]: https://github.com/flferretti\n\n## License\n\n[BSD3](https://choosealicense.com/licenses/bsd-3-clause/)\n",
"bugtrack_url": null,
"license": "BSD",
"summary": "A differentiable physics engine and multibody dynamics library for control and robot learning.",
"version": "0.2.0",
"project_urls": {
"Changelog": "https://github.com/ami-iit/jaxsim/releases",
"Documentation": "https://jaxsim.readthedocs.io",
"Homepage": "https://github.com/ami-iit/jaxsim",
"Source": "https://github.com/ami-iit/jaxsim",
"Tracker": "https://github.com/ami-iit/jaxsim/issues"
},
"split_keywords": [
"physics",
" physics engine",
" jax",
" rigid body dynamics",
" featherstone",
" reinforcement learning",
" robot",
" robotics",
" sdf",
" urdf"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "b4420567d3bdf7e7ce360269ddaf6a5403306b8a3cc54f4f2284fcf8cc266f29",
"md5": "ec8d11b6d86d77e733ab325eb7d60695",
"sha256": "78118738b5a6f858ede927f9fedbdb5f859cc92318481fdf58c3fe426ec721b3"
},
"downloads": -1,
"filename": "jaxsim-0.2.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "ec8d11b6d86d77e733ab325eb7d60695",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.11",
"size": 119352,
"upload_time": "2024-04-03T15:56:40",
"upload_time_iso_8601": "2024-04-03T15:56:40.747551Z",
"url": "https://files.pythonhosted.org/packages/b4/42/0567d3bdf7e7ce360269ddaf6a5403306b8a3cc54f4f2284fcf8cc266f29/jaxsim-0.2.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "8331eda33be294f66794c5103eec144f21667cd87a1319a8f0d01bf60fb0b5a0",
"md5": "32f962c7deb97b5628e928fd0600d7e0",
"sha256": "4247148a46875541d26f07eb49b95ed6ba755a0eb266e7eaca463b9c7d874d90"
},
"downloads": -1,
"filename": "jaxsim-0.2.0.tar.gz",
"has_sig": false,
"md5_digest": "32f962c7deb97b5628e928fd0600d7e0",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.11",
"size": 182817,
"upload_time": "2024-04-03T15:56:43",
"upload_time_iso_8601": "2024-04-03T15:56:43.470992Z",
"url": "https://files.pythonhosted.org/packages/83/31/eda33be294f66794c5103eec144f21667cd87a1319a8f0d01bf60fb0b5a0/jaxsim-0.2.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-04-03 15:56:43",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "ami-iit",
"github_project": "jaxsim",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "jaxsim"
}
JSON
