diff --git a/README.md b/README.md index 9bcfc3c6e..a8ff58edd 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,33 @@ # JAXsim -**A scalable physics engine implemented with JAX. With JIT batteries 🔋** +**A scalable physics engine and multibody dynamics library implemented with JAX. With JIT batteries 🔋** -⚠ This project is still experimental, APIs could change without notice. ️⚠ +> [!WARNING] +> This project is still experimental, APIs could change without notice. -⚠ This simulator currently focuses on locomotion applications. Only contacts with ground are supported. ️⚠ +> [!NOTE] +> This simulator currently focuses on locomotion applications. Only contacts with ground are supported. ## Features - Physics engine in reduced coordinates implemented with [JAX][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). +- 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. -- Continuous soft contacts model with no friction cone approximations. -- Full support of inertial properties of bodies. +- 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 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. +- High-level classes wrapping the low-level functional RBDAs with support of [multiple velocities representations][notation]. +- Default validation of JAX pytrees to prevent JIT re-compilations. +- Preliminary support for automatic differentiation of RBDAs. [jax]: https://github.com/google/jax/ +[notation]: https://research.tue.nl/en/publications/multibody-dynamics-notation-version-2 ## Installation @@ -37,37 +37,39 @@ You can install the project with [`pypa/pip`][pip], preferably in a [virtual env pip install jaxsim ``` -Have a look to [`setup.cfg`](setup.cfg) for a complete list of optional dependencies. -You can install all of them by specifying `jaxsim[all]`. +Check [`setup.cfg`](setup.cfg) for the complete list of optional dependencies. +Install all of them with `jaxsim[all]`. -**Note:** if you need GPU support, please follow the official [installation instruction][jax_gpu] of JAX. +**Note:** For GPU support, follow the official [installation instruction][jax_gpu] of JAX. [pip]: https://github.com/pypa/pip/ -[venv]: https://docs.python.org/3.8/tutorial/venv.html +[venv]: https://docs.python.org/3/tutorial/venv.html [jax_gpu]: https://github.com/google/jax/#installation ## Credits The physics module of JAXsim is based on the theory of the [Rigid Body Dynamics Algorithms][RBDA] -book authored by Roy Featherstone. +book by Roy Featherstone. We structured part of our logic following its accompanying [code][spatial_v2]. The physics engine is developed entirely in Python using [JAX][jax]. [RBDA]: https://link.springer.com/book/10.1007/978-1-4899-7560-7 [spatial_v2]: http://royfeatherstone.org/spatial/index.html#spatial-software -The inspiration of developing JAXsim stems from [`google/brax`][brax]. +The inspiration for developing JAXsim originally stemmed from early versions of [`google/brax`][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è +- JAXsim simulates multibody dynamics in reduced coordinates, while brax v1 uses maximal coordinates. +- The new v2 APIs of brax (and the new [MJX][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, under the assumption that the model is described with the +- JAXsim supports SDF (and, indirectly, URDF) models, assuming the model is described with the recent [Pose Frame Semantics][PFS]. -- 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. +- 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 being thoroughly tested. [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 @@ -80,18 +82,20 @@ For major changes, please open an issue first to discuss what you would like to ```bibtex @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}, + 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}, } ``` -## Maintainers +## People -| [][df] | [@diegoferigo][df] | -|:---------------------------------------------------------------:|:------------------:| +| Author | Maintainers | +|:------:|:-----------:| +| [][df] | [][ff] [][df] | [df]: https://github.com/diegoferigo +[ff]: https://github.com/flferretti ## License diff --git a/src/jaxsim/simulation/simulator.py b/src/jaxsim/simulation/simulator.py index 6b7e9b77c..fa96c3650 100644 --- a/src/jaxsim/simulation/simulator.py +++ b/src/jaxsim/simulation/simulator.py @@ -63,7 +63,7 @@ class JaxSim(Vmappable): """The JaxSim simulator.""" # Step size stored in ns in order to prevent floats approximation - step_size_ns: jtp.Int = dataclasses.field( + step_size_ns: Static[jtp.Int] = dataclasses.field( default_factory=lambda: jnp.array(1_000_000, dtype=jnp.uint64) ) diff --git a/src/jaxsim/typing.py b/src/jaxsim/typing.py index 82355cf33..94b9508be 100644 --- a/src/jaxsim/typing.py +++ b/src/jaxsim/typing.py @@ -1,7 +1,6 @@ from typing import Any, Dict, Hashable, List, NamedTuple, Tuple, Union import jax.numpy as jnp -import numpy as np import numpy.typing as npt # JAX types @@ -35,6 +34,6 @@ Tensor = Union[npt.NDArray, ArrayJax] Vector = Array Matrix = Array -Bool = bool +Bool = Union[bool, ArrayJax] Int = Union[int, IntJax] Float = Union[float, FloatJax]