This release is a pre-release and may not be stable for production use.
JAXsim
A scalable physics engine and multibody dynamics library implemented with JAX. With JIT batteries 🔋
[!WARNING] This project is still experimental, APIs could change without notice.
[!NOTE] This simulator currently focuses on locomotion applications. Only contacts with ground are supported.
Features
- Physics engine in reduced coordinates implemented with JAX in Python.
- JIT compilation of Python code for increased performance.
- Transparent support to execute logic on CPUs, GPUs, and TPUs.
- Parallel multi-body simulations on hardware accelerators for significantly increased throughput.
- Support for SDF models (and, upon conversion, URDF models).
- Collision detection between bodies and uneven ground surface.
- Soft contacts model supporting full friction cone and sticking / slipping transition.
- Complete support for inertial properties of rigid bodies.
- Revolute, prismatic, and fixed joints support.
- Integrators: forward Euler, semi-implicit Euler, Runge-Kutta 4.
- High-level classes for object-oriented programming.
- High-level classes to compute multi-body dynamics quantities from the simulation state.
- High-level classes wrapping the low-level functional RBDAs with support of multiple velocities representations.
- Default validation of JAX pytrees to prevent JIT re-compilations.
- Preliminary support for automatic differentiation of RBDAs.
Documentation
The JAXsim API documentation is available at kaxsim.readthedocs.io.
Installation
You can install the project with pypa/pip, preferably in a virtual environment:
pip install jaxsim
Check setup.cfg for the complete list of optional dependencies.
Install all of them with jaxsim[all].
Note: For GPU support, follow the official installation instruction of JAX.
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.
Contributing
Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
Citing
@software{ferigo_jaxsim_2022,
author = {Diego Ferigo and Silvio Traversaro and Daniele Pucci},
title = {{JAXsim}: A Physics Engine in Reduced Coordinates and Multibody Dynamics Library for Control and Robot Learning},
url = {http://github.com/ami-iit/jaxsim},
year = {2022},
}
People
| Author | Maintainers |
|---|---|
 |
 |
License
2 maintainers
Metadata
2 maintainers
Release files for jaxsim 0.1.dev388
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jaxsim-0.1.dev388.tar.gz | 150.8 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jaxsim-0.1.dev388-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 245.3 kB
Release files / jaxsim-0.1.dev388.tar.gz
| Download URL | jaxsim-0.1.dev388.tar.gz |
|---|---|
| Size | 150.8 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
023507c2a01852263486d1b5f77345a5cd4828a03b5c0cc0e58c40d5401550be
|
|
BLAKE2b-256 checksum How to use checksums |
1c7d812c2c3c5fb1c8c90edf8edac1421f298c7cf3fe60f9341d4d80f6dd070b
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.11.7
|
Release files / jaxsim-0.1.dev388-py3-none-any.whl
| Download URL | jaxsim-0.1.dev388-py3-none-any.whl |
|---|---|
| Size | 94.6 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
165574234985763a8d2d7635dd1ed6490123ffe67bf16b7c60747a244405b536
|
|
BLAKE2b-256 checksum How to use checksums |
68a873191e3af30d3086b4e5b2b6f6c3a2b392946f99959c6c9a87724471310f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.2 CPython/3.11.7
|
