Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added usage section in ReadMe and basic example for simulator #314

Open
wants to merge 13 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 100 additions & 37 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
# 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.
**JaxSim** is a **differentiable physics engine** and **multibody dynamics library** built with JAX, tailored for control and robotic learning applications.

<div align="center">
<br/>
Expand All @@ -17,41 +15,105 @@ Its design facilitates research and accelerates prototyping in the intersection
</div>

## Features
- Reduced-coordinate physics engine for **fixed-base** and **floating-base** robots.
- Multibody dynamics library for model-based control algorithms.
- Fully Python-based, leveraging [jax][jax] following a functional programming paradigm.
- Seamless execution on CPUs, GPUs, and TPUs.
- Supports JIT compilation and automatic vectorization for high performance.
- Compatible with SDF models and URDF (via [sdformat][sdformat] conversion).

## Usage

### Using JaxSim as simulator


```python
import jax
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib

# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')
ndof = len(joints)

# Build and reduce the model
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.0001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)

# Initialize data and simulation
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)
T = jnp.arange(start=0, stop=1.0, step=model.time_step)
tau = jnp.zeros(ndof)

# Simulate
for t in T:
data, _ = js.model.step(model=model, data=data, link_forces=None, joint_force_references=tau)
Comment on lines +30 to +63
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think about compressing the example code in a collapsible section?
I also removed an unused import and reformatted

Suggested change
```python
import jax
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib
# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')
ndof = len(joints)
# Build and reduce the model
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.0001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)
# Initialize data and simulation
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)
T = jnp.arange(start=0, stop=1.0, step=model.time_step)
tau = jnp.zeros(ndof)
# Simulate
for t in T:
data, _ = js.model.step(model=model, data=data, link_forces=None, joint_force_references=tau)
<details>
<summary><b>Using JaxSim as a simulator</b></summary>
```python
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib
# Load the iCub model.
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')
# Build and reduce the model.
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)
# Initialize data and simulation.
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)
T = jnp.arange(start=0, stop=1.0, step=model.time_step)
tau = jnp.zeros(model.dofs())
# Simulate.
for _t in T:
data, _ = js.model.step(
model=model, data=data, link_forces=None, joint_force_references=tau
)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SInce the example are quite short I would prefer to have them directly in the readme, so that it is directly visible that it is possible to use jaxsim within few lines of code, what do you think ?


- 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.
- Supports running open-loop and full closed-loop control architectures on hardware accelerators.
### Using JaxSim as a multibody dynamics library
``` python
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib

# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')

# Build and reduce the model
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.0001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)

# Initialize model data
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)

# Frame and dynamics computations
frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")
W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index) # Frame transformation
W_J_F = js.frame.jacobian(model=model, data=data, frame_index=frame_index) # Frame Jacobian

# Dynamics properties
M = js.model.free_floating_mass_matrix(model=model, data=data) # Mass matrix
h = js.model.free_floating_bias_forces(model=model, data=data) # Bias forces
g = js.model.free_floating_gravity_forces(model=model, data=data) # Gravity forces
C = js.model.free_floating_coriolis_matrix(model=model, data=data) # Coriolis matrix

# Print dynamics results
print(f"M: shape={M.shape}, h: shape={h.shape}, g: shape={g.shape}, C: shape={C.shape}")
Comment on lines +68 to +107
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Collapsible section as above. Reformatted and used f-string syntax to print the shapes

Suggested change
``` python
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib
# Load the iCub model
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')
# Build and reduce the model
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.0001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)
# Initialize model data
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)
# Frame and dynamics computations
frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")
W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index) # Frame transformation
W_J_F = js.frame.jacobian(model=model, data=data, frame_index=frame_index) # Frame Jacobian
# Dynamics properties
M = js.model.free_floating_mass_matrix(model=model, data=data) # Mass matrix
h = js.model.free_floating_bias_forces(model=model, data=data) # Bias forces
g = js.model.free_floating_gravity_forces(model=model, data=data) # Gravity forces
C = js.model.free_floating_coriolis_matrix(model=model, data=data) # Coriolis matrix
# Print dynamics results
print(f"M: shape={M.shape}, h: shape={h.shape}, g: shape={g.shape}, C: shape={C.shape}")
<details>
<summary><b>Using JaxSim as a multibody dynamics library</b></summary>
``` python
import jax.numpy as jnp
import jaxsim.api as js
import icub_models
import pathlib
# Load the iCub model.
model_path = icub_models.get_model_file("iCubGazeboV2_5")
joints = ('torso_pitch', 'torso_roll', 'torso_yaw', 'l_shoulder_pitch',
'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow', 'r_shoulder_pitch',
'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow', 'l_hip_pitch',
'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',
'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch',
'r_ankle_roll')
# Build and reduce the model.
model_description = pathlib.Path(model_path)
full_model = js.model.JaxSimModel.build_from_model_description(
model_description=model_description, time_step=0.0001, is_urdf=True
)
model = js.model.reduce(model=full_model, considered_joints=joints)
# Initialize model data.
data = js.data.JaxSimModelData.zero(model=model).reset_base_position(
base_position=jnp.array([0.0, 0.0, 1.0])
)
# Frame and dynamics computations.
frame_index = js.frame.name_to_idx(model=model, frame_name="l_foot")
W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index) # Frame transformation
W_J_F = js.frame.jacobian(model=model, data=data, frame_index=frame_index) # Frame Jacobian
# Dynamics properties.
M = js.model.free_floating_mass_matrix(model=model, data=data) # Mass matrix
h = js.model.free_floating_bias_forces(model=model, data=data) # Bias forces
g = js.model.free_floating_gravity_forces(model=model, data=data) # Gravity forces
C = js.model.free_floating_coriolis_matrix(model=model, data=data) # Coriolis matrix
# Print dynamics results.
print(f"{M.shape=} \t {h.shape=} \t {g.shape=} \t {C.shape=}")
```

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comments above about the collapse


### JaxSim for robot learning
```
### Additional features

- Being developed with JAX, all the RBDAs support automatic differentiation both in forward and reverse modes.
- Full support for automatic differentiation of RBDAs (forward and reverse modes) with JAX.
- 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.
- Check the example folder for additional usecase !

[jax]: https://github.com/google/jax/
[sdformat]: https://github.com/gazebosim/sdformat
Expand All @@ -65,12 +127,6 @@ Its design facilitates research and accelerates prototyping in the intersection
> 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>
Expand Down Expand Up @@ -142,6 +198,13 @@ pip install --no-deps -e .
[venv]: https://docs.python.org/3/tutorial/venv.html
[jax_gpu]: https://github.com/google/jax/#installation

## Documentation

The JaxSim API documentation is available at [jaxsim.readthedocs.io][readthedocs].

[readthedocs]: https://jaxsim.readthedocs.io/


## Overview

<details>
Expand Down
2 changes: 2 additions & 0 deletions environment.yml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ dependencies:
- pytest
- pytest-icdiff
- robot_descriptions
- icub-models
# [viz]
- lxml
- mediapy
Expand All @@ -55,6 +56,7 @@ dependencies:
- sphinx-multiversion
- sphinx_rtd_theme
- sphinx-toolbox
- icub-models
# ========================================
# Other dependencies for GitHub Codespaces
# ========================================
Expand Down
2 changes: 2 additions & 0 deletions examples/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,13 @@ This folder contains Jupyter notebooks that demonstrate the practical usage of J
| :--- | :---: | :--- |
| [`jaxsim_as_multibody_dynamics_library`](./jaxsim_as_multibody_dynamics_library.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_as_multibody_dynamics] | An example demonstrating how to use JaxSim as a multibody dynamics library. |
| [`jaxsim_as_physics_engine.ipynb`](./jaxsim_as_physics_engine.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_as_physics_engine] | An example demonstrating how to simulate vectorized models in parallel. |
| [`jaxsim_as_physics_engine_advanced.ipynb`](./jaxsim_as_physics_engine_advanced.ipynb) | [![Open In Colab][colab_badge]][jaxsim_as_physics_engine_advanced] | An example showcasing advanced JaxSim usage, such as customizing the integrator, contact model, and more. |
| [`jaxsim_for_robot_controllers.ipynb`](./jaxsim_for_robot_controllers.ipynb) | [![Open In Colab][colab_badge]][ipynb_jaxsim_closed_loop] | A basic example showing how to simulate a PD controller with gravity compensation for a 2-DOF cart-pole. |

[colab_badge]: https://colab.research.google.com/assets/colab-badge.svg
[ipynb_jaxsim_closed_loop]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_for_robot_controllers.ipynb
[ipynb_jaxsim_as_physics_engine]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_physics_engine.ipynb
[jaxsim_as_physics_engine_advanced]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_physics_engine_advanced.ipynb
[ipynb_jaxsim_as_multibody_dynamics]: https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_multibody_dynamics_library.ipynb

## How to run the examples
Expand Down
Loading