Merge branch 'ami-iit:main' into motor_dynamics

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
 
-| [<img src="https://github.com/diegoferigo.png" width="40">][df] | [@diegoferigo][df] |
-|:---------------------------------------------------------------:|:------------------:|
+| 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
 

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)
     )
 

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]