Skip to main content
Pre-release

This release is a pre-release and may not be stable for production use.

JAXsim

A scalable physics engine implemented with JAX. With JIT batteries 🔋

⚠ This project is still experimental, APIs could change without notice. ️⚠

⚠ This simulator currently focuses on locomotion applications. Only contacts with ground are supported. ️⚠

Features

  • Physics engine in reduced coordinates implemented with JAX in Python.
  • Supported JIT compilation of Python code for increased performance.
  • Transparent support to execute the simulation on CPUs, GPUs, and TPUs.
  • Possibility to run parallel multi-body simulations on hardware accelerators for significantly increased throughput.
  • Support of SDF models (and, upon conversion, URDF models).
  • Collision detection between bodies and uneven ground surface.
  • Continuous soft contacts model with no friction cone approximations.
  • Full support of inertial properties of bodies.
  • Revolute, prismatic, and fixed joints support.
  • Integrators: forward Euler, semi-implicit Euler, Runge-Kutta 4.
  • High-level classes to compute multi-body dynamics quantities from simulation state.
  • High-level classes supporting both object-oriented and functional programming.
  • Optional validation of JAX pytrees to prevent JIT re-compilation.

Planned features:

  • Reinforcement Learning module developed in JAX.
  • Finalization of differentiable physics through the simulation.

Installation

You can install the project with pypa/pip, preferably in a virtual environment:

pip install jaxsim

Have a look to setup.cfg for a complete list of optional dependencies. You can install all of them by specifying jaxsim[all].

Note: if you need GPU support, please 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 authored 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 of developing JAXsim stems from google/brax. Here below we summarize the differences between the projects:

  • JAXsim simulates multibody dynamics in reduced coordinates, while brax uses maximal coordinates.
  • 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, under the assumption that 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.
  • While supported thanks to the usage of JAX, differentiating through the simulator has not yet been studied.

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 for Control and Robot Learning},
  url = {http://github.com/ami-iit/jaxsin},
  year = {2022},
}

Maintainers

@diegoferigo

License

BSD3

Metadata

Release files for jaxsim 0.1.dev149

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for jaxsim 0.1.dev149
File Size Uploaded
jaxsim-0.1.dev149.tar.gz 55.8 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for jaxsim 0.1.dev149
File Interpreter ABI Platform
jaxsim-0.1.dev149-py3-none-any.whl Python 3 none any Details

Total release size: 124.3 kB

Release files / jaxsim-0.1.dev149.tar.gz

Download URL jaxsim-0.1.dev149.tar.gz
Size 55.8 kB
Tags Source
SHA-256 checksum
How to use checksums
cda2fc61e0a6364d7d7fd05575b0090a4abc3b03581a60cd0807d5df2ada65d7
BLAKE2b-256 checksum
How to use checksums
8aa819a5a166b0fc6962dfb7c75c64c9d5a23eba0521ea6fdd3d5fc9e5e86e47
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/4.0.1 CPython/3.11.3

Release files / jaxsim-0.1.dev149-py3-none-any.whl

Download URL jaxsim-0.1.dev149-py3-none-any.whl
Size 68.4 kB
Tags Python 3
SHA-256 checksum
How to use checksums
ec23cbc9a4603dd491589f55b6e86fa743829080c385c1fe75f9241a92d8048c
BLAKE2b-256 checksum
How to use checksums
dccf9f5a9e7ea02713c273aa0a3a45d55d85dd0625fe931b0a2dd9577ad937d0
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/4.0.1 CPython/3.11.3

Release history Release notifications | RSS feed

0.9.0

2 release files

0.8.2

2 release files

0.8.1

2 release files

0.8.0

2 release files

0.7.0

2 release files

0.6.1

2 release files

0.6.0

2 release files

0.5.0

2 release files

0.4.2

2 release files

0.4.1

2 release files

0.4.0

2 release files

0.3.0

2 release files

0.2.0

2 release files

0.1

2 release files

This release

0.1.dev149 This release

2 release files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page