condapixipipconda)pixi)
][df] [
][ff] | [][ff] |
[df]: https://github.com/diegoferigo
[ff]: https://github.com/flferretti
## License
[BSD3](https://choosealicense.com/licenses/bsd-3-clause/)
================================================
FILE: docs/Makefile
================================================
SPHINXOPTS ?=
SPHINXBUILD ?= sphinx-build
SOURCEDIR = .
BUILDDIR = _build
SPHINXPROJ = JAXsim
# Put it first so that "make" without argument is like "make help".
help:
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
.PHONY: help Makefile
# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
================================================
FILE: docs/conf.py
================================================
# Configuration file for the Sphinx documentation builder.
import os
import sys
if os.environ.get("READTHEDOCS"):
checkout_name = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
os.environ["CONDA_PREFIX"] = os.path.realpath(
os.path.join("..", "..", "conda", checkout_name)
)
import jaxsim
# -- Version information
sys.path.insert(0, os.path.abspath("."))
sys.path.insert(0, os.path.abspath("../"))
sys.path.insert(0, os.path.abspath("../../"))
module_path = os.path.abspath("../src/")
sys.path.insert(0, module_path)
__version__ = jaxsim._version.__version__
# -- Project information
project = "JAXsim"
copyright = "2022, Artificial and Mechanical Intelligence"
author = "Artificial and Mechanical Intelligence"
release = version = __version__
# -- General configuration
extensions = [
"sphinx.ext.duration",
"sphinx.ext.doctest",
"sphinx.ext.autodoc",
"sphinx.ext.autosummary",
"sphinx.ext.intersphinx",
"sphinx.ext.mathjax",
"sphinx.ext.ifconfig",
"sphinx.ext.viewcode",
"sphinx_rtd_theme",
"sphinx.ext.napoleon",
"sphinx_autodoc_typehints",
"sphinx_multiversion",
"myst_nb",
"sphinx_gallery.gen_gallery",
"sphinxcontrib.collections",
"sphinx_design",
]
# -- Options for intersphinx extension
language = "en"
html_theme = "sphinx_book_theme"
templates_path = ["_templates"]
html_title = f"JAXsim {version}"
master_doc = "index"
autodoc_typehints_format = "short"
autodoc_typehints = "description"
autosummary_generate = True
epub_show_urls = "footnote"
# Enable postponed evaluation of annotations (PEP 563)
autodoc_type_aliases = {
"jaxsim.typing.PyTree": "jaxsim.typing.PyTree",
"jaxsim.typing.Vector": "jaxsim.typing.Vector",
"jaxsim.typing.Matrix": "jaxsim.typing.Matrix",
"jaxsim.typing.Array": "jaxsim.typing.Array",
"jaxsim.typing.Int": "jaxsim.typing.Int",
"jaxsim.typing.Bool": "jaxsim.typing.Bool",
"jaxsim.typing.Float": "jaxsim.typing.Float",
"jaxsim.typing.ScalarLike": "jaxsim.typing.ScalarLike",
"jaxsim.typing.ArrayLike": "jaxsim.typing.ArrayLike",
"jaxsim.typing.VectorLike": "jaxsim.typing.VectorLike",
"jaxsim.typing.MatrixLike": "jaxsim.typing.MatrixLike",
"jaxsim.typing.IntLike": "jaxsim.typing.IntLike",
"jaxsim.typing.BoolLike": "jaxsim.typing.BoolLike",
"jaxsim.typing.FloatLike": "jaxsim.typing.FloatLike",
}
# -- Options for sphinx-collections
collections = {
"examples": {"driver": "copy_folder", "source": "../examples/", "ignore": "assets"}
}
# -- Options for sphinx-gallery ----------------------------------------------
sphinx_gallery_conf = {
"examples_dirs": "../examples",
"gallery_dirs": "../generated_examples/",
"doc_module": "jaxsim",
}
# -- Options for myst -------------------------------------------------------
myst_enable_extensions = [
"amsmath",
"dollarmath",
]
nb_execution_mode = "auto"
nb_execution_raise_on_error = True
nb_render_image_options = {
"scale": "60",
}
nb_execution_timeout = 180
source_suffix = [".rst", ".md", ".ipynb"]
# Ignore header warnings
suppress_warnings = ["myst.header"]
================================================
FILE: docs/examples.rst
================================================
.. _collections:
Example Notebooks
=================
.. toctree::
:glob:
:hidden:
:maxdepth: 1
_collections/examples/README.md
.. raw:: html
![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rcEwprINtwZ3"
},
"source": [
"## Prepare environment\n",
"\n",
"First, we need to install the necessary packages and import their resources."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "u4xL7dbBtwZ3",
"outputId": "1a088e28-e005-4910-928c-cb641e589ab5"
},
"outputs": [],
"source": [
"# @title Imports and setup\n",
"from IPython.display import clear_output\n",
"import sys\n",
"\n",
"IS_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"# Install JAX, sdformat, and other notebook dependencies.\n",
"if IS_COLAB:\n",
" !{sys.executable} -m pip install --pre -qU jaxsim\n",
" !{sys.executable} -m pip install robot_descriptions>=1.16.0\n",
" !apt install -qq lsb-release wget gnupg\n",
" !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
" !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
" !apt -qq update\n",
" !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
"\n",
" clear_output()\n",
"\n",
"import os\n",
"import pathlib\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import jaxsim.api as js\n",
"import jaxsim.math\n",
"from jaxsim import logging\n",
"from jaxsim import VelRepr\n",
"\n",
"logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
"print(f\"Running on {jax.devices()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fN8Xg4QgtwZ4"
},
"source": [
"## Robot model\n",
"\n",
"JaxSim allows loading robot descriptions from both [SDF][sdformat] and [URDF][urdf] files.\n",
"\n",
"In this example, we will use the [ErgoCub][ergocub] humanoid robot model. If you have a URDF/SDF file for your robot that is compatible with [`gazebosim/sdformat`][sdformat_github][1], it should work out-of-the-box with JaxSim.\n",
"\n",
"[sdformat]: http://sdformat.org/\n",
"[urdf]: http://wiki.ros.org/urdf/\n",
"[ergocub]: https://ergocub.eu/\n",
"[sdformat_github]: https://github.com/gazebosim/sdformat\n",
"\n",
"---\n",
"\n",
"[1]: JaxSim validates robot descriptions using the command `gz sdf -p /path/to/file.urdf`. Ensure this command runs successfully on your file.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rB0BFxyPtwZ5"
},
"outputs": [],
"source": [
"# @title Fetch the URDF file\n",
"\n",
"try:\n",
" os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.7\"\n",
"\n",
" import robot_descriptions.ergocub_description\n",
"\n",
"finally:\n",
" _ = os.environ.pop(\"ROBOT_DESCRIPTION_COMMIT\", None)\n",
"\n",
"model_description_path = pathlib.Path(\n",
" robot_descriptions.ergocub_description.URDF_PATH.replace(\n",
" \"ergoCubSN002\", \"ergoCubSN001\"\n",
" )\n",
")\n",
"\n",
"clear_output()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "jeTUZic8twZ5"
},
"source": [
"### Create the model and its data\n",
"\n",
"The dynamics of a generic floating-base model are governed by the following equations of motion:\n",
"\n",
"$$\n",
"M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + \\sum_{L_i \\in \\mathcal{L}} J_{W,L_i}^\\top(\\mathbf{q}) \\: \\mathbf{f}_i\n",
".\n",
"$$\n",
"\n",
"Here, the system state is represented by:\n",
"\n",
"- $\\mathbf{q} = ({}^W \\mathbf{p}_B, \\, \\mathbf{s}) \\in \\text{SE}(3) \\times \\mathbb{R}^n$ is the generalized position.\n",
"- $\\boldsymbol{\\nu} = (\\boldsymbol{v}_{W,B}, \\, \\boldsymbol{\\omega}_{W,B}, \\, \\dot{\\mathbf{s}}) \\in \\mathbb{R}^{6+n}$ is the generalized velocity.\n",
"\n",
"The inputs to the system are:\n",
"\n",
"- $\\boldsymbol{\\tau} \\in \\mathbb{R}^n$ are the joint torques.\n",
"- $\\mathbf{f}_i \\in \\mathbb{R}^6$ is the 6D force applied to the link $L_i$.\n",
"\n",
"JaxSim exposes functional APIs to operate over the following two main data structures:\n",
"\n",
"- **`JaxSimModel`** stores all the constant information parsed from the model description.\n",
"- **`JaxSimModelData`** holds the state of model.\n",
"\n",
"Additionally, JaxSim includes a utility class, **`JaxSimModelReferences`**, for managing and manipulating system inputs.\n",
"\n",
"---\n",
"\n",
"This notebook uses the notation summarized in the following report. Please refer to this document if you have any questions or if something is unclear.\n",
"\n",
"> Traversaro and Saccon, **Multibody dynamics notation**, 2019, [URL](https://pure.tue.nl/ws/portalfiles/portal/139293126/A_Multibody_Dynamics_Notation_Revision_2_.pdf)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WYgBAxU0twZ6"
},
"outputs": [],
"source": [
"# Create the model from the model description.\n",
"# JaxSim removes all fixed joints by lumping together their parent and child links.\n",
"full_model = js.model.JaxSimModel.build_from_model_description(\n",
" model_description=model_description_path\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DdaETmDStwZ6"
},
"source": [
"It is often useful to work with only a subset of joints, referred to as the _considered joints_. JaxSim allows to reduce a model so that the computation of the rigid body dynamics quantities is simplified.\n",
"\n",
"By default, the positions of the removed joints are considered to be zero. If this is not the case, the `reduce` function accepts a dictionary `dict[str, float]` to specify custom joint positions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QuhG7Zv5twZ7"
},
"outputs": [],
"source": [
"model = js.model.reduce(\n",
" model=full_model,\n",
" considered_joints=tuple(\n",
" j\n",
" for j in full_model.joint_names()\n",
" # Remove sensor joints.\n",
" if \"camera\" not in j\n",
" # Remove head and hands.\n",
" and \"neck\" not in j\n",
" and \"wrist\" not in j\n",
" and \"thumb\" not in j\n",
" and \"index\" not in j\n",
" and \"middle\" not in j\n",
" and \"ring\" not in j\n",
" and \"pinkie\" not in j\n",
" # Remove upper body.\n",
" and \"torso\" not in j and \"elbow\" not in j and \"shoulder\" not in j\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "RLvAit_i2ZiA",
"outputId": "ea3954af-b9b9-46ac-d9cb-20b99b1eac94"
},
"outputs": [],
"source": [
"# Print model quantities.\n",
"print(f\"Model name: {model.name()}\")\n",
"print(f\"Number of links: {model.number_of_links()}\")\n",
"print(f\"Number of joints: {model.number_of_joints()}\")\n",
"\n",
"print()\n",
"print(f\"Links:\\n{model.link_names()}\")\n",
"\n",
"print()\n",
"print(f\"Joints:\\n{model.joint_names()}\")\n",
"\n",
"print()\n",
"print(f\"Frames:\\n{model.frame_names()}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Xp8V5on5twZ8",
"outputId": "cc1564db-ae91-4dba-92c9-b8b87bd65f10"
},
"outputs": [],
"source": [
"# Create a random data object from the reduced model.\n",
"data = js.data.random_model_data(model=model)\n",
"\n",
"# Print the default state.\n",
"W_H_B, s = data.generalized_position\n",
"ν = data.generalized_velocity\n",
"\n",
"print(f\"W_H_B: shape={W_H_B.shape}\\n{W_H_B}\\n\")\n",
"print(f\"s: shape={s.shape}\\n{s}\\n\")\n",
"print(f\"ν: shape={ν.shape}\\n{ν}\\n\") # noqa: RUF001"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "XLx3sv9VtwZ9",
"outputId": "28f5f070-e37e-464e-d84e-2944cfdc28dc"
},
"outputs": [],
"source": [
"# Create a random link forces matrix.\n",
"link_forces = jax.random.uniform(\n",
" minval=-10.0,\n",
" maxval=10.0,\n",
" shape=(model.number_of_links(), 6),\n",
" key=jax.random.PRNGKey(0),\n",
")\n",
"\n",
"# Create a random joint force references vector.\n",
"# Note that these are called 'references' because the actual joint forces that\n",
"# are actuated might differ due to effects like joint friction.\n",
"joint_force_references = jax.random.uniform(\n",
" minval=-10.0, maxval=10.0, shape=(model.dofs(),), key=jax.random.PRNGKey(0)\n",
")\n",
"\n",
"# Create the references object.\n",
"references = js.references.JaxSimModelReferences.build(\n",
" model=model,\n",
" data=data,\n",
" link_forces=link_forces,\n",
" joint_force_references=joint_force_references,\n",
")\n",
"\n",
"print(f\"link_forces: shape={references.link_forces(model=model, data=data).shape}\")\n",
"print(f\"joint_force_references: shape={references.joint_force_references(model=model).shape}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AaG817vP4LfT"
},
"source": [
"## Robot Kinematics\n",
"\n",
"JaxSim offers functional APIs for computing kinematic quantities:\n",
"\n",
"- **`jaxsim.api.model`**: vectorized functions operating on the whole model.\n",
"- **`jaxsim.api.link`**: functions operating on individual links.\n",
"- **`jaxsim.api.frame`**: functions operating on individual frames. \n",
"\n",
"Due to JAX limitations on vectorizable data types, many APIs operate on indices instead of names. Since using indices can be error prone, JaxSim provides conversion functions for both links:\n",
"\n",
"- **jaxsim.api.link.names_to_idxs()**\n",
"- **jaxsim.api.link.idxs_to_names()**\n",
"\n",
"and frames: \n",
"\n",
"- **jaxsim.api.frame.names_to_idxs()**\n",
"- **jaxsim.api.frame.idxs_to_names()**\n",
"\n",
"We recommend using names whenever possible to avoid hard-to-trace errors.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QxImwfZz7pz-"
},
"outputs": [],
"source": [
"# Find the index of a link.\n",
"link_name = \"l_ankle_2\"\n",
"link_index = js.link.name_to_idx(model=model, link_name=link_name)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "C22Iqu2i4G-I",
"outputId": "94376151-177d-410f-f375-b7b8bd080992"
},
"outputs": [],
"source": [
"# @title Link Pose\n",
"\n",
"# Compute its pose w.r.t. the world frame through forward kinematics.\n",
"W_H_L = js.link.transform(model=model, data=data, link_index=link_index)\n",
"\n",
"print(f\"Transform of '{link_name}': shape={W_H_L.shape}\\n{W_H_L}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "DnSpE_f97RkX",
"outputId": "a3f6b535-4ae5-49f4-8921-7fe4dda5debb"
},
"outputs": [],
"source": [
"# @title Link 6D Velocity\n",
"\n",
"# JaxSim allows to select the so-called representation of the frame velocity.\n",
"L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body)\n",
"LW_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Mixed)\n",
"W_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Inertial)\n",
"\n",
"print(f\"Body-fixed velocity L_v_WL={L_v_WL}\")\n",
"print(f\"Mixed velocity: LW_v_WL={LW_v_WL}\")\n",
"print(f\"Inertial-fixed velocity: W_v_WL={W_v_WL}\")\n",
"\n",
"# These can also be computed passing through the link free-floating Jacobian.\n",
"# This type of Jacobian has a input velocity representation that corresponds\n",
"# the velocity representation of ν, and an output velocity representation that\n",
"# corresponds to the velocity representation of the desired 6D velocity.\n",
"\n",
"# You can use the following context manager to easily switch between representations.\n",
"with data.switch_velocity_representation(VelRepr.Body):\n",
"\n",
" # Body-fixed generalized velocity.\n",
" B_ν = data.generalized_velocity\n",
"\n",
" # Free-floating Jacobian accepting a body-fixed generalized velocity and\n",
" # returning an inertial-fixed link velocity.\n",
" W_J_WL_B = js.link.jacobian(\n",
" model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Inertial\n",
" )\n",
"\n",
"# Now the following relation should hold.\n",
"assert jnp.allclose(W_v_WL, W_J_WL_B @ B_ν)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SSoziCShtwZ9"
},
"outputs": [],
"source": [
"# Find the index of a frame.\n",
"frame_name = \"l_foot_front\"\n",
"frame_index = js.frame.name_to_idx(model=model, frame_name=frame_name)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "fVp_xP_1twZ9",
"outputId": "cfaa0569-d768-4708-c98c-a5867c056d04"
},
"outputs": [],
"source": [
"# @title Frame Pose\n",
"\n",
"# Compute its pose w.r.t. the world frame through forward kinematics.\n",
"W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index)\n",
"\n",
"print(f\"Transform of '{frame_name}': shape={W_H_F.shape}\\n{W_H_F}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QqaqxneEFYiW"
},
"outputs": [],
"source": [
"# @title Frame 6D Velocity\n",
"\n",
"# JaxSim allows to select the so-called representation of the frame velocity.\n",
"F_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Body)\n",
"FW_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Mixed)\n",
"W_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Inertial)\n",
"\n",
"print(f\"Body-fixed velocity F_v_WF={F_v_WF}\")\n",
"print(f\"Mixed velocity: FW_v_WF={FW_v_WF}\")\n",
"print(f\"Inertial-fixed velocity: W_v_WF={W_v_WF}\")\n",
"\n",
"# These can also be computed passing through the frame free-floating Jacobian.\n",
"# This type of Jacobian has a input velocity representation that corresponds\n",
"# the velocity representation of ν, and an output velocity representation that\n",
"# corresponds to the velocity representation of the desired 6D velocity.\n",
"\n",
"# You can use the following context manager to easily switch between representations.\n",
"with data.switch_velocity_representation(VelRepr.Body):\n",
"\n",
" # Body-fixed generalized velocity.\n",
" B_ν = data.generalized_velocity\n",
"\n",
" # Free-floating Jacobian accepting a body-fixed generalized velocity and\n",
" # returning an inertial-fixed link velocity.\n",
" W_J_WF_B = js.frame.jacobian(\n",
" model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Inertial\n",
" )\n",
"\n",
"# Now the following relation should hold.\n",
"assert jnp.allclose(W_v_WF, W_J_WF_B @ B_ν)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "d_vp6D74GoVZ",
"outputId": "798b9283-792e-4339-b56c-df2595fac974"
},
"source": [
"## Robot Dynamics\n",
"\n",
"JaxSim provides all the quantities involved in the equations of motion, restated here:\n",
"\n",
"$$\n",
"M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + \\sum_{L_i \\in \\mathcal{L}} J_{W,L_i}^\\top(\\mathbf{q}) \\: \\mathbf{f}_i\n",
".\n",
"$$\n",
"\n",
"Specifically, it can compute:\n",
"\n",
"- $M(\\mathbf{q}) \\in \\mathbb{R}^{(6+n)\\times(6+n)}$: the mass matrix.\n",
"- $\\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) \\in \\mathbb{R}^{6+n}$: the vector of bias forces.\n",
"- $B \\in \\mathbb{R}^{(6+n) \\times n}$ the joint selector matrix.\n",
"- $J_{W,L} \\in \\mathbb{R}^{6 \\times (6+n)}$ the Jacobian of link $L$.\n",
"\n",
"Often, for convenience, link Jacobians are stacked together. Since JaxSim efficiently computes the Jacobians for all links, using the stacked version is recommended when needed:\n",
"\n",
"$$\n",
"M(\\mathbf{q}) \\dot{\\boldsymbol{\\nu}} + \\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = B \\boldsymbol{\\tau} + J_{W,\\mathcal{L}}^\\top(\\mathbf{q}) \\: \\mathbf{f}_\\mathcal{L}\n",
".\n",
"$$\n",
"\n",
"Furthermore, there are applications that require unpacking the vector of bias forces as follow:\n",
"\n",
"$$\n",
"\\mathbf{h}(\\mathbf{q}, \\boldsymbol{\\nu}) = C(\\mathbf{q}, \\boldsymbol{\\nu}) \\boldsymbol{\\nu} + \\mathbf{g}(\\mathbf{q})\n",
",\n",
"$$\n",
"\n",
"where:\n",
"\n",
"- $\\mathbf{g}(\\mathbf{q}) \\in \\mathbb{R}^{6+n}$: the vector of gravity forces.\n",
"- $C(\\mathbf{q}, \\boldsymbol{\\nu}) \\in \\mathbb{R}^{(6+n)\\times(6+n)}$: the Coriolis matrix.\n",
"\n",
"Here below we report the functions to compute all these quantities. Note that all quantities depend on the active velocity representation of `data`. As it was done for the link velocity, it is possible to change the representation associated to all the computed quantities by operating within the corresponding context manager. Here below we consider the default representation of data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "oOKJOVfsH4Ki"
},
"outputs": [],
"source": [
"print(\"Velocity representation of data:\", data.velocity_representation, \"\\n\")\n",
"\n",
"# Compute the mass matrix.\n",
"M = js.model.free_floating_mass_matrix(model=model, data=data)\n",
"print(f\"M: shape={M.shape}\")\n",
"\n",
"# Compute the vector of bias forces.\n",
"h = js.model.free_floating_bias_forces(model=model, data=data)\n",
"print(f\"h: shape={h.shape}\")\n",
"\n",
"# Compute the vector of gravity forces.\n",
"g = js.model.free_floating_gravity_forces(model=model, data=data)\n",
"print(f\"g: shape={g.shape}\")\n",
"\n",
"# Compute the Coriolis matrix.\n",
"C = js.model.free_floating_coriolis_matrix(model=model, data=data)\n",
"print(f\"C: shape={C.shape}\")\n",
"\n",
"# Create a the joint selector matrix.\n",
"B = jnp.block([jnp.zeros(shape=(model.dofs(), 6)), jnp.eye(model.dofs())]).T\n",
"print(f\"B: shape={B.shape}\")\n",
"\n",
"# Compute the stacked tensor of link Jacobians.\n",
"J = js.model.generalized_free_floating_jacobian(model=model, data=data)\n",
"print(f\"J: shape={J.shape}\")\n",
"\n",
"# Extract the joint forces from the references object.\n",
"τ = references.joint_force_references(model=model)\n",
"print(f\"τ: shape={τ.shape}\")\n",
"\n",
"# Extract the link forces from the references object.\n",
"f_L = references.link_forces(model=model, data=data)\n",
"print(f\"f_L: shape={f_L.shape}\")\n",
"\n",
"# The following relation should hold.\n",
"assert jnp.allclose(h, C @ ν + g)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "FlNo8dNWKKtu",
"outputId": "313e939b-f88f-4407-c9ee-b5b3b7443061"
},
"source": [
"### Forward Dynamics\n",
"\n",
"$$\n",
"\\dot{\\boldsymbol{\\nu}} = \\text{FD}(\\mathbf{q}, \\boldsymbol{\\nu}, \\boldsymbol{\\tau}, \\mathbf{f}_{\\mathcal{L}})\n",
"$$\n",
"\n",
"JaxSim provides two alternative methods to compute the forward dynamics:\n",
"\n",
"1. Operate on the quantities of the equations of motion.\n",
"2. Call the recursive Articulated Body Algorithm (ABA).\n",
"\n",
"The physics engine provided by JaxSim exploits the efficient calculation of the forward dynamics with ABA for simulating the trajectories of the system dynamics."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LXARuRu1Ly1K"
},
"outputs": [],
"source": [
"ν̇_eom = jnp.linalg.pinv(M) @ (B @ τ - h + jnp.einsum(\"l6g,l6->g\", J, f_L))\n",
"\n",
"v̇_WB, s̈ = js.model.forward_dynamics_aba(\n",
" model=model, data=data, link_forces=f_L, joint_forces=joint_force_references\n",
")\n",
"\n",
"ν̇_aba = jnp.hstack([v̇_WB, s̈])\n",
"print(f\"ν̇: shape={ν̇_aba.shape}\") # noqa: RUF001\n",
"\n",
"# The following relation should hold.\n",
"assert jnp.allclose(ν̇_eom, ν̇_aba)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "g5GOYXDnLySU",
"outputId": "ad4ce77d-d06f-473a-9c32-040680d76aa5"
},
"source": [
"### Inverse Dynamics\n",
"\n",
"$$\n",
"(\\boldsymbol{\\tau}, \\, \\mathbf{f}_B) = \\text{ID}(\\mathbf{q}, \\boldsymbol{\\nu}, \\dot{\\boldsymbol{\\nu}}, \\mathbf{f}_{\\mathcal{L}})\n",
"$$\n",
"\n",
"JaxSim offers two methods to compute inverse dynamics:\n",
"\n",
"- Directly use the quantities from the equations of motion.\n",
"- Use the Recursive Newton-Euler Algorithm (RNEA).\n",
"\n",
"Unlike many other implementations, JaxSim's RNEA for floating-base systems is the true inverse of $\\text{FD}$. It also computes the 6D force applied to the base link that generates the base acceleration."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "UTae5MjhaP2H"
},
"outputs": [],
"source": [
"f_B, τ_rnea = js.model.inverse_dynamics(\n",
" model=model,\n",
" data=data,\n",
" base_acceleration=v̇_WB,\n",
" joint_accelerations=s̈,\n",
" # To check that f_B works, let's remove the force applied\n",
" # to the base link from the link forces.\n",
" link_forces=f_L.at[0].set(jnp.zeros(6))\n",
")\n",
"\n",
"print(f\"f_B: shape={f_B.shape}\")\n",
"print(f\"τ_rnea: shape={τ_rnea.shape}\")\n",
"\n",
"# The following relations should hold.\n",
"assert jnp.allclose(τ_rnea, τ)\n",
"assert jnp.allclose(f_B, link_forces[0])"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "gYZ1jK1Neg1H",
"outputId": "0de79770-1e18-4027-bb47-5713bc1b4a72"
},
"source": [
"### Centroidal Dynamics\n",
"\n",
"Centroidal dynamics is a useful simplification often employed in planning and control applications. It represents the dynamics projected onto a mixed frame associated with the center of mass (CoM):\n",
"\n",
"$$\n",
"G = G[W] = ({}^W \\mathbf{p}_{\\text{CoM}}, [W])\n",
".\n",
"$$\n",
"\n",
"The governing equations for centroidal dynamics take into account the 6D centroidal momentum:\n",
"\n",
"$$\n",
"{}_G \\mathbf{h} =\n",
"\\begin{bmatrix}\n",
"{}_G \\mathbf{h}^l \\\\ {}_G \\mathbf{h}^\\omega\n",
"\\end{bmatrix} =\n",
"\\begin{bmatrix}\n",
"m \\, {}^W \\dot{\\mathbf{p}}_\\text{CoM} \\\\ {}_G \\mathbf{h}^\\omega\n",
"\\end{bmatrix}\n",
"\\in \\mathbb{R}^6\n",
".\n",
"$$\n",
"\n",
"The equations of centroidal dynamics can be expressed as:\n",
"\n",
"$$\n",
"{}_G \\dot{\\mathbf{h}} =\n",
"m \\,\n",
"\\begin{bmatrix}\n",
"{}^W \\mathbf{g} \\\\ \\mathbf{0}_3\n",
"\\end{bmatrix} +\n",
"\\sum_{C_i \\in \\mathcal{C}} {}_G \\mathbf{X}^{C_i} \\, {}_{C_i} \\mathbf{f}_i\n",
".\n",
"$$\n",
"\n",
"While centroidal dynamics can function independently by considering the total mass $m \\in \\mathbb{R}$ of the robot and the transformations for 6D contact forces ${}_G \\mathbf{X}^{C_i}$ corresponding to the pose ${}^G \\mathbf{H}_{C_i} \\in \\text{SE}(3)$ of the contact frames, advanced kino-dynamic methods may require a relationship between full kinematics and centroidal dynamics. This is typically achieved through the _Centroidal Momentum Matrix_ (also known as the _centroidal momentum Jacobian_):\n",
"\n",
"$$\n",
"{}_G \\mathbf{h} = J_\\text{CMM}(\\mathbf{q}) \\, \\boldsymbol{\\nu}\n",
".\n",
"$$\n",
"\n",
"JaxSim offers APIs to compute all these quantities (and many more) in the `jaxsim.api.com` package."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rrSfxp8lh9YZ"
},
"outputs": [],
"source": [
"# Number of contact points.\n",
"n_cp = len(model.kin_dyn_parameters.contact_parameters.body)\n",
"print(\"Number of contact points:\", n_cp, \"\\n\")\n",
"\n",
"# Compute the centroidal momentum.\n",
"J_CMM = js.com.centroidal_momentum_jacobian(model=model, data=data)\n",
"G_h = J_CMM @ ν\n",
"print(f\"G_h: shape={G_h.shape}\")\n",
"print(f\"J_CMM: shape={J_CMM.shape}\")\n",
"\n",
"# The following relation should hold.\n",
"assert jnp.allclose(G_h, js.com.centroidal_momentum(model=model, data=data))\n",
"\n",
"# If we consider all contact points of the model as active\n",
"# (discourages since they might be too many), the 6D transforms of\n",
"# collidable points can be computed as follows:\n",
"W_H_C = js.contact.transforms(model=model, data=data)\n",
"\n",
"# Compute the pose of the G frame.\n",
"W_p_CoM = js.com.com_position(model=model, data=data)\n",
"G_H_W = jaxsim.math.Transform.inverse(jnp.eye(4).at[0:3, 3].set(W_p_CoM))\n",
"\n",
"# Convert from SE(3) to the transforms for 6D forces.\n",
"G_Xf_C = jax.vmap(\n",
" lambda W_H_Ci: jaxsim.math.Adjoint.from_transform(\n",
" transform=G_H_W @ W_H_Ci, inverse=True\n",
" )\n",
")(W_H_C)\n",
"print(f\"G_Xf_C: shape={G_Xf_C.shape}\")\n",
"\n",
"# Let's create random 3D linear forces applied to the contact points.\n",
"C_fl = jax.random.uniform(\n",
" minval=-10.0,\n",
" maxval=10.0,\n",
" shape=(n_cp, 3),\n",
" key=jax.random.PRNGKey(0),\n",
")\n",
"\n",
"# Compute the 3D gravity vector and the total mass of the robot.\n",
"m = js.model.total_mass(model=model)\n",
"\n",
"# The centroidal dynamics can be computed as follows.\n",
"G_ḣ = 0\n",
"G_ḣ += m * jnp.hstack([0, 0, model.gravity, 0, 0, 0])\n",
"G_ḣ += jnp.einsum(\"c66,c6->6\", G_Xf_C, jnp.hstack([C_fl, jnp.zeros_like(C_fl)]))\n",
"print(f\"G_ḣ: shape={G_ḣ.shape}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Ot6HePB_twaE",
"outputId": "02a6abae-257e-45ee-e9de-6a607cdbeb9a"
},
"source": [
"## Contact Frames\n",
"\n",
"Many control and planning applications require projecting the floating-base dynamics into the contact space or computing quantities related to active contact points, such as enforcing holonomic constraints.\n",
"\n",
"The underlying theory for these applications becomes clearer in a mixed representation. Specifically, the position, linear velocity, and linear acceleration of contact points in their corresponding mixed frame align with the numerical derivatives of their coordinate vectors.\n",
"\n",
"Key methodologies in this area may involve the Delassus matrix:\n",
"\n",
"$$\n",
"\\Psi(\\mathbf{q}) = J_{W,C}(\\mathbf{q}) \\, M(\\mathbf{q})^{-1} \\, J_{W,C}^T(\\mathbf{q})\n",
"$$\n",
"\n",
"or the linear acceleration of a contact point:\n",
"\n",
"$$\n",
"{}^W \\ddot{\\mathbf{p}}_C = \\frac{\\text{d} (J^l_{W,C} \\boldsymbol{\\nu})}{\\text{d}t}\n",
"= \\dot{J}^l_{W,C} \\boldsymbol{\\nu} + J^l_{W,C} \\dot{\\boldsymbol{\\nu}}\n",
".\n",
"$$\n",
"\n",
"JaxSim offers APIs to compute all these quantities (and many more) in the `jaxsim.api.contact` package."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LITRC3STliKR"
},
"outputs": [],
"source": [
"with (\n",
" data.switch_velocity_representation(VelRepr.Mixed),\n",
" references.switch_velocity_representation(VelRepr.Mixed),\n",
"):\n",
"\n",
" # Compute the mixed generalized velocity.\n",
" BW_ν = data.generalized_velocity\n",
"\n",
" # Compute the mixed generalized acceleration.\n",
" BW_ν̇ = jnp.hstack(\n",
" js.model.forward_dynamics(\n",
" model=model,\n",
" data=data,\n",
" link_forces=references.link_forces(model=model, data=data),\n",
" joint_forces=references.joint_force_references(model=model),\n",
" )\n",
" )\n",
"\n",
" # Compute the mass matrix in mixed representation.\n",
" BW_M = js.model.free_floating_mass_matrix(model=model, data=data)\n",
"\n",
" # Compute the contact Jacobian and its derivative.\n",
" Jl_WC = js.contact.jacobian(model=model, data=data)[:, 0:3, :]\n",
" J̇l_WC = js.contact.jacobian_derivative(model=model, data=data)[:, 0:3, :]\n",
"\n",
"# Compute the Delassus matrix.\n",
"Ψ = jnp.vstack(Jl_WC) @ jnp.linalg.lstsq(BW_M, jnp.vstack(Jl_WC).T)[0]\n",
"print(f\"Ψ: shape={Ψ.shape}\")\n",
"\n",
"# Compute the transforms of the mixed frames implicitly associated\n",
"# to each collidable point.\n",
"W_H_C = js.contact.transforms(model=model, data=data)\n",
"print(f\"W_H_C: shape={W_H_C.shape}\")\n",
"\n",
"# Compute the linear velocity of the collidable points.\n",
"with data.switch_velocity_representation(VelRepr.Mixed):\n",
" W_ṗ_B = js.contact.collidable_point_velocities(model=model, data=data)[:, 0:3]\n",
" print(f\"W_ṗ_B: shape={W_ṗ_B.shape}\")\n",
"\n",
"# Compute the linear acceleration of the collidable points.\n",
"W_p̈_C = 0\n",
"W_p̈_C += jnp.einsum(\"c3g,g->c3\", J̇l_WC, BW_ν)\n",
"W_p̈_C += jnp.einsum(\"c3g,g->c3\", Jl_WC, BW_ν̇)\n",
"print(f\"W_p̈_C: shape={W_p̈_C.shape}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LITRC3STliKR"
},
"source": [
"## Conclusions\n",
"\n",
"This notebook provided an overview of the main APIs in JaxSim for its use as a multibody dynamics library. Here are a few key points to remember:\n",
"\n",
"- Explore all the modules in the `jaxsim.api` package to discover the full range of APIs available. Many more functionalities exist beyond what was covered in this notebook.\n",
"- All APIs follow a functional approach, consistent with the JAX programming style.\n",
"- This functional design allows for easy application of `jax.vmap` to execute functions in parallel on hardware accelerators.\n",
"- Since the entire multibody dynamics library is built with JAX, it natively supports `jax.grad`, `jax.jacfwd`, and `jax.jacrev` transformations, enabling automatic differentiation through complex logic without additional effort.\n",
"\n",
"Have fun!"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuClass": "premium",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "comodo_jaxsim",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.8"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
================================================
FILE: examples/jaxsim_as_physics_engine.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "H-WgcgGQaTG7"
},
"source": [
"# JaxSim as a hardware-accelerated parallel physics engine\n",
"\n",
"This notebook shows how to use the key APIs to load a robot model and simulate multiple trajectories simultaneously.\n",
"\n",
"\n",
" \n",
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SgOSnrSscEkt"
},
"source": [
"## Prepare the environment"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fdqvAqMDaTG9"
},
"outputs": [],
"source": [
"# @title Imports and setup\n",
"import os\n",
"import sys\n",
"from IPython.display import clear_output\n",
"\n",
"IS_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"# Install JAX and Gazebo\n",
"if IS_COLAB:\n",
" !{sys.executable} -m pip install --pre -qU jaxsim\n",
" !apt install -qq lsb-release wget gnupg\n",
" !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
" !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
" !apt -qq update\n",
" !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
"\n",
" clear_output()\n",
"\n",
"# Set environment variable to avoid GPU out of memory errors\n",
"%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",
"\n",
"\n",
"# ================\n",
"# Notebook imports\n",
"# ================\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import jaxsim.api as js\n",
"from jaxsim import logging\n",
"import pathlib\n",
"\n",
"logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
"print(f\"Running on {jax.devices()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NqjuZKvOaTG_"
},
"source": [
"## Prepare the simulation\n",
"\n",
"JaxSim supports loading robot descriptions from both [SDF][sdformat] and [URDF][urdf] files. In this example, we will load the [ergoCub][ergocub] model urdf.\n",
"\n",
"[sdformat]: http://sdformat.org/\n",
"[urdf]: http://wiki.ros.org/urdf/\n",
"[ergocub]: https://ergocub.eu/\n",
"[rod]: https://github.com/gbionics/rod\n",
"\n",
"### Create the model and its data\n",
" To define a simulation we need two main objects:\n",
"\n",
"- `model`: an object that defines the dynamics of the system.\n",
"- `data`: an object that contains the state of the system.\n",
"\n",
"\n",
"The `JaxSimModel` object contains the simulation time step, the integrator and the contact model.\n",
"To see the advanced usage, check the advanced example, where you will see how to pass explicitly an integrator class and state to the `model` object and how to change the contact model."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Create the model "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "etQ577cFaTHA"
},
"outputs": [],
"source": [
"# Create the JaxSim model.\n",
"try:\n",
" os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.7\"\n",
"\n",
" import robot_descriptions.ergocub_description\n",
"\n",
"finally:\n",
" _ = os.environ.pop(\"ROBOT_DESCRIPTION_COMMIT\", None)\n",
"\n",
"model_description_path = pathlib.Path(\n",
" robot_descriptions.ergocub_description.URDF_PATH.replace(\n",
" \"ergoCubSN002\", \"ergoCubSN001\"\n",
" )\n",
")\n",
"\n",
"clear_output()\n",
"\n",
"full_model = js.model.JaxSimModel.build_from_model_description(\n",
" model_description=model_description_path,\n",
" time_step=0.0001,\n",
" is_urdf=True\n",
")\n",
"\n",
"joints_list = tuple(('l_shoulder_pitch', 'l_shoulder_roll', 'l_shoulder_yaw', 'l_elbow',\n",
" 'r_shoulder_pitch', 'r_shoulder_roll', 'r_shoulder_yaw', 'r_elbow',\n",
" 'l_hip_pitch', 'l_hip_roll', 'l_hip_yaw', 'l_knee', 'l_ankle_pitch', 'l_ankle_roll',\n",
" 'r_hip_pitch', 'r_hip_roll', 'r_hip_yaw', 'r_knee', 'r_ankle_pitch', 'r_ankle_roll'))\n",
"\n",
"model = js.model.reduce(\n",
" model=full_model,\n",
" considered_joints=joints_list\n",
")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Create the data object \n",
"\n",
"The data object is never changed by reference. Anytime you call a method aimed at modifying data, like `reset_base_position`, a new data object will be returned with the updated attributes while the original data will not be changed."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create the data of a single model.\n",
"data = js.data.JaxSimModelData.build(model=model, base_position=jnp.array([0.0, 0.0, 1.0]))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Simulation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create a random JAX key.\n",
"\n",
"key = jax.random.PRNGKey(seed=0)\n",
"\n",
"# Initialize the simulated time.\n",
"T = jnp.arange(start=0, stop=0.3, step=model.time_step)\n",
"\n",
"# Simulate\n",
"for _t in T:\n",
" data = js.model.step(\n",
" model=model,\n",
" data=data,\n",
" link_forces=None,\n",
" joint_force_references=None,\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Vectorized simulation \n",
"\n",
"We will now vectorize the simulation on batched data using `jax.vmap`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# first we have to vmap the function\n",
"\n",
"import functools\n",
"from typing import Any\n",
"\n",
"\n",
"@jax.jit\n",
"def step_single(\n",
" model: js.model.JaxSimModel,\n",
" data: js.data.JaxSimModelData,\n",
") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
"\n",
" # Close step over static arguments.\n",
" return js.model.step(\n",
" model=model,\n",
" data=data,\n",
" link_forces=None,\n",
" joint_force_references=None,\n",
" )\n",
"\n",
"\n",
"@jax.jit\n",
"@functools.partial(jax.vmap, in_axes=(None, 0))\n",
"def step_parallel(\n",
" model: js.model.JaxSimModel,\n",
" data: js.data.JaxSimModelData,\n",
") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
"\n",
" return step_single(\n",
" model=model, data=data\n",
" )\n",
"\n",
"\n",
"# Then we have to create the vector of initial state\n",
"batch_size = 5\n",
"data_batch_t0 = jax.vmap(\n",
" lambda pos: js.data.JaxSimModelData.build(model=model, base_position=pos))(jnp.tile(jnp.array([0.0, 0.0, 1.0]), (batch_size, 1)))\n",
"\n",
"data = data_batch_t0\n",
"for _t in T:\n",
" data = step_parallel(model, data)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuClass": "premium",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "jaxsim",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.1"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
================================================
FILE: examples/jaxsim_as_physics_engine_advanced.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "H-WgcgGQaTG7"
},
"source": [
"# JaxSim as a hardware-accelerated parallel physics engine-advanced usage\n",
"\n",
"JaxSim is developed to optimize synthetic data generation by sampling trajectories using hardware accelerators such as GPUs and TPUs.\n",
"\n",
"In this notebook, you'll learn how to use the key APIs to load a simple robot model (a sphere) and simulate multiple trajectories in parallel on GPUs.\n",
"\n",
"\n",
" \n",
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SgOSnrSscEkt"
},
"source": [
"## Prepare the environment"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fdqvAqMDaTG9"
},
"outputs": [],
"source": [
"# @title Imports and setup\n",
"import sys\n",
"from IPython.display import clear_output\n",
"\n",
"IS_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"# Install JAX and Gazebo\n",
"if IS_COLAB:\n",
" !{sys.executable} -m pip install --pre -qU jaxsim[viz]\n",
" !apt install -qq lsb-release wget gnupg\n",
" !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
" !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
" !apt -qq update\n",
" !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
"\n",
" clear_output()\n",
"\n",
"# Set environment variable to avoid GPU out of memory errors\n",
"%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",
"\n",
"# ================\n",
"# Notebook imports\n",
"# ================\n",
"\n",
"import os\n",
"\n",
"if sys.platform == 'darwin':\n",
" os.environ[\"MUJOCO_GL\"] = \"glfw\"\n",
"else:\n",
" os.environ[\"MUJOCO_GL\"] = \"egl\"\n",
"\n",
"import jax\n",
"\n",
"import jax.numpy as jnp\n",
"import jaxsim.api as js\n",
"import rod\n",
"from jaxsim import logging\n",
"from rod.builder.primitives import SphereBuilder\n",
"\n",
"logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
"print(f\"Running on {jax.devices()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QtCCUhdpdGFH"
},
"source": [
"## Prepare the simulation\n",
"\n",
"JaxSim supports loading robot descriptions from both [SDF][sdformat] and [URDF][urdf] files. This is done using the [`gbionics/rod`][rod] library, which processes these formats.\n",
"\n",
"The `rod` library also allows creating in-memory models that can be serialized to SDF or URDF. We'll use this functionality to build a sphere model, which will later be used to create the JaxSim model.\n",
"\n",
"[sdformat]: http://sdformat.org/\n",
"[urdf]: http://wiki.ros.org/urdf/\n",
"[rod]: https://github.com/gbionics/rod"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "0emoMQhCaTG_"
},
"outputs": [],
"source": [
"# @title Create the model description of a sphere\n",
"\n",
"# Create a SDF model.\n",
"# The builder takes care to compute the right inertia tensor for you.\n",
"rod_sdf = rod.Sdf(\n",
" version=\"1.7\",\n",
" model=SphereBuilder(radius=0.10, mass=1.0, name=\"sphere\")\n",
" .build_model()\n",
" .add_link()\n",
" .add_inertial()\n",
" .add_visual()\n",
" .add_collision()\n",
" .build(),\n",
")\n",
"\n",
"# Rod allows to update the frames w.r.t. the poses are expressed.\n",
"rod_sdf.model.switch_frame_convention(\n",
" frame_convention=rod.FrameConvention.Urdf, explicit_frames=True\n",
")\n",
"\n",
"# Serialize the model to a SDF string.\n",
"model_sdf_string = rod_sdf.serialize(pretty=True)\n",
"print(model_sdf_string)\n",
"\n",
"# JaxSim currently only supports collisions between points attached to bodies\n",
"# and a ground surface modeled as a heightmap sampled from a smooth function.\n",
"# While this approach is universal as it applies to generic meshes, the number\n",
"# of considered points greatly affects the performance. Spheres, by default,\n",
"# are discretized with 250 points. It's too much for this simple example.\n",
"# This number can be decreased with the following environment variable.\n",
"os.environ[\"JAXSIM_COLLISION_SPHERE_POINTS\"] = \"50\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NqjuZKvOaTG_"
},
"source": [
"### Create the model and its data\n",
"\n",
"JAXsim offers a simple functional API in order to interact in a memory-efficient way with the simulation. Four main elements are used to define a simulation:\n",
"\n",
"- `model`: an object that defines the dynamics of the system.\n",
"- `data`: an object that contains the state of the system.\n",
"- `integrator` *(Optional)*: an object that defines the integration method.\n",
"- `integrator_metadata` *(Optional)*: an object that contains the state of the integrator.\n",
"\n",
"The `JaxSimModel` object contains the simulation time step, the integrator and the contact model.\n",
"In this example, we will explicitly pass an integrator class to the `model` object and we will use the default `SoftContacts` contact model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "etQ577cFaTHA"
},
"outputs": [],
"source": [
"# Create the JaxSim model.\n",
"# This is shared among all the parallel instances.\n",
"model = js.model.JaxSimModel.build_from_model_description(\n",
" model_description=model_sdf_string,\n",
" time_step=0.001,\n",
")\n",
"\n",
"# Create the data of a single model.\n",
"# We will create a vectorized instance later.\n",
"data_single = js.data.JaxSimModelData.zero(model=model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "o86Teq5piVGj"
},
"outputs": [],
"source": [
"# Initialize the simulated time.\n",
"T = jnp.arange(start=0, stop=1.0, step=model.time_step)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "V6IeD2B3m4F0"
},
"source": [
"## Sample a batch of trajectories in parallel\n",
"\n",
"With the provided resources, you can step through an open-loop trajectory on a single model using `jaxsim.api.model.step`.\n",
"\n",
"In this notebook, we'll focus on running parallel steps. We'll use JAX's automatic vectorization to apply the step function to batched data.\n",
"\n",
"Note that these parallel simulations are independent — models don't interact, so there's no need to avoid initial collisions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vtEn0aIzr_2j"
},
"outputs": [],
"source": [
"# @title Generate batched initial data\n",
"\n",
"# Create a random JAX key.\n",
"key = jax.random.PRNGKey(seed=0)\n",
"\n",
"# Split subkeys for sampling random initial data.\n",
"batch_size = 9\n",
"row_length = int(jnp.sqrt(batch_size))\n",
"row_dist = 0.3 * row_length\n",
"key, *subkeys = jax.random.split(key=key, num=batch_size + 1)\n",
"\n",
"# Create the batched data by sampling the height from [0.5, 0.6] meters.\n",
"data_batch_t0 = jax.vmap(\n",
" lambda key: js.data.random_model_data(\n",
" model=model,\n",
" key=key,\n",
" base_pos_bounds=([0, 0, 0.3], [0, 0, 1.2]),\n",
" base_vel_lin_bounds=(0, 0),\n",
" base_vel_ang_bounds=(0, 0),\n",
" )\n",
")(jnp.vstack(subkeys))\n",
"\n",
"x, y = jnp.meshgrid(\n",
" jnp.linspace(-row_dist, row_dist, num=row_length),\n",
" jnp.linspace(-row_dist, row_dist, num=row_length),\n",
")\n",
"xy_coordinate = jnp.stack([x.flatten(), y.flatten()], axis=-1)\n",
"\n",
"# Reset the x and y position to a grid.\n",
"data_batch_t0 = data_batch_t0.replace(\n",
" model=model,\n",
" base_position=data_batch_t0.base_position.at[:, :2].set(xy_coordinate),\n",
")\n",
"\n",
"print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position[0:10])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0tQPfsl6uxHm"
},
"outputs": [],
"source": [
"# @title Create parallel step function\n",
"\n",
"import functools\n",
"from typing import Any\n",
"\n",
"\n",
"@jax.jit\n",
"def step_single(\n",
" model: js.model.JaxSimModel,\n",
" data: js.data.JaxSimModelData,\n",
") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
"\n",
" # Close step over static arguments.\n",
" return js.model.step(\n",
" model=model,\n",
" data=data,\n",
" link_forces=None,\n",
" joint_force_references=None,\n",
" )\n",
"\n",
"\n",
"@jax.jit\n",
"@functools.partial(jax.vmap, in_axes=(None, 0))\n",
"def step_parallel(\n",
" model: js.model.JaxSimModel,\n",
" data: js.data.JaxSimModelData,\n",
") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
"\n",
" return step_single(\n",
" model=model, data=data\n",
" )\n",
"\n",
"\n",
"# The first run will be slow since JAX needs to JIT-compile the functions.\n",
"_ = step_single(model, data_single)\n",
"_ = step_parallel(model, data_batch_t0)\n",
"\n",
"# Benchmark the execution of a single step.\n",
"print(\"\\nSingle simulation step:\")\n",
"%timeit step_single(model, data_single)\n",
"\n",
"# On hardware accelerators, there's a range of batch_size values where\n",
"# increasing the number of parallel instances doesn't affect computation time.\n",
"# This range depends on the GPU/TPU specifications.\n",
"print(f\"\\nParallel simulation steps (batch_size={batch_size} on {jax.devices()[0]}):\")\n",
"%timeit step_parallel(model, data_batch_t0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VNwzT2JQ1n15"
},
"outputs": [],
"source": [
"# @title Run parallel simulation\n",
"\n",
"data = data_batch_t0\n",
"data_trajectory_list = []\n",
"\n",
"for _ in T:\n",
"\n",
" data = step_parallel(model, data)\n",
" data_trajectory_list.append(data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Y6n720Cr3G44"
},
"source": [
"## Visualize trajectory"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BLPODyKr3Lyg"
},
"outputs": [],
"source": [
"# Convert a list of PyTrees to a batched PyTree.\n",
"# This operation is called 'tree transpose' in JAX.\n",
"data_trajectory = jax.tree.map(lambda *leafs: jnp.stack(leafs), *data_trajectory_list)\n",
"\n",
"print(f\"W_p_B: shape={data_trajectory.base_position.shape}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-jxJXy5r3RMt"
},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"\n",
"\n",
"plt.plot(T, data_trajectory.base_position[:, :, 2])\n",
"plt.grid(True)\n",
"plt.xlabel(\"Time [s]\")\n",
"plt.ylabel(\"Height [m]\")\n",
"plt.title(\"Height trajectory of the sphere\")\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import jaxsim.mujoco\n",
"\n",
"mjcf_string, assets = jaxsim.mujoco.ModelToMjcf.convert(\n",
" model.built_from,\n",
" cameras=jaxsim.mujoco.loaders.MujocoCamera.build_from_target_view(\n",
" camera_name=\"sphere_cam\",\n",
" lookat=[0, 0, 0.3],\n",
" distance=4,\n",
" azimuth=150,\n",
" elevation=-10,\n",
" ),\n",
")\n",
"\n",
"# Create a helper for each parallel instance.\n",
"mj_model_helpers = [\n",
" jaxsim.mujoco.MujocoModelHelper.build_from_xml(\n",
" mjcf_description=mjcf_string, assets=assets\n",
" )\n",
" for _ in range(batch_size)\n",
"]\n",
"\n",
"# Create the video recorder.\n",
"recorder = jaxsim.mujoco.MujocoVideoRecorder(\n",
" model=mj_model_helpers[0].model,\n",
" data=[helper.data for helper in mj_model_helpers],\n",
" fps=int(1 / model.time_step),\n",
" width=320 * 2,\n",
" height=240 * 2,\n",
")\n",
"\n",
"for data_t in data_trajectory_list:\n",
"\n",
" for helper, base_position, base_quaternion, joint_position in zip(\n",
" mj_model_helpers,\n",
" data_t.base_position,\n",
" data_t.base_orientation,\n",
" data_t.joint_positions,\n",
" strict=True,\n",
" ):\n",
" helper.set_base_position(position=base_position)\n",
" helper.set_base_orientation(orientation=base_quaternion)\n",
"\n",
" if model.dofs() > 0:\n",
" helper.set_joint_positions(\n",
" positions=joint_position, joint_names=model.joint_names()\n",
" )\n",
"\n",
" # Record a new video frame.\n",
" recorder.record_frame(camera_name=\"sphere_cam\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import mediapy as media\n",
"\n",
"media.show_video(recorder.frames, fps=recorder.fps)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuClass": "premium",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "jaxpypi",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.1"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
================================================
FILE: examples/jaxsim_for_robot_controllers.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "EhPy6FgiZH4d"
},
"source": [
"# JaxSim for developing closed-loop robot controllers\n",
"\n",
"Originally developed as a **hardware-accelerated physics engine**, JaxSim has expanded its capabilities to become a full-featured **JAX-based multibody dynamics library**.\n",
"\n",
"In this notebook, you'll explore how to combine these two core features. Specifically, you'll learn how to load a robot model and design a model-based controller for closed-loop simulations.\n",
"\n",
"\n",
" \n",
""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vsf1AlxdZH4f"
},
"outputs": [],
"source": [
"# @title Prepare the environment\n",
"from IPython.display import clear_output\n",
"import sys\n",
"\n",
"IS_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"# Install JAX, sdformat, and other notebook dependencies.\n",
"if IS_COLAB:\n",
" !{sys.executable} -m pip install --pre -qU jaxsim[viz]\n",
" !apt install -qq lsb-release wget gnupg\n",
" !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
" !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
" !apt -qq update\n",
" !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
"\n",
" clear_output()\n",
"\n",
"# ================\n",
"# Notebook imports\n",
"# ================\n",
"\n",
"import os\n",
"\n",
"if sys.platform == 'darwin':\n",
" os.environ[\"MUJOCO_GL\"] = \"glfw\"\n",
"else:\n",
" os.environ[\"MUJOCO_GL\"] = \"egl\"\n",
"\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import jaxsim.mujoco\n",
"from jaxsim import logging\n",
"\n",
"logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
"print(f\"Running on {jax.devices()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kN-b9nOsZH4g"
},
"source": [
"We will use a simple cartpole model for this example. The cartpole model is a 2D model with a cart that can move horizontally and a pole that can rotate around the cart. The state of the cartpole is given by the position of the cart, the angle of the pole, the velocity of the cart, and the angular velocity of the pole. The control input is the horizontal force applied to the cart."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5aLqrZDqR5LA"
},
"source": [
"## Prepare the simulation\n",
"\n",
"JaxSim supports loading robot models from both [SDF][sdformat] and [URDF][urdf] files, utilizing the [`gbionics/rod`][rod] library for processing these formats.\n",
"\n",
"The `rod` library library can read URDF files and validates them internally using [`gazebosim/sdformat`][sdformat_github]. In this example, we'll load a cart-pole model, which will be used to create the JaxSim simulation model.\n",
"\n",
"[sdformat]: http://sdformat.org/\n",
"[urdf]: http://wiki.ros.org/urdf/\n",
"[rod]: https://github.com/gbionics/rod\n",
"[sdformat_github]: https://github.com/gazebosim/sdformat"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.path.abspath(\"\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PZM7hEvFZH4h"
},
"outputs": [],
"source": [
"# @title Load the URDF model\n",
"import pathlib\n",
"import urllib\n",
"\n",
"# Retrieve the file\n",
"url = \"https://raw.githubusercontent.com/gbionics/jaxsim/refs/heads/main/examples/assets/cartpole.urdf\"\n",
"model_path, _ = urllib.request.urlretrieve(url)\n",
"model_urdf_string = pathlib.Path(model_path).read_text()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "M5XsKehvZH4j"
},
"outputs": [],
"source": [
"# @title Create the model and its data\n",
"\n",
"import jaxsim.api as js\n",
"\n",
"# Create the model from the model description.\n",
"model = js.model.JaxSimModel.build_from_model_description(\n",
" model_description=model_urdf_string,\n",
" time_step=0.010,\n",
")\n",
"\n",
"# Create the data storing the simulation state.\n",
"data_zero = js.data.JaxSimModelData.zero(model=model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jk9csR5ETgn1"
},
"outputs": [],
"source": [
"# @title Define simulation parameters\n",
"\n",
"# Initialize the simulated time.\n",
"T = jnp.arange(start=0, stop=5.0, step=model.time_step)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bo6Ke5nAWL-S"
},
"source": [
"## Prepare the MuJoCo renderer\n",
"\n",
"For visualization purpose, we use the passive viewer of the MuJoCo simulator. It allows to either open an interactive windows when used locally or record a video when used in notebooks."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "j1_I2i5TZH4n"
},
"outputs": [],
"source": [
"# Create the MJCF resources from the URDF.\n",
"mjcf_string, assets = jaxsim.mujoco.UrdfToMjcf.convert(\n",
" urdf=model.built_from,\n",
" # Create the camera used by the recorder.\n",
" cameras=jaxsim.mujoco.loaders.MujocoCamera.build_from_target_view(\n",
" camera_name=\"cartpole_camera\",\n",
" lookat=js.link.com_position(\n",
" model=model,\n",
" data=data_zero,\n",
" link_index=js.link.name_to_idx(model=model, link_name=\"cart\"),\n",
" in_link_frame=False,\n",
" ),\n",
" distance=3,\n",
" azimuth=150,\n",
" elevation=-10,\n",
" ),\n",
")\n",
"\n",
"# Create a helper to operate on the MuJoCo model and data.\n",
"mj_model_helper = jaxsim.mujoco.MujocoModelHelper.build_from_xml(\n",
" mjcf_description=mjcf_string, assets=assets\n",
")\n",
"\n",
"# Create the video recorder.\n",
"recorder = jaxsim.mujoco.MujocoVideoRecorder(\n",
" model=mj_model_helper.model,\n",
" data=mj_model_helper.data,\n",
" fps=int(1 / model.time_step),\n",
" width=320 * 2,\n",
" height=240 * 2,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DpRvvGujZH4o"
},
"source": [
"## Open-loop simulation\n",
"\n",
"Now, let's run a simulation to demonstrate the open-loop dynamics of the system."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gSWzcsKWZH4p"
},
"outputs": [],
"source": [
"import mediapy as media\n",
"\n",
"\n",
"# Create a random joint position.\n",
"# For a random full state, you can use jaxsim.api.data.random_model_data.\n",
"random_joint_positions = jax.random.uniform(\n",
" minval=-1.0,\n",
" maxval=1.0,\n",
" shape=(model.dofs(),),\n",
" key=jax.random.PRNGKey(0),\n",
")\n",
"\n",
"# Reset the state to the random joint positions.\n",
"data = js.data.JaxSimModelData.build(model=model, joint_positions=random_joint_positions)\n",
"\n",
"for _ in T:\n",
"\n",
" # Step the JaxSim simulation.\n",
" data = js.model.step(\n",
" model=model,\n",
" data=data,\n",
" joint_force_references=None,\n",
" link_forces=None,\n",
" )\n",
"\n",
" # Update the MuJoCo data.\n",
" mj_model_helper.set_joint_positions(\n",
" positions=data.joint_positions, joint_names=model.joint_names()\n",
" )\n",
"\n",
" # Record a new video frame.\n",
" recorder.record_frame(camera_name=\"cartpole_camera\")\n",
"\n",
"\n",
"# Play the video.\n",
"media.show_video(recorder.frames, fps=recorder.fps)\n",
"recorder.frames = []"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "j1rguK3UZH4p"
},
"source": [
"## Closed-loop simulation\n",
"\n",
"Next, let's design a simple computed torque controller. The equations of motion for the cart-pole system are given by:\n",
"\n",
"$$\n",
"M_{ss}(\\mathbf{s}) \\, \\ddot{\\mathbf{s}} + \\mathbf{h}_s(\\mathbf{s}, \\dot{\\mathbf{s}}) = \\boldsymbol{\\tau}\n",
",\n",
"$$\n",
"\n",
"where:\n",
"\n",
"- $\\mathbf{s} \\in \\mathbb{R}^n$ are the joint positions.\n",
"- $\\dot{\\mathbf{s}} \\in \\mathbb{R}^n$ are the joint velocities.\n",
"- $\\ddot{\\mathbf{s}} \\in \\mathbb{R}^n$ are the joint accelerations.\n",
"- $\\boldsymbol{\\tau} \\in \\mathbb{R}^n$ are the joint torques.\n",
"- $M_{ss} \\in \\mathbb{R}^{n \\times n}$ is the mass matrix.\n",
"- $\\mathbf{h}_s \\in \\mathbb{R}^n$ is the vector of bias forces.\n",
"\n",
"JaxSim computes these quantities for floating-base systems, so we specifically focus on the joint-related portions by marking them with subscripts.\n",
"\n",
"Since no external forces or joint friction are present, we can extend a PD controller with a feed-forward term that includes gravity compensation:\n",
"\n",
"$$\n",
"\\begin{cases}\n",
"\\boldsymbol{\\tau} &= M_{ss} \\, \\ddot{\\mathbf{s}}^* + \\mathbf{h}_s \\\\\n",
"\\ddot{\\mathbf{s}}^* &= \\ddot{\\mathbf{s}}^\\text{des} - k_p(\\mathbf{s} - \\mathbf{s}^{\\text{des}}) - k_d(\\mathbf{s}^{\\text{des}} - \\dot{\\mathbf{s}}^{\\text{des}})\n",
"\\end{cases}\n",
"\\quad\n",
",\n",
"$$\n",
"\n",
"where $\\tilde{\\mathbf{s}} = \\left(\\mathbf{s} - \\mathbf{s}^\\text{des}\\right)$ is the joint position error.\n",
"\n",
"With this control law, the closed-loop system dynamics simplifies to:\n",
"\n",
"$$\n",
"\\ddot{\\tilde{\\mathbf{s}}} = -k_p \\tilde{\\mathbf{s}} - k_d \\dot{\\tilde{\\mathbf{s}}}\n",
",\n",
"$$\n",
"\n",
"which converges asymptotically to zero, ensuring stability."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rfMTCMyGZH4q"
},
"outputs": [],
"source": [
"# @title Create the computed torque controller\n",
"\n",
"# Define the PD gains\n",
"kp = 10.0\n",
"kd = 6.0\n",
"\n",
"\n",
"def computed_torque_controller(\n",
" data: js.data.JaxSimModelData,\n",
" s_des: jax.Array,\n",
" s_dot_des: jax.Array,\n",
") -> jax.Array:\n",
"\n",
" # Compute the gravity compensation term.\n",
" hs = js.model.free_floating_bias_forces(model=model, data=data)[6:]\n",
"\n",
" # Compute the joint-related portion of the floating-base mass matrix.\n",
" Mss = js.model.free_floating_mass_matrix(model=model, data=data)[6:, 6:]\n",
"\n",
" # Get the current joint positions and velocities.\n",
" s = data.joint_positions\n",
" ṡ = data.joint_velocities\n",
"\n",
" # Compute the actuated joint torques.\n",
" s_star = -kp * (s - s_des) - kd * (ṡ - s_dot_des)\n",
" τ = Mss @ s_star + hs\n",
"\n",
" return τ"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ERAUisywZH4q"
},
"source": [
"Now, we can use the `pd_controller` function to compute the torque to apply to the cartpole. Our aim is to stabilize the cartpole in the upright position, so we set the desired position `q_d` to 0 and the desired velocity `q_dot_d` to 0."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8YmDdGDVZH4q"
},
"outputs": [],
"source": [
"# @title Run the simulation\n",
"\n",
"# Initialize the data.\n",
"\n",
"# Set the joint positions.\n",
"data = js.data.JaxSimModelData.build(model=model, joint_positions=jnp.array([-0.25, jnp.deg2rad(160)]), joint_velocities=jnp.array([3.00, jnp.deg2rad(10) / model.time_step]))\n",
"\n",
"for _ in T:\n",
"\n",
" # Get the actuated torques from the computed torque controller.\n",
" τ = computed_torque_controller(\n",
" data=data,\n",
" s_des=jnp.array([0.0, 0.0]),\n",
" s_dot_des=jnp.array([0.0, 0.0]),\n",
" )\n",
"\n",
" # Step the JaxSim simulation.\n",
" data = js.model.step(\n",
" model=model,\n",
" data=data,\n",
" joint_force_references=τ,\n",
" )\n",
"\n",
" # Update the MuJoCo data.\n",
" mj_model_helper.set_joint_positions(\n",
" positions=data.joint_positions, joint_names=model.joint_names()\n",
" )\n",
"\n",
" # Record a new video frame.\n",
" recorder.record_frame(camera_name=\"cartpole_camera\")\n",
"\n",
"media.show_video(recorder.frames, fps=recorder.fps)\n",
"recorder.frames = []"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sZ76QqeWeMQz"
},
"source": [
"## Conclusions\n",
"\n",
"In this notebook, we explored how to use JaxSim for developing a closed-loop controller for a robot model. Key takeaways include:\n",
"\n",
"- We performed an open-loop simulation to understand the dynamics of the system without control.\n",
"- We implemented a computed torque controller with PD feedback and a feed-forward gravity compensation term, enabling the stabilization of the system by controlling joint torques.\n",
"- The closed-loop simulation can leverage hardware acceleration on GPUs and TPUs, with the ability to use `jax.vmap` for parallel sampling through automatic vectorization.\n",
"\n",
"JaxSim's closed-loop support can be extended to more advanced, model-based reactive controllers and planners for trajectory optimization. To explore optimization-based methods, consider the following JAX-based projects for hardware-accelerated control and planning:\n",
"\n",
"- [`deepmind/optax`](https://github.com/google-deepmind/optax)\n",
"- [`google/jaxopt`](https://github.com/google/jaxopt)\n",
"- [`patrick-kidger/lineax`](https://github.com/patrick-kidger/lineax)\n",
"- [`patrick-kidger/optimistix`](https://github.com/patrick-kidger/optimistix)\n",
"- [`kevin-tracy/qpax`](https://github.com/kevin-tracy/qpax)\n",
"\n",
"Additionally, if your controllers or planners require the derivatives of the dynamics with respect to the state or inputs, you can obtain them using automatic differentiation directly through JaxSim's API."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuClass": "premium",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "comodo_jaxsim",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.8"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
================================================
FILE: pyproject.toml
================================================
[project]
name = "jaxsim"
dynamic = ["version"]
requires-python = ">= 3.10"
description = "A differentiable physics engine and multibody dynamics library for control and robot learning."
authors = [
{ name = "Diego Ferigo", email = "dgferigo@gmail.com" },
{ name = "Filippo Luca Ferretti", email = "filippoluca.ferretti@outlook.com" },
]
maintainers = [
{ name = "Filippo Luca Ferretti", email = "filippo.ferretti@outlook.com" },
]
license = "BSD-3-Clause"
license-files = ["LICENSE"]
keywords = [
"physics",
"physics engine",
"jax",
"rigid body dynamics",
"featherstone",
"reinforcement learning",
"robot",
"robotics",
"sdf",
"urdf",
]
classifiers = [
"Development Status :: 4 - Beta",
"Framework :: Robot Framework",
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"Operating System :: POSIX :: Linux",
"Operating System :: MacOS",
"Operating System :: Microsoft",
"Programming Language :: Python :: 3 :: Only",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
"Programming Language :: Python :: Implementation :: CPython",
"Topic :: Games/Entertainment :: Simulation",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Scientific/Engineering :: Physics",
"Topic :: Software Development",
]
dependencies = [
"coloredlogs",
"jax >= 0.4.34",
"jaxlib >= 0.4.34",
"jaxlie >= 1.3.0",
"jax_dataclasses >= 1.4.0",
"pptree",
"optax >= 0.2.3",
"qpax",
"rod >= 0.4.1",
"typing_extensions ; python_version < '3.12'",
"trimesh",
]
[project.optional-dependencies]
testing = [
"chex >= 0.1.91",
"idyntree >= 12.2.1",
"pytest >=6.0",
"pytest-benchmark",
"pytest-icdiff",
"pytest-xdist",
"robot-descriptions >= 1.16.0",
"icub-models",
]
viz = [
"lxml",
"mediapy",
"mujoco >= 3.0.0",
"scipy >= 1.14.0",
]
all = [
"jaxsim[testing,viz]",
]
[project.readme]
file = "README.md"
content-type = "text/markdown"
[project.urls]
Changelog = "https://github.com/gbionics/jaxsim/releases"
Documentation = "https://jaxsim.readthedocs.io"
Source = "https://github.com/gbionics/jaxsim"
Tracker = "https://github.com/gbionics/jaxsim/issues"
# ===========
# Build tools
# ===========
[build-system]
build-backend = "hatchling.build"
requires = [
"hatchling",
"hatch-vcs",
]
[tool.hatch.version]
source = "vcs"
raw-options = { local_scheme = "dirty-tag" }
[tool.hatch.build.targets.wheel]
packages = ["src/jaxsim"]
[tool.hatch.build.hooks.vcs]
version-file = "src/jaxsim/_version.py"
# =================
# Style and testing
# =================
[tool.black]
line-length = 88
[tool.isort]
multi_line_output = 3
profile = "black"
[tool.pytest.ini_options]
addopts = "-rsxX -v --strict-markers --benchmark-skip --benchmark-warmup=ON"
minversion = "6.0"
testpaths = [
"tests",
]
# ==================
# Ruff configuration
# ==================
[tool.ruff]
exclude = [
".git",
".pixi",
".pytest_cache",
".ruff_cache",
".idea",
".vscode",
".devcontainer",
"__pycache__",
]
preview = true
[tool.ruff.lint]
# https://docs.astral.sh/ruff/rules/
select = [
"B",
"D",
"E",
"F",
"I",
"W",
"RUF",
"UP",
"YTT",
]
ignore = [
"B008", # Function call in default argument
"B024", # Abstract base class without abstract methods
"D100", # Missing docstring in public module
"D104", # Missing docstring in public package
"D105", # Missing docstring in magic method
"D200", # One-line docstring should fit on one line with quotes
"D202", # No blank lines allowed after function docstring
"D203", # Incorrect blank line before class
"D205", # 1 blank line required between summary line and description
"D212", # Multi-line docstring summary should start at the first line
"D411", # Missing blank line before section
"D413", # Missing blank line after last section
"E402", # Module level import not at top of file
"E501", # Line too long
"E731", # Do not assign a `lambda` expression, use a `def`
"E741", # Ambiguous variable name
"I001", # Import block is unsorted or unformatted
"RUF003", # Ambiguous unicode character in comment
]
[tool.ruff.lint.per-file-ignores]
# Ignore `E402` (import violations) in all `__init__.py` files
"**/{tests,docs,tools}/*" = ["E402"]
"**/{tests,examples}/*" = ["B007", "D100", "D102", "D103"]
"__init__.py" = ["F401", "RUF067"]
"docs/conf.py" = ["F401"]
"src/jaxsim/exceptions.py" = ["D401"]
"src/jaxsim/logging.py" = ["D101", "D103"]
# ==================
# Pixi configuration
# ==================
[tool.pixi.workspace]
channels = ["conda-forge"]
platforms = ["linux-64", "linux-aarch64", "osx-arm64", "osx-64"]
requires-pixi = ">=0.39.0"
preview = ["pixi-build"]
[tool.pixi.environments]
# We resolve only two groups: cpu and gpu.
# Then, multiple environments can be created from these groups.
default = { features = ["test", "examples"] }
gpu = { features = ["test", "examples", "gpu"] }
# ---------------
# feature.default
# ---------------
# Dependencies from conda-forge.
[tool.pixi.dependencies]
#
# Matching `project.dependencies`.
#
coloredlogs = "*"
jax = "*"
jaxlib = "*"
jaxlie = "*"
jax-dataclasses = "*"
pptree = "*"
optax = "*"
qpax = "*"
rod = ">=0.4.1"
trimesh = "*"
typing_extensions = "*"
#
# Optional dependencies.
#
lxml = "*"
mediapy = "*"
mujoco = "*"
scipy = "*"
#
# Additional dependencies.
#
pip = "*"
hatchling = "*"
hatch-vcs = "*"
# Dependencies from PyPI.
[tool.pixi.pypi-dependencies]
jaxsim = { path = "./", editable = true }
[tool.pixi.pypi-options]
no-build-isolation = ["jaxsim"]
# ------------
# feature.test
# ------------
[tool.pixi.feature.test.tasks]
pipcheck = "pip check"
benchmark = { cmd = "pytest --benchmark-only --benchmark-warmup=ON", depends-on = ["pipcheck"] }
test = { cmd = "pytest", depends-on = ["pipcheck"] }
[tool.pixi.feature.test.dependencies]
black-jupyter = "*"
chex = ">=0.1.91"
idyntree = "*"
isort = "*"
pre-commit = "*"
pytest = "*"
pytest-benchmark = "*"
pytest-icdiff = "*"
pytest-xdist = "*"
robot_descriptions = ">=1.16.0"
# ----------------
# feature.examples
# ----------------
[tool.pixi.feature.examples.tasks]
examples = { cmd = "jupyter notebook ./examples" }
[tool.pixi.feature.examples.dependencies]
notebook = "*"
robot_descriptions = ">=1.16.0"
# -----------
# feature.gpu
# -----------
[tool.pixi.feature.gpu]
platforms = ["linux-64"]
system-requirements = { cuda = "13" }
[tool.pixi.feature.gpu.dependencies]
jaxlib = { version = "*", build = "*cuda*" }
[tool.pixi.feature.gpu.tasks]
test-gpu = { cmd = "pytest --gpu-only", depends-on = ["pipcheck"] }
================================================
FILE: src/jaxsim/__init__.py
================================================
from . import logging
from ._version import __version__
# Follow upstream development in https://github.com/google/jax/pull/13304
def _jnp_options() -> None:
import os
import jax
# Check if running on TPU.
is_tpu = jax.devices()[0].platform == "tpu"
# Check if running on Metal.
is_metal = jax.devices()[0].platform == "METAL"
# Enable by default 64-bit precision to get accurate physics.
# Users can enforce 32-bit precision by setting the following variable to 0.
use_x64 = os.environ.get("JAX_ENABLE_X64", "1") != "0"
# Notify the user if unsupported 64-bit precision was enforced on TPU.
if (is_tpu or is_metal) and use_x64:
msg = f"64-bit precision is not allowed on {jax.devices()[0].platform.upper}. Enforcing 32bit precision."
logging.warning(msg)
use_x64 = False
if is_metal:
logging.warning(
"JAX Metal backend is experimental. Some functionalities may not be available."
)
# Enable 64-bit precision in JAX.
if use_x64:
logging.info("Enabling JAX to use 64-bit precision")
jax.config.update("jax_enable_x64", True)
# Warn about experimental usage of 32-bit precision.
else:
logging.warning(
"Using 32-bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
)
def _np_options() -> None:
import numpy as np
np.set_printoptions(precision=5, suppress=True, linewidth=150, threshold=10_000)
def _is_editable() -> bool:
import importlib.util
import pathlib
import site
# Get the ModuleSpec of jaxsim.
jaxsim_spec = importlib.util.find_spec(name="jaxsim")
# This can be None. If it's None, assume non-editable installation.
if jaxsim_spec.origin is None:
return False
# Get the folder containing the jaxsim package.
jaxsim_package_dir = str(pathlib.Path(jaxsim_spec.origin).parent.parent)
# The installation is editable if the package dir is not in any {site|dist}-packages.
return jaxsim_package_dir not in site.getsitepackages()
def _get_default_logging_level() -> logging.LoggingLevel:
"""
Get the default logging level.
Returns:
The logging level to set.
"""
import os
import sys
# Allow to override the default logging level with an environment variable.
if overriden_logging_level := os.environ.get("JAXSIM_LOGGING_LEVEL"):
try:
return logging.LoggingLevel[overriden_logging_level.upper()]
except KeyError as exc:
msg = "Invalid logging level defined in JAXSIM_LOGGING_LEVEL"
raise RuntimeError(msg) from exc
# If running under a debugger, set the logging level to DEBUG.
if getattr(sys, "gettrace", lambda: None)():
return logging.LoggingLevel.DEBUG
# If not running under a debugger, set the logging level to INFO or WARNING.
# INFO for editable installations, WARNING for non-editable installations.
# This is to avoid too verbose logging in non-editable installations.
return (
logging.LoggingLevel.INFO
if _is_editable() # noqa: F821
else logging.LoggingLevel.WARNING
)
# Configure the logger with the default logging level.
logging.configure(level=_get_default_logging_level())
# Configure JAX.
_jnp_options()
# Initialize the numpy print options.
_np_options()
del _jnp_options
del _np_options
del _get_default_logging_level
del _is_editable
from . import terrain # isort:skip
from . import api, logging, math, rbda
from .api.common import VelRepr
================================================
FILE: src/jaxsim/api/__init__.py
================================================
from . import common # isort:skip
from . import model, data # isort:skip
from . import (
actuation_model,
com,
contact,
frame,
integrators,
joint,
kin_dyn_parameters,
link,
ode,
references,
)
================================================
FILE: src/jaxsim/api/actuation_model.py
================================================
import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.typing as jtp
def compute_resultant_torques(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
joint_force_references: jtp.Vector | None = None,
) -> jtp.Vector:
"""
Compute the resultant torques acting on the joints.
Args:
model: The model to consider.
data: The data of the considered model.
joint_force_references: The joint force references to apply.
Returns:
The resultant torques acting on the joints.
"""
# Build joint torques if not provided.
τ_references = (
jnp.atleast_1d(joint_force_references.squeeze())
if joint_force_references is not None
else jnp.zeros_like(data.joint_positions)
).astype(float)
# ====================
# Enforce joint limits
# ====================
τ_position_limit = jnp.zeros_like(τ_references).astype(float)
if model.dofs() > 0:
# Stiffness and damper parameters for the joint position limits.
k_j = jnp.array(
model.kin_dyn_parameters.joint_parameters.position_limit_spring
).astype(float)
d_j = jnp.array(
model.kin_dyn_parameters.joint_parameters.position_limit_damper
).astype(float)
# Compute the joint position limit violations.
lower_violation = jnp.clip(
data.joint_positions
- model.kin_dyn_parameters.joint_parameters.position_limits_min,
max=0.0,
)
upper_violation = jnp.clip(
data.joint_positions
- model.kin_dyn_parameters.joint_parameters.position_limits_max,
min=0.0,
)
# Compute the joint position limit torque.
τ_position_limit -= jnp.diag(k_j) @ (lower_violation + upper_violation)
τ_position_limit -= (
jnp.positive(τ_position_limit) * jnp.diag(d_j) @ data.joint_velocities
)
# ====================
# Joint friction model
# ====================
τ_friction = jnp.zeros_like(τ_references).astype(float)
# Apply joint friction only if enabled in the actuation parameters.
if model.dofs() > 0 and model.actuation_params.enable_friction:
# Static and viscous joint friction parameters
kc = jnp.array(
model.kin_dyn_parameters.joint_parameters.friction_static
).astype(float)
kv = jnp.array(
model.kin_dyn_parameters.joint_parameters.friction_viscous
).astype(float)
# Compute the joint friction torque.
τ_friction = -(
jnp.diag(kc) @ jnp.sign(data.joint_velocities)
+ jnp.diag(kv) @ data.joint_velocities
)
# ===============================
# Compute the total joint forces.
# ===============================
τ_total = τ_references + τ_friction + τ_position_limit
τ_lim = tn_curve_fn(model=model, data=data)
τ_total = jnp.clip(τ_total, -τ_lim, τ_lim)
return τ_total
def tn_curve_fn(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
"""
Compute the torque limits using the tn curve.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The torque limits.
"""
τ_max = model.actuation_params.torque_max # Max torque (Nm)
ω_th = model.actuation_params.omega_th # Threshold speed (rad/s)
ω_max = model.actuation_params.omega_max # Max speed for torque drop-off (rad/s)
abs_vel = jnp.abs(data.joint_velocities)
τ_lim = jnp.where(
abs_vel <= ω_th,
τ_max,
jnp.where(
abs_vel <= ω_max, τ_max * (1 - (abs_vel - ω_th) / (ω_max - ω_th)), 0.0
),
)
return τ_lim
================================================
FILE: src/jaxsim/api/com.py
================================================
import jax
import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.math
import jaxsim.typing as jtp
from .common import VelRepr
@jax.jit
@js.common.named_scope
def com_position(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
"""
Compute the position of the center of mass of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The position of the center of mass of the model w.r.t. the world frame.
"""
m = js.model.total_mass(model=model)
W_H_L = data._link_transforms
W_H_B = data._base_transform
B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
def B_p̃_LCoM(i) -> jtp.Vector:
m = js.link.mass(model=model, link_index=i)
L_p_LCoM = js.link.com_position(
model=model, data=data, link_index=i, in_link_frame=True
)
return m * B_H_W @ W_H_L[i] @ jnp.hstack([L_p_LCoM, 1])
com_links = jax.vmap(B_p̃_LCoM)(jnp.arange(model.number_of_links()))
B_p̃_CoM = (1 / m) * com_links.sum(axis=0)
B_p̃_CoM = B_p̃_CoM.at[3].set(1)
return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
@jax.jit
@js.common.named_scope
def com_linear_velocity(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
r"""
Compute the linear velocity of the center of mass of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The linear velocity of the center of mass of the model in the
active representation.
Note:
The linear velocity of the center of mass is expressed in the mixed frame
:math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
active velocity representation is either inertial-fixed or mixed,
and :math:`[C] = [B]` if the active velocity representation is body-fixed.
"""
# Extract the linear component of the 6D average centroidal velocity.
# This is expressed in G[B] in body-fixed representation, and in G[W] in
# inertial-fixed or mixed representation.
G_vl_WG = average_centroidal_velocity(model=model, data=data)[0:3]
return G_vl_WG
@jax.jit
@js.common.named_scope
def centroidal_momentum(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
r"""
Compute the centroidal momentum of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The centroidal momentum of the model.
Note:
The centroidal momentum is expressed in the mixed frame
:math:`({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`C = W` if the
active velocity representation is either inertial-fixed or mixed,
and :math:`C = B` if the active velocity representation is body-fixed.
"""
ν = data.generalized_velocity
G_J = centroidal_momentum_jacobian(model=model, data=data)
return G_J @ ν
@jax.jit
@js.common.named_scope
def centroidal_momentum_jacobian(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Matrix:
r"""
Compute the Jacobian of the centroidal momentum of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The Jacobian of the centroidal momentum of the model.
Note:
The frame corresponding to the output representation of this Jacobian is either
:math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
or :math:`G[B]`, if the active velocity representation is body-fixed.
Note:
This Jacobian is also known in the literature as Centroidal Momentum Matrix.
"""
# Compute the Jacobian of the total momentum with body-fixed output representation.
# We convert the output representation either to G[W] or G[B] below.
B_Jh = js.model.total_momentum_jacobian(
model=model, data=data, output_vel_repr=VelRepr.Body
)
W_H_B = data._base_transform
B_H_W = jaxsim.math.Transform.inverse(W_H_B)
W_p_CoM = com_position(model=model, data=data)
match data.velocity_representation:
case VelRepr.Inertial | VelRepr.Mixed:
W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM) # noqa: F841
case VelRepr.Body:
W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM) # noqa: F841
case _:
raise ValueError(data.velocity_representation)
# Compute the transform for 6D forces.
G_Xf_B = jaxsim.math.Adjoint.from_transform(transform=B_H_W @ W_H_G).T
return G_Xf_B @ B_Jh
@jax.jit
@js.common.named_scope
def locked_centroidal_spatial_inertia(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
):
"""
Compute the locked centroidal spatial inertia of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The locked centroidal spatial inertia of the model.
"""
with data.switch_velocity_representation(VelRepr.Body):
B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
W_H_B = data._base_transform
W_p_CoM = com_position(model=model, data=data)
match data.velocity_representation:
case VelRepr.Inertial | VelRepr.Mixed:
W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM) # noqa: F841
case VelRepr.Body:
W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM) # noqa: F841
case _:
raise ValueError(data.velocity_representation)
B_H_G = jaxsim.math.Transform.inverse(W_H_B) @ W_H_G
B_Xv_G = jaxsim.math.Adjoint.from_transform(transform=B_H_G)
G_Xf_B = B_Xv_G.transpose()
return G_Xf_B @ B_Mbb_B @ B_Xv_G
@jax.jit
@js.common.named_scope
def average_centroidal_velocity(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
r"""
Compute the average centroidal velocity of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The average centroidal velocity of the model.
Note:
The average velocity is expressed in the mixed frame
:math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
active velocity representation is either inertial-fixed or mixed,
and :math:`[C] = [B]` if the active velocity representation is body-fixed.
"""
ν = data.generalized_velocity
G_J = average_centroidal_velocity_jacobian(model=model, data=data)
return G_J @ ν
@jax.jit
@js.common.named_scope
def average_centroidal_velocity_jacobian(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Matrix:
r"""
Compute the Jacobian of the average centroidal velocity of the model.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The Jacobian of the average centroidal velocity of the model.
Note:
The frame corresponding to the output representation of this Jacobian is either
:math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
or :math:`G[B]`, if the active velocity representation is body-fixed.
"""
G_J = centroidal_momentum_jacobian(model=model, data=data)
G_Mbb = locked_centroidal_spatial_inertia(model=model, data=data)
return jnp.linalg.inv(G_Mbb) @ G_J
@jax.jit
@js.common.named_scope
def bias_acceleration(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Vector:
r"""
Compute the bias linear acceleration of the center of mass.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The bias linear acceleration of the center of mass in the active representation.
Note:
The bias acceleration is expressed in the mixed frame
:math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
active velocity representation is either inertial-fixed or mixed,
and :math:`[C] = [B]` if the active velocity representation is body-fixed.
"""
# Compute the pose of all links with forward kinematics.
W_H_L = data._link_transforms
# Compute the bias acceleration of all links by zeroing the generalized velocity
# in the active representation.
v̇_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
def other_representation_to_body(
C_v̇_WL: jtp.Vector, C_v_WC: jtp.Vector, L_H_C: jtp.Matrix, L_v_LC: jtp.Vector
) -> jtp.Vector:
"""
Convert the body-fixed representation of the link bias acceleration
C_v̇_WL expressed in a generic frame C to the body-fixed representation L_v̇_WL.
"""
L_X_C = jaxsim.math.Adjoint.from_transform(transform=L_H_C)
C_X_L = jaxsim.math.Adjoint.inverse(L_X_C)
L_v̇_WL = L_X_C @ (C_v̇_WL + jaxsim.math.Cross.vx(C_X_L @ L_v_LC) @ C_v_WC)
return L_v̇_WL
# We need here to get the body-fixed bias acceleration of the links.
# Since it's computed in the active representation, we need to convert it to body.
match data.velocity_representation:
case VelRepr.Body:
L_a_bias_WL = v̇_bias_WL
case VelRepr.Inertial:
C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL # noqa: F841
C_v_WC = W_v_WW = jnp.zeros(6) # noqa: F841
L_H_C = L_H_W = jax.vmap(jaxsim.math.Transform.inverse)(W_H_L) # noqa: F841
L_v_LC = L_v_LW = jax.vmap( # noqa: F841
lambda i: -js.link.velocity(
model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
)
)(jnp.arange(model.number_of_links()))
L_a_bias_WL = jax.vmap(
lambda i: other_representation_to_body(
C_v̇_WL=C_v̇_WL[i],
C_v_WC=C_v_WC,
L_H_C=L_H_C[i],
L_v_LC=L_v_LC[i],
)
)(jnp.arange(model.number_of_links()))
case VelRepr.Mixed:
C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL # noqa: F841
C_v_WC = LW_v_W_LW = jax.vmap( # noqa: F841
lambda i: js.link.velocity(
model=model, data=data, link_index=i, output_vel_repr=VelRepr.Mixed
)
.at[3:6]
.set(jnp.zeros(3))
)(jnp.arange(model.number_of_links()))
L_H_C = L_H_LW = jax.vmap( # noqa: F841
lambda W_H_L: jaxsim.math.Transform.inverse(
W_H_L.at[0:3, 3].set(jnp.zeros(3))
)
)(W_H_L)
L_v_LC = L_v_L_LW = jax.vmap( # noqa: F841
lambda i: -js.link.velocity(
model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
)
.at[0:3]
.set(jnp.zeros(3))
)(jnp.arange(model.number_of_links()))
L_a_bias_WL = jax.vmap(
lambda i: other_representation_to_body(
C_v̇_WL=C_v̇_WL[i],
C_v_WC=C_v_WC[i],
L_H_C=L_H_C[i],
L_v_LC=L_v_LC[i],
)
)(jnp.arange(model.number_of_links()))
case _:
raise ValueError(data.velocity_representation)
# Compute the bias of the 6D momentum derivative.
def bias_momentum_derivative_term(
link_index: jtp.Int, L_a_bias_WL: jtp.Vector
) -> jtp.Vector:
# Get the body-fixed 6D inertia matrix.
L_M_L = js.link.spatial_inertia(model=model, link_index=link_index)
# Compute the body-fixed 6D velocity.
L_v_WL = js.link.velocity(
model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body
)
# Compute the world-to-link transformations for 6D forces.
W_Xf_L = jaxsim.math.Adjoint.from_transform(
transform=W_H_L[link_index], inverse=True
).T
# Compute the contribution of the link to the bias acceleration of the CoM.
W_ḣ_bias_link_contribution = W_Xf_L @ (
L_M_L @ L_a_bias_WL + jaxsim.math.Cross.vx_star(L_v_WL) @ L_M_L @ L_v_WL
)
return W_ḣ_bias_link_contribution
# Sum the contributions of all links to the bias acceleration of the CoM.
W_ḣ_bias = jax.vmap(bias_momentum_derivative_term)(
jnp.arange(model.number_of_links()), L_a_bias_WL
).sum(axis=0)
# Compute the total mass of the model.
m = js.model.total_mass(model=model)
# Compute the position of the CoM.
W_p_CoM = com_position(model=model, data=data)
match data.velocity_representation:
# G := G[W] = (W_p_CoM, [W])
case VelRepr.Inertial | VelRepr.Mixed:
W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
GW_Xf_W = jaxsim.math.Adjoint.from_transform(W_H_GW).T
GW_ḣ_bias = GW_Xf_W @ W_ḣ_bias
GW_v̇l_com_bias = GW_ḣ_bias[0:3] / m
return GW_v̇l_com_bias
# G := G[B] = (W_p_CoM, [B])
case VelRepr.Body:
GB_Xf_W = jaxsim.math.Adjoint.from_transform(
transform=data._base_transform.at[0:3].set(W_p_CoM)
).T
GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias
GB_v̇l_com_bias = GB_ḣ_bias[0:3] / m
return GB_v̇l_com_bias
case _:
raise ValueError(data.velocity_representation)
================================================
FILE: src/jaxsim/api/common.py
================================================
import abc
import contextlib
import dataclasses
import enum
import functools
from collections.abc import Callable, Iterator
from typing import ParamSpec, TypeVar
import jax
import jax.numpy as jnp
import jax_dataclasses
from jax_dataclasses import Static
import jaxsim.typing as jtp
from jaxsim.math import Adjoint
from jaxsim.utils import JaxsimDataclass, Mutability
try:
from typing import Self
except ImportError:
from typing_extensions import Self
_P = ParamSpec("_P")
_R = TypeVar("_R")
def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
"""Apply a JAX named scope to a function for improved profiling and clarity."""
@functools.wraps(fn)
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
with jax.named_scope(name or fn.__name__):
return fn(*args, **kwargs)
return wrapper
@enum.unique
class VelRepr(enum.IntEnum):
"""
Enumeration of all supported 6D velocity representations.
"""
Body = enum.auto()
Mixed = enum.auto()
Inertial = enum.auto()
@jax_dataclasses.pytree_dataclass
class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
"""
Base class for model data structures with velocity representation.
"""
velocity_representation: Static[VelRepr] = dataclasses.field(
default=VelRepr.Inertial, kw_only=True
)
@contextlib.contextmanager
def switch_velocity_representation(
self, velocity_representation: VelRepr
) -> Iterator[Self]:
"""
Context manager to temporarily switch the velocity representation.
Args:
velocity_representation: The new velocity representation.
Yields:
The same object with the new velocity representation.
"""
original_representation = self.velocity_representation
try:
# First, we replace the velocity representation.
with self.mutable_context(
mutability=Mutability.MUTABLE_NO_VALIDATION,
restore_after_exception=True,
):
self.velocity_representation = velocity_representation
# Then, we yield the data with changed representation.
# We run this in a mutable context with restoration so that any exception
# occurring, we restore the original object in case it was modified.
with self.mutable_context(
mutability=self.mutability(), restore_after_exception=True
):
yield self
finally:
with self.mutable_context(
mutability=Mutability.MUTABLE_NO_VALIDATION,
restore_after_exception=True,
):
self.velocity_representation = original_representation
@staticmethod
@functools.partial(jax.jit, static_argnames=["other_representation", "is_force"])
def inertial_to_other_representation(
array: jtp.Array,
other_representation: VelRepr,
transform: jtp.Matrix,
*,
is_force: bool,
) -> jtp.Array:
r"""
Convert a 6D quantity from inertial-fixed to another representation.
Args:
array: The 6D quantity to convert.
other_representation: The representation to convert to.
transform:
The :math:`W \mathbf{H}_O` transform, where :math:`O` is the
reference frame of the other representation.
is_force: Whether the quantity is a 6D force or a 6D velocity.
Returns:
The 6D quantity in the other representation.
"""
W_array = array
W_H_O = transform
match other_representation:
case VelRepr.Inertial:
return W_array
case VelRepr.Body:
if not is_force:
O_Xv_W = Adjoint.from_transform(transform=W_H_O, inverse=True)
O_array = jnp.einsum("...ij,...j->...i", O_Xv_W, W_array)
else:
O_Xf_W = Adjoint.from_transform(transform=W_H_O).swapaxes(-1, -2)
O_array = jnp.einsum("...ij,...j->...i", O_Xf_W, W_array)
return O_array
case VelRepr.Mixed:
W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
if not is_force:
OW_Xv_W = Adjoint.from_transform(transform=W_H_OW, inverse=True)
OW_array = jnp.einsum("...ij,...j->...i", OW_Xv_W, W_array)
else:
OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).swapaxes(-1, -2)
OW_array = jnp.einsum("...ij,...j->...i", OW_Xf_W, W_array)
return OW_array
case _:
raise ValueError(other_representation)
@staticmethod
@functools.partial(jax.jit, static_argnames=["other_representation", "is_force"])
def other_representation_to_inertial(
array: jtp.Array,
other_representation: VelRepr,
transform: jtp.Matrix,
*,
is_force: bool,
) -> jtp.Array:
r"""
Convert a 6D quantity from another representation to inertial-fixed.
Args:
array: The 6D quantity to convert.
other_representation: The representation to convert from.
transform:
The `math:W \mathbf{H}_O` transform, where `math:O` is the
reference frame of the other representation.
is_force: Whether the quantity is a 6D force or a 6D velocity.
Returns:
The 6D quantity in the inertial-fixed representation.
"""
O_array = array
W_H_O = transform
match other_representation:
case VelRepr.Inertial:
return O_array
case VelRepr.Body:
if not is_force:
W_Xv_O = Adjoint.from_transform(W_H_O)
W_array = jnp.einsum("...ij,...j->...i", W_Xv_O, O_array)
else:
W_Xf_O = Adjoint.from_transform(
transform=W_H_O, inverse=True
).swapaxes(-1, -2)
W_array = jnp.einsum("...ij,...j->...i", W_Xf_O, O_array)
return W_array
case VelRepr.Mixed:
W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
if not is_force:
W_Xv_BW = Adjoint.from_transform(W_H_OW)
W_array = jnp.einsum("...ij,...j->...i", W_Xv_BW, O_array)
else:
W_Xf_BW = Adjoint.from_transform(
transform=W_H_OW, inverse=True
).swapaxes(-1, -2)
W_array = jnp.einsum("...ij,...j->...i", W_Xf_BW, O_array)
return W_array
case _:
raise ValueError(other_representation)
================================================
FILE: src/jaxsim/api/contact.py
================================================
from __future__ import annotations
import functools
import jax
import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.exceptions
import jaxsim.typing as jtp
from jaxsim import logging
from jaxsim.math import Adjoint, Cross, Transform
from jaxsim.rbda.contacts import SoftContacts
from .common import VelRepr
@jax.jit
@js.common.named_scope
def collidable_point_kinematics(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> tuple[jtp.Matrix, jtp.Matrix]:
"""
Compute the position and 3D velocity of the collidable points in the world frame.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The position and velocity of the collidable points in the world frame.
Note:
The collidable point velocity is the plain coordinate derivative of the position.
If we attach a frame C = (p_C, [C]) to the collidable point, it corresponds to
the linear component of the mixed 6D frame velocity.
"""
W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
model=model,
link_transforms=data._link_transforms,
link_velocities=data._link_velocities,
)
return W_p_Ci, W_ṗ_Ci
@jax.jit
@js.common.named_scope
def collidable_point_positions(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Matrix:
"""
Compute the position of the collidable points in the world frame.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The position of the collidable points in the world frame.
"""
W_p_Ci, _ = collidable_point_kinematics(model=model, data=data)
return W_p_Ci
@jax.jit
@js.common.named_scope
def collidable_point_velocities(
model: js.model.JaxSimModel, data: js.data.JaxSimModelData
) -> jtp.Matrix:
"""
Compute the 3D velocity of the collidable points in the world frame.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The 3D velocity of the collidable points.
"""
_, W_ṗ_Ci = collidable_point_kinematics(model=model, data=data)
return W_ṗ_Ci
@functools.partial(jax.jit, static_argnames=["link_names"])
@js.common.named_scope
def in_contact(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
link_names: tuple[str, ...] | None = None,
) -> jtp.Vector:
"""
Return whether the links are in contact with the terrain.
Args:
model: The model to consider.
data: The data of the considered model.
link_names:
The names of the links to consider. If None, all links are considered.
Returns:
A boolean vector indicating whether the links are in contact with the terrain.
"""
if link_names is not None and set(link_names).difference(model.link_names()):
raise ValueError("One or more link names are not part of the model")
# Get the indices of the enabled collidable points.
indices_of_enabled_collidable_points = (
model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
)
parent_link_idx_of_enabled_collidable_points = jnp.array(
model.kin_dyn_parameters.contact_parameters.body, dtype=int
)[indices_of_enabled_collidable_points]
W_p_Ci = collidable_point_positions(model=model, data=data)
terrain_height = jax.vmap(lambda x, y: model.terrain.height(x=x, y=y))(
W_p_Ci[:, 0], W_p_Ci[:, 1]
)
below_terrain = W_p_Ci[:, 2] <= terrain_height
link_idxs = (
js.link.names_to_idxs(link_names=link_names, model=model)
if link_names is not None
else jnp.arange(model.number_of_links())
)
links_in_contact = jax.vmap(
lambda link_index: jnp.where(
parent_link_idx_of_enabled_collidable_points == link_index,
below_terrain,
jnp.zeros_like(below_terrain, dtype=bool),
).any()
)(link_idxs)
return links_in_contact
def estimate_good_soft_contacts_parameters(
*args, **kwargs
) -> jaxsim.rbda.contacts.ContactParamsTypes:
"""
Estimate good soft contacts parameters. Deprecated, use `estimate_good_contact_parameters` instead.
"""
msg = "This method is deprecated, please use `{}`."
logging.warning(msg.format(estimate_good_contact_parameters.__name__))
return estimate_good_contact_parameters(*args, **kwargs)
def estimate_good_contact_parameters(
model: js.model.JaxSimModel,
*,
standard_gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
static_friction_coefficient: jtp.FloatLike = 0.5,
number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
damping_ratio: jtp.FloatLike = 1.0,
max_penetration: jtp.FloatLike | None = None,
) -> jaxsim.rbda.contacts.ContactParamsTypes:
"""
Estimate good contact parameters.
Args:
model: The model to consider.
standard_gravity: The standard gravity acceleration.
static_friction_coefficient: The static friction coefficient.
number_of_active_collidable_points_steady_state:
The number of active collidable points in steady state.
damping_ratio: The damping ratio.
max_penetration: The maximum penetration allowed.
Returns:
The estimated good contacts parameters.
Note:
This is primarily a convenience function for soft-like contact models.
However, it provides with some good default parameters also for the other ones.
Note:
This method provides a good set of contacts parameters.
The user is encouraged to fine-tune the parameters based on the
specific application.
"""
if max_penetration is None:
zero_data = js.data.JaxSimModelData.build(model=model)
W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
if model.floating_base():
W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
W_pz_CoM = W_pz_CoM - W_pz_C.min()
# Consider as default a 1% of the model center of mass height.
max_penetration = 0.01 * W_pz_CoM
nc = number_of_active_collidable_points_steady_state
return model.contact_model._parameters_class().build_default_from_jaxsim_model(
model=model,
standard_gravity=standard_gravity,
static_friction_coefficient=static_friction_coefficient,
max_penetration=max_penetration,
number_of_active_collidable_points_steady_state=nc,
damping_ratio=damping_ratio,
)
@jax.jit
@js.common.named_scope
def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
r"""
Return the pose of the enabled collidable points.
Args:
model: The model to consider.
data: The data of the considered model.
Returns:
The stacked SE(3) matrices of all enabled collidable points.
Note:
Each collidable point is implicitly associated with a frame
:math:`C = ({}^W p_C, [L])`, where :math:`{}^W p_C` is the position of the
collidable point and :math:`[L]` is the orientation frame of the link it is
rigidly attached to.
"""
# Get the indices of the enabled collidable points.
indices_of_enabled_collidable_points = (
model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
)
parent_link_idx_of_enabled_collidable_points = jnp.array(
model.kin_dyn_parameters.contact_parameters.body, dtype=int
)[indices_of_enabled_collidable_points]
# Get the transforms of the parent link of all collidable points.
W_H_L = data._link_transforms[parent_link_idx_of_enabled_collidable_points]
L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
indices_of_enabled_collidable_points
]
# Build the link-to-point transform from the displacement between the link frame L
# and the implicit contact frame C.
L_H_C = jax.vmap(jnp.eye(4).at[0:3, 3].set)(L_p_Ci)
# Compose the work-to-link and link-to-point transforms.
return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
@js.common.named_scope
def jacobian(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
output_vel_repr: VelRepr | None = None,
) -> jtp.Array:
r"""
Return the free-floating Jacobian of the enabled collidable points.
Args:
model: The model to consider.
data: The data of the considered model.
output_vel_repr:
The output velocity representation of the free-floating jacobian.
Returns:
The stacked :math:`6 \times (6+n)` free-floating jacobians of the frames associated to the
enabled collidable points.
Note:
Each collidable point is implicitly associated with a frame
:math:`C = ({}^W p_C, [L])`, where :math:`{}^W p_C` is the position of the
collidable point and :math:`[L]` is the orientation frame of the link it is
rigidly attached to.
"""
output_vel_repr = (
output_vel_repr if output_vel_repr is not None else data.velocity_representation
)
# Get the indices of the enabled collidable points.
indices_of_enabled_collidable_points = (
model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
)
parent_link_idx_of_enabled_collidable_points = jnp.array(
model.kin_dyn_parameters.contact_parameters.body, dtype=int
)[indices_of_enabled_collidable_points]
# Compute the Jacobians of all links.
W_J_WL = js.model.generalized_free_floating_jacobian(
model=model, data=data, output_vel_repr=VelRepr.Inertial
)
# Compute the contact Jacobian.
# In inertial-fixed output representation, the Jacobian of the parent link is also
# the Jacobian of the frame C implicitly associated with the collidable point.
W_J_WC = W_J_WL[parent_link_idx_of_enabled_collidable_points]
# Adjust the output representation.
match output_vel_repr:
case VelRepr.Inertial:
O_J_WC = W_J_WC
case VelRepr.Body:
W_H_C = transforms(model=model, data=data)
def body_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
C_X_W = jaxsim.math.Adjoint.from_transform(
transform=W_H_C, inverse=True
)
C_J_WC = C_X_W @ W_J_WC
return C_J_WC
O_J_WC = jax.vmap(body_jacobian)(W_H_C, W_J_WC)
case VelRepr.Mixed:
W_H_C = transforms(model=model, data=data)
def mixed_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
CW_X_W = jaxsim.math.Adjoint.from_transform(
transform=W_H_CW, inverse=True
)
CW_J_WC = CW_X_W @ W_J_WC
return CW_J_WC
O_J_WC = jax.vmap(mixed_jacobian)(W_H_C, W_J_WC)
case _:
raise ValueError(output_vel_repr)
return O_J_WC
@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
@js.common.named_scope
def jacobian_derivative(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
output_vel_repr: VelRepr | None = None,
) -> jtp.Matrix:
r"""
Compute the derivative of the free-floating jacobian of the enabled collidable points.
Args:
model: The model to consider.
data: The data of the considered model.
output_vel_repr:
The output velocity representation of the free-floating jacobian derivative.
Returns:
The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the enabled collidable points.
Note:
The input representation of the free-floating jacobian derivative is the active
velocity representation.
"""
output_vel_repr = (
output_vel_repr if output_vel_repr is not None else data.velocity_representation
)
indices_of_enabled_collidable_points = (
model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
)
# Get the index of the parent link and the position of the collidable point.
parent_link_idx_of_enabled_collidable_points = jnp.array(
model.kin_dyn_parameters.contact_parameters.body, dtype=int
)[indices_of_enabled_collidable_points]
L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
indices_of_enabled_collidable_points
]
# Get the transforms of all the parent links.
W_H_Li = data._link_transforms
# Get the link velocities.
W_v_WLi = data._link_velocities
# =====================================================
# Compute quantities to adjust the input representation
# =====================================================
def compute_T(model: js.model.JaxSimModel, X: jtp.Matrix) -> jtp.Matrix:
In = jnp.eye(model.dofs())
T = jax.scipy.linalg.block_diag(X, In)
return T
def compute_Ṫ(model: js.model.JaxSimModel, Ẋ: jtp.Matrix) -> jtp.Matrix:
On = jnp.zeros(shape=(model.dofs(), model.dofs()))
Ṫ = jax.scipy.linalg.block_diag(Ẋ, On)
return Ṫ
# Compute the operator to change the representation of ν, and its
# time derivative.
match data.velocity_representation:
case VelRepr.Inertial:
W_H_W = jnp.eye(4)
W_X_W = Adjoint.from_transform(transform=W_H_W)
W_Ẋ_W = jnp.zeros((6, 6))
T = compute_T(model=model, X=W_X_W)
Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
case VelRepr.Body:
W_H_B = data._base_transform
W_X_B = Adjoint.from_transform(transform=W_H_B)
B_v_WB = data.base_velocity
B_vx_WB = Cross.vx(B_v_WB)
W_Ẋ_B = W_X_B @ B_vx_WB
T = compute_T(model=model, X=W_X_B)
Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
case VelRepr.Mixed:
W_H_B = data._base_transform
W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
W_X_BW = Adjoint.from_transform(transform=W_H_BW)
BW_v_WB = data.base_velocity
BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
BW_vx_W_BW = Cross.vx(BW_v_W_BW)
W_Ẋ_BW = W_X_BW @ BW_vx_W_BW
T = compute_T(model=model, X=W_X_BW)
Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_BW)
case _:
raise ValueError(data.velocity_representation)
# =====================================================
# Compute quantities to adjust the output representation
# =====================================================
with data.switch_velocity_representation(VelRepr.Inertial):
# Compute the Jacobian of the parent link in inertial representation.
W_J_WL_W = js.model.generalized_free_floating_jacobian(
model=model,
data=data,
)
# Compute the Jacobian derivative of the parent link in inertial representation.
W_J̇_WL_W = js.model.generalized_free_floating_jacobian_derivative(
model=model,
data=data,
)
def compute_O_J̇_WC_I(
L_p_C: jtp.Vector,
parent_link_idx: jtp.Int,
W_H_L: jtp.Matrix,
) -> jtp.Matrix:
match output_vel_repr:
case VelRepr.Inertial:
O_X_W = W_X_W = Adjoint.from_transform( # noqa: F841
transform=jnp.eye(4)
)
O_Ẋ_W = W_Ẋ_W = jnp.zeros((6, 6)) # noqa: F841
case VelRepr.Body:
L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
W_H_C = W_H_L[parent_link_idx] @ L_H_C
O_X_W = C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
W_v_WC = W_v_WLi[parent_link_idx]
W_vx_WC = Cross.vx(W_v_WC)
O_Ẋ_W = C_Ẋ_W = -C_X_W @ W_vx_WC # noqa: F841
case VelRepr.Mixed:
L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
W_H_C = W_H_L[parent_link_idx] @ L_H_C
W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
CW_H_W = Transform.inverse(W_H_CW)
O_X_W = CW_X_W = Adjoint.from_transform(transform=CW_H_W)
CW_v_WC = CW_X_W @ W_v_WLi[parent_link_idx]
W_v_W_CW = jnp.zeros(6).at[0:3].set(CW_v_WC[0:3])
W_vx_W_CW = Cross.vx(W_v_W_CW)
O_Ẋ_W = CW_Ẋ_W = -CW_X_W @ W_vx_W_CW # noqa: F841
case _:
raise ValueError(output_vel_repr)
O_J̇_WC_I = jnp.zeros(shape=(6, 6 + model.dofs()))
O_J̇_WC_I += O_Ẋ_W @ W_J_WL_W[parent_link_idx] @ T
O_J̇_WC_I += O_X_W @ W_J̇_WL_W[parent_link_idx] @ T
O_J̇_WC_I += O_X_W @ W_J_WL_W[parent_link_idx] @ Ṫ
return O_J̇_WC_I
O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, None))(
L_p_Ci, parent_link_idx_of_enabled_collidable_points, W_H_Li
)
return O_J̇_WC
@jax.jit
@js.common.named_scope
def link_contact_forces(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
link_forces: jtp.MatrixLike | None = None,
joint_torques: jtp.VectorLike | None = None,
) -> tuple[jtp.Matrix, dict[str, jtp.Matrix]]:
"""
Compute the 6D contact forces of all links of the model in inertial representation.
Args:
model: The model to consider.
data: The data of the considered model.
link_forces:
The 6D external forces to apply to the links expressed in inertial representation
joint_torques:
The joint torques acting on the joints.
Returns:
A `(nL, 6)` array containing the stacked 6D contact forces of the links,
expressed in inertial representation.
"""
# Compute the contact forces for each collidable point with the active contact model.
W_f_C, aux_dict = model.contact_model.compute_contact_forces(
model=model,
data=data,
**(
dict(link_forces=link_forces, joint_force_references=joint_torques)
if not isinstance(model.contact_model, SoftContacts)
else {}
),
)
# Compute the 6D forces applied to the links equivalent to the forces applied
# to the frames associated to the collidable points.
W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
return W_f_L, aux_dict
def link_forces_from_contact_forces(
model: js.model.JaxSimModel,
*,
contact_forces: jtp.MatrixLike,
) -> jtp.Matrix:
"""
Compute the link forces from the contact forces.
Args:
model: The robot model considered by the contact model.
contact_forces: The contact forces computed by the contact model.
Returns:
The 6D contact forces applied to the links and expressed in the frame of
the velocity representation of data.
"""
# Get the object storing the contact parameters of the model.
contact_parameters = model.kin_dyn_parameters.contact_parameters
# Extract the indices corresponding to the enabled collidable points.
indices_of_enabled_collidable_points = (
contact_parameters.indices_of_enabled_collidable_points
)
# Convert the contact forces to a JAX array.
W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
# Construct the vector defining the parent link index of each collidable point.
# We use this vector to sum the 6D forces of all collidable points rigidly
# attached to the same link.
parent_link_index_of_collidable_points = jnp.array(
contact_parameters.body, dtype=int
)[indices_of_enabled_collidable_points]
# Create the mask that associate each collidable point to their parent link.
# We use this mask to sum the collidable points to the right link.
mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
model.number_of_links()
)
# Sum the forces of all collidable points rigidly attached to a body.
# Since the contact forces W_f_C are expressed in the world frame,
# we don't need any coordinate transformation.
W_f_L = mask.T @ W_f_C
return W_f_L
================================================
FILE: src/jaxsim/api/data.py
================================================
from __future__ import annotations
import dataclasses
import functools
from collections.abc import Sequence
try:
from typing import Self, override
except ImportError:
from typing_extensions import override, Self
import jax
import jax.numpy as jnp
import jax.scipy.spatial.transform
import jax_dataclasses
import jaxsim.api as js
import jaxsim.math
import jaxsim.rbda
import jaxsim.typing as jtp
from . import common
from .common import VelRepr
@jax_dataclasses.pytree_dataclass
class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
"""
Class storing the state of the physics model dynamics.
Attributes:
joint_positions: The vector of joint positions.
joint_velocities: The vector of joint velocities.
base_position: The 3D position of the base link.
base_quaternion: The quaternion defining the orientation of the base link.
base_linear_velocity:
The linear velocity of the base link in inertial-fixed representation.
base_angular_velocity:
The angular velocity of the base link in inertial-fixed representation.
base_transform: The base transform.
joint_transforms: The joint transforms.
link_transforms: The link transforms.
link_velocities: The link velocities in inertial-fixed representation.
"""
# Joint state
_joint_positions: jtp.Vector
_joint_velocities: jtp.Vector
# Base state
_base_quaternion: jtp.Vector
_base_linear_velocity: jtp.Vector
_base_angular_velocity: jtp.Vector
_base_position: jtp.Vector
# Cached computations.
_base_transform: jtp.Matrix = dataclasses.field(repr=False, default=None)
_joint_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
_link_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
_link_velocities: jtp.Matrix = dataclasses.field(repr=False, default=None)
# Extended state for soft and rigid contact models.
contact_state: dict[str, jtp.Array] = dataclasses.field(default_factory=dict)
@staticmethod
def build(
model: js.model.JaxSimModel,
base_position: jtp.VectorLike | None = None,
base_quaternion: jtp.VectorLike | None = None,
joint_positions: jtp.VectorLike | None = None,
base_linear_velocity: jtp.VectorLike | None = None,
base_angular_velocity: jtp.VectorLike | None = None,
joint_velocities: jtp.VectorLike | None = None,
contact_state: dict[str, jtp.Array] | None = None,
velocity_representation: VelRepr = VelRepr.Mixed,
) -> JaxSimModelData:
"""
Create a `JaxSimModelData` object with the given state.
Args:
model: The model for which to create the state.
base_position: The base position.
base_quaternion: The base orientation as a quaternion.
joint_positions: The joint positions.
base_linear_velocity:
The base linear velocity in the selected representation.
base_angular_velocity:
The base angular velocity in the selected representation.
joint_velocities: The joint velocities.
velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
contact_state: The optional contact state.
Returns:
A `JaxSimModelData` initialized with the given state.
"""
base_position = jnp.array(
base_position if base_position is not None else jnp.zeros(3),
dtype=float,
).squeeze()
base_quaternion = jnp.array(
(
base_quaternion
if base_quaternion is not None
else jnp.array([1.0, 0, 0, 0])
),
dtype=float,
).squeeze()
base_linear_velocity = jnp.array(
base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3),
dtype=float,
).squeeze()
base_angular_velocity = jnp.array(
(
base_angular_velocity
if base_angular_velocity is not None
else jnp.zeros(3)
),
dtype=float,
).squeeze()
joint_positions = jnp.atleast_1d(
jnp.array(
(
joint_positions
if joint_positions is not None
else jnp.zeros(model.dofs())
),
dtype=float,
).squeeze()
)
joint_velocities = jnp.atleast_1d(
jnp.array(
(
joint_velocities
if joint_velocities is not None
else jnp.zeros(model.dofs())
),
dtype=float,
).squeeze()
)
W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
translation=base_position, quaternion=base_quaternion
)
W_v_WB = JaxSimModelData.other_representation_to_inertial(
array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
other_representation=velocity_representation,
transform=W_H_B,
is_force=False,
).astype(float)
joint_transforms = model.kin_dyn_parameters.joint_transforms(
joint_positions=joint_positions, base_transform=W_H_B
)
link_transforms, link_velocities_inertial = (
jaxsim.rbda.forward_kinematics_model(
model=model,
base_position=base_position,
base_quaternion=base_quaternion,
joint_positions=joint_positions,
base_linear_velocity_inertial=W_v_WB[0:3],
base_angular_velocity_inertial=W_v_WB[3:6],
joint_velocities=joint_velocities,
joint_transforms=joint_transforms,
)
)
contact_state = contact_state or {}
if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
contact_state["tangential_deformation"] = contact_state.get(
"tangential_deformation",
jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
)
model_data = JaxSimModelData(
velocity_representation=velocity_representation,
_base_quaternion=base_quaternion,
_base_position=base_position,
_joint_positions=joint_positions,
_base_linear_velocity=W_v_WB[0:3],
_base_angular_velocity=W_v_WB[3:6],
_joint_velocities=joint_velocities,
_base_transform=W_H_B,
_joint_transforms=joint_transforms,
_link_transforms=link_transforms,
_link_velocities=link_velocities_inertial,
contact_state=contact_state,
)
if not model_data.valid(model=model):
raise ValueError(
"The built state is not compatible with the model.", model_data
)
return model_data
@staticmethod
def zero(
model: js.model.JaxSimModel,
velocity_representation: VelRepr = VelRepr.Mixed,
) -> JaxSimModelData:
"""
Create a `JaxSimModelData` object with zero state.
Args:
model: The model for which to create the state.
velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
Returns:
A `JaxSimModelData` initialized with zero state.
"""
return JaxSimModelData.build(
model=model, velocity_representation=velocity_representation
)
# ==================
# Extract quantities
# ==================
@property
def joint_positions(self) -> jtp.Vector:
"""
Get the joint positions.
Returns:
The joint positions.
"""
return self._joint_positions
@property
def joint_velocities(self) -> jtp.Vector:
"""
Get the joint velocities.
Returns:
The joint velocities.
"""
return self._joint_velocities
@property
def base_quaternion(self) -> jtp.Vector:
"""
Get the base quaternion.
Returns:
The base quaternion.
"""
return self._base_quaternion
@property
def base_position(self) -> jtp.Vector:
"""
Get the base position.
Returns:
The base position.
"""
return self._base_position
@property
def base_orientation(self) -> jtp.Matrix:
"""
Get the base orientation.
Returns:
The base orientation.
"""
# Extract the base quaternion.
W_Q_B = self.base_quaternion
# Always normalize the quaternion to avoid numerical issues.
# If the active scheme does not integrate the quaternion on its manifold,
# we introduce a Baumgarte stabilization to let the quaternion converge to
# a unit quaternion. In this case, it is not guaranteed that the quaternion
# stored in the state is a unit quaternion.
norm = jaxsim.math.safe_norm(W_Q_B, axis=-1, keepdims=True)
W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
return W_Q_B
@property
def base_velocity(self) -> jtp.Vector:
"""
Get the base 6D velocity.
Returns:
The base 6D velocity in the active representation.
"""
W_v_WB = jnp.concatenate(
[self._base_linear_velocity, self._base_angular_velocity], axis=-1
)
W_H_B = self._base_transform
return (
JaxSimModelData.inertial_to_other_representation(
array=W_v_WB,
other_representation=self.velocity_representation,
transform=W_H_B,
is_force=False,
)
.squeeze()
.astype(float)
)
@property
def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
r"""
Get the generalized position
:math:`\mathbf{q} = ({}^W \mathbf{H}_B, \mathbf{s}) \in \text{SO}(3) \times \mathbb{R}^n`.
Returns:
A tuple containing the base transform and the joint positions.
"""
return self._base_transform, self.joint_positions
@property
def generalized_velocity(self) -> jtp.Vector:
r"""
Get the generalized velocity.
:math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\, \boldsymbol{\omega}_{W,B};\, \mathbf{s}) \in \mathbb{R}^{6+n}`
Returns:
The generalized velocity in the active representation.
"""
return (
jnp.hstack([self.base_velocity, self.joint_velocities])
.squeeze()
.astype(float)
)
@property
def base_transform(self) -> jtp.Matrix:
"""
Get the base transform.
Returns:
The base transform.
"""
return self._base_transform
# ================
# Store quantities
# ================
@js.common.named_scope
@jax.jit
def reset_base_quaternion(
self, model: js.model.JaxSimModel, base_quaternion: jtp.VectorLike
) -> Self:
"""
Reset the base quaternion.
Args:
model: The JaxSim model to use.
base_quaternion: The base orientation as a quaternion.
Returns:
The updated `JaxSimModelData` object.
"""
W_Q_B = jnp.array(base_quaternion, dtype=float)
norm = jaxsim.math.safe_norm(W_Q_B, axis=-1)
W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
return self.replace(model=model, base_quaternion=W_Q_B)
@js.common.named_scope
@jax.jit
def reset_base_pose(
self, model: js.model.JaxSimModel, base_pose: jtp.MatrixLike
) -> Self:
"""
Reset the base pose.
Args:
model: The JaxSim model to use.
base_pose: The base pose as an SE(3) matrix.
Returns:
The updated `JaxSimModelData` object.
"""
base_pose = jnp.array(base_pose)
W_p_B = base_pose[0:3, 3]
W_Q_B = jaxsim.math.Quaternion.from_dcm(dcm=base_pose[0:3, 0:3])
return self.replace(
model=model,
base_position=W_p_B,
base_quaternion=W_Q_B,
)
@override
def replace(
self,
model: js.model.JaxSimModel,
joint_positions: jtp.Vector | None = None,
joint_velocities: jtp.Vector | None = None,
base_quaternion: jtp.Vector | None = None,
base_linear_velocity: jtp.Vector | None = None,
base_angular_velocity: jtp.Vector | None = None,
base_position: jtp.Vector | None = None,
*,
contact_state: dict[str, jtp.Array] | None = None,
validate: bool = False,
) -> Self:
"""
Replace the attributes of the `JaxSimModelData` object.
"""
if joint_positions is None:
joint_positions = self.joint_positions
if joint_velocities is None:
joint_velocities = self.joint_velocities
if base_quaternion is None:
base_quaternion = self.base_quaternion
if base_position is None:
base_position = self.base_position
if contact_state is None:
contact_state = self.contact_state
# Normalize the quaternion to avoid numerical issues.
base_quaternion_norm = jaxsim.math.safe_norm(
base_quaternion, axis=-1, keepdims=True
)
base_quaternion = base_quaternion / jnp.where(
base_quaternion_norm == 0, 1.0, base_quaternion_norm
)
joint_positions = jnp.atleast_2d(joint_positions.squeeze()).astype(float)
joint_velocities = jnp.atleast_2d(joint_velocities.squeeze()).astype(float)
base_quaternion = jnp.atleast_2d(base_quaternion.squeeze()).astype(float)
base_position = jnp.atleast_2d(base_position.squeeze()).astype(float)
base_transform = jaxsim.math.Transform.from_quaternion_and_translation(
translation=base_position, quaternion=base_quaternion
).reshape((-1, 4, 4))
joint_transforms = jax.vmap(model.kin_dyn_parameters.joint_transforms)(
joint_positions=joint_positions,
base_transform=base_transform,
)
if base_linear_velocity is None and base_angular_velocity is None:
base_linear_velocity_inertial = self._base_linear_velocity
base_angular_velocity_inertial = self._base_angular_velocity
else:
if base_linear_velocity is None:
base_linear_velocity = self.base_velocity[:3]
if base_angular_velocity is None:
base_angular_velocity = self.base_velocity[3:]
base_linear_velocity = jnp.atleast_1d(base_linear_velocity.squeeze())
base_angular_velocity = jnp.atleast_1d(base_angular_velocity.squeeze())
W_v_WB = JaxSimModelData.other_representation_to_inertial(
array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
other_representation=self.velocity_representation,
transform=base_transform,
is_force=False,
).astype(float)
base_linear_velocity_inertial, base_angular_velocity_inertial = (
W_v_WB[..., :3],
W_v_WB[..., 3:],
)
link_transforms, link_velocities = jax.vmap(
jaxsim.rbda.forward_kinematics_model, in_axes=(None,)
)(
model,
base_position=base_position,
base_quaternion=base_quaternion,
joint_positions=joint_positions,
joint_velocities=joint_velocities,
base_linear_velocity_inertial=jnp.atleast_2d(base_linear_velocity_inertial),
base_angular_velocity_inertial=jnp.atleast_2d(
base_angular_velocity_inertial
),
joint_transforms=joint_transforms,
)
# Adjust the output shapes.
joint_positions = joint_positions.reshape(self._joint_positions.shape)
joint_velocities = joint_velocities.reshape(self._joint_velocities.shape)
base_quaternion = base_quaternion.reshape(self._base_quaternion.shape)
base_linear_velocity_inertial = base_linear_velocity_inertial.reshape(
self._base_linear_velocity.shape
)
base_angular_velocity_inertial = base_angular_velocity_inertial.reshape(
self._base_angular_velocity.shape
)
base_position = base_position.reshape(self._base_position.shape)
base_transform = base_transform.reshape(self._base_transform.shape)
joint_transforms = joint_transforms.reshape(self._joint_transforms.shape)
link_transforms = link_transforms.reshape(self._link_transforms.shape)
link_velocities = link_velocities.reshape(self._link_velocities.shape)
return super().replace(
_joint_positions=joint_positions,
_joint_velocities=joint_velocities,
_base_quaternion=base_quaternion,
_base_linear_velocity=base_linear_velocity_inertial,
_base_angular_velocity=base_angular_velocity_inertial,
_base_position=base_position,
_base_transform=base_transform,
_joint_transforms=joint_transforms,
_link_transforms=link_transforms,
_link_velocities=link_velocities,
contact_state=contact_state,
validate=validate,
)
def valid(self, model: js.model.JaxSimModel) -> bool:
"""
Check if the `JaxSimModelData` is valid for a given `JaxSimModel`.
Args:
model: The `JaxSimModel` to validate the `JaxSimModelData` against.
Returns:
`True` if the `JaxSimModelData` is valid for the given model,
`False` otherwise.
"""
if self._joint_positions.shape != (model.dofs(),):
return False
if self._joint_velocities.shape != (model.dofs(),):
return False
if self._base_position.shape != (3,):
return False
if self._base_quaternion.shape != (4,):
return False
if self._base_linear_velocity.shape != (3,):
return False
if self._base_angular_velocity.shape != (3,):
return False
return True
@functools.partial(jax.jit, static_argnames=["velocity_representation", "base_rpy_seq"])
def random_model_data(
model: js.model.JaxSimModel,
*,
key: jax.Array | None = None,
velocity_representation: VelRepr | None = None,
base_pos_bounds: tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
] = ((-1, -1, 0.5), 1.0),
base_rpy_bounds: tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
] = (-jnp.pi, jnp.pi),
base_rpy_seq: str = "XYZ",
joint_pos_bounds: (
tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
]
| None
) = None,
base_vel_lin_bounds: tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
] = (-1.0, 1.0),
base_vel_ang_bounds: tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
] = (-1.0, 1.0),
joint_vel_bounds: tuple[
jtp.FloatLike | Sequence[jtp.FloatLike],
jtp.FloatLike | Sequence[jtp.FloatLike],
] = (-1.0, 1.0),
) -> JaxSimModelData:
"""
Randomly generate a `JaxSimModelData` object.
Args:
model: The target model for the random data.
key: The random key.
velocity_representation: The velocity representation to use.
base_pos_bounds: The bounds for the base position.
base_rpy_bounds:
The bounds for the euler angles used to build the base orientation.
base_rpy_seq:
The sequence of axes for rotation (using `Rotation` from scipy).
joint_pos_bounds:
The bounds for the joint positions (reading the joint limits if None).
base_vel_lin_bounds: The bounds for the base linear velocity.
base_vel_ang_bounds: The bounds for the base angular velocity.
joint_vel_bounds: The bounds for the joint velocities.
Returns:
A `JaxSimModelData` object with random data.
"""
key = key if key is not None else jax.random.PRNGKey(seed=0)
k1, k2, k3, k4, k5, k6 = jax.random.split(key, num=6)
p_min = jnp.array(base_pos_bounds[0], dtype=float)
p_max = jnp.array(base_pos_bounds[1], dtype=float)
rpy_min = jnp.array(base_rpy_bounds[0], dtype=float)
rpy_max = jnp.array(base_rpy_bounds[1], dtype=float)
v_min = jnp.array(base_vel_lin_bounds[0], dtype=float)
v_max = jnp.array(base_vel_lin_bounds[1], dtype=float)
ω_min = jnp.array(base_vel_ang_bounds[0], dtype=float)
ω_max = jnp.array(base_vel_ang_bounds[1], dtype=float)
ṡ_min, ṡ_max = joint_vel_bounds
base_position = jax.random.uniform(key=k1, shape=(3,), minval=p_min, maxval=p_max)
base_quaternion = jaxsim.math.Quaternion.to_wxyz(
xyzw=jax.scipy.spatial.transform.Rotation.from_euler(
seq=base_rpy_seq,
angles=jax.random.uniform(
key=k2, shape=(3,), minval=rpy_min, maxval=rpy_max
),
).as_quat()
)
(
joint_positions,
joint_velocities,
base_linear_velocity,
base_angular_velocity,
) = (None,) * 4
if model.number_of_joints() > 0:
s_min, s_max = (
jnp.array(joint_pos_bounds, dtype=float)
if joint_pos_bounds is not None
else (None, None)
)
joint_positions = (
js.joint.random_joint_positions(model=model, key=k3)
if (s_min is None or s_max is None)
else jax.random.uniform(
key=k3, shape=(model.dofs(),), minval=s_min, maxval=s_max
)
)
joint_velocities = jax.random.uniform(
key=k4, shape=(model.dofs(),), minval=ṡ_min, maxval=ṡ_max
)
if model.floating_base():
base_linear_velocity = jax.random.uniform(
key=k5, shape=(3,), minval=v_min, maxval=v_max
)
base_angular_velocity = jax.random.uniform(
key=k6, shape=(3,), minval=ω_min, maxval=ω_max
)
return JaxSimModelData.build(
model=model,
base_position=base_position,
base_quaternion=base_quaternion,
joint_positions=joint_positions,
joint_velocities=joint_velocities,
base_linear_velocity=base_linear_velocity,
base_angular_velocity=base_angular_velocity,
**(
{"velocity_representation": velocity_representation}
if velocity_representation is not None
else {}
),
)
================================================
FILE: src/jaxsim/api/frame.py
================================================
import functools
from collections.abc import Sequence
import jax
import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.typing as jtp
from jaxsim import exceptions
from jaxsim.math import Adjoint, Cross
from .common import VelRepr
# =======================
# Index-related functions
# =======================
@jax.jit
@js.common.named_scope
def idx_of_parent_link(
model: js.model.JaxSimModel, *, frame_index: jtp.IntLike
) -> jtp.Int:
"""
Get the index of the link to which the frame is rigidly attached.
Args:
model: The model to consider.
frame_index: The index of the frame.
Returns:
The index of the frame's parent link.
"""
n_l = model.number_of_links()
n_f = len(model.frame_names())
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
return jnp.array(model.kin_dyn_parameters.frame_parameters.body)[
frame_index - model.number_of_links()
]
@functools.partial(jax.jit, static_argnames="frame_name")
@js.common.named_scope
def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
"""
Convert the name of a frame to its index.
Args:
model: The model to consider.
frame_name: The name of the frame.
Returns:
The index of the frame.
"""
if frame_name not in model.frame_names():
raise ValueError(f"Frame '{frame_name}' not found in the model.")
return (
jnp.array(
model.number_of_links()
+ model.kin_dyn_parameters.frame_parameters.name.index(frame_name)
)
.astype(int)
.squeeze()
)
def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
"""
Convert the index of a frame to its name.
Args:
model: The model to consider.
frame_index: The index of the frame.
Returns:
The name of the frame.
"""
n_l = model.number_of_links()
n_f = len(model.frame_names())
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
return model.kin_dyn_parameters.frame_parameters.name[
frame_index - model.number_of_links()
]
@functools.partial(jax.jit, static_argnames=["frame_names"])
@js.common.named_scope
def names_to_idxs(
model: js.model.JaxSimModel, *, frame_names: Sequence[str]
) -> jax.Array:
"""
Convert a sequence of frame names to their corresponding indices.
Args:
model: The model to consider.
frame_names: The names of the frames.
Returns:
The indices of the frames.
"""
return jnp.array(
[name_to_idx(model=model, frame_name=name) for name in frame_names]
).astype(int)
def idxs_to_names(
model: js.model.JaxSimModel, *, frame_indices: Sequence[jtp.IntLike]
) -> tuple[str, ...]:
"""
Convert a sequence of frame indices to their corresponding names.
Args:
model: The model to consider.
frame_indices: The indices of the frames.
Returns:
The names of the frames.
"""
return tuple(idx_to_name(model=model, frame_index=idx) for idx in frame_indices)
# ==========
# Frame APIs
# ==========
@jax.jit
@js.common.named_scope
def transform(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
frame_index: jtp.IntLike,
) -> jtp.Matrix:
"""
Compute the SE(3) transform from the world frame to the specified frame.
Args:
model: The model to consider.
data: The data of the considered model.
frame_index: The index of the frame for which the transform is requested.
Returns:
The 4x4 matrix representing the transform.
"""
n_l = model.number_of_links()
n_f = len(model.frame_names())
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
# Compute the necessary transforms.
L = idx_of_parent_link(model=model, frame_index=frame_index)
W_H_L = js.link.transform(model=model, data=data, link_index=L)
# Get the static frame pose wrt the parent link.
L_H_F = model.kin_dyn_parameters.frame_parameters.transform[
frame_index - model.number_of_links()
]
# Combine the transforms computing the frame pose.
return W_H_L @ L_H_F
@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
@js.common.named_scope
def velocity(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
frame_index: jtp.IntLike,
output_vel_repr: VelRepr | None = None,
) -> jtp.Vector:
"""
Compute the 6D velocity of the frame.
Args:
model: The model to consider.
data: The data of the considered model.
frame_index: The index of the frame.
output_vel_repr:
The output velocity representation of the frame velocity.
Returns:
The 6D velocity of the frame in the specified velocity representation.
"""
n_l = model.number_of_links()
n_f = model.number_of_frames()
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
output_vel_repr = (
output_vel_repr if output_vel_repr is not None else data.velocity_representation
)
# Get the frame jacobian having I as input representation (taken from data)
# and O as output representation, specified by the user (or taken from data).
O_J_WF_I = jacobian(
model=model,
data=data,
frame_index=frame_index,
output_vel_repr=output_vel_repr,
)
# Get the generalized velocity in the input velocity representation.
I_ν = data.generalized_velocity
# Compute the frame velocity in the output velocity representation.
return O_J_WF_I @ I_ν
@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
@js.common.named_scope
def jacobian(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
frame_index: jtp.IntLike,
output_vel_repr: VelRepr | None = None,
) -> jtp.Matrix:
r"""
Compute the free-floating jacobian of the frame.
Args:
model: The model to consider.
data: The data of the considered model.
frame_index: The index of the frame.
output_vel_repr:
The output velocity representation of the free-floating jacobian.
Returns:
The :math:`6 \times (6+n)` free-floating jacobian of the frame.
Note:
The input representation of the free-floating jacobian is the active
velocity representation.
"""
n_l = model.number_of_links()
n_f = model.number_of_frames()
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
output_vel_repr = (
output_vel_repr if output_vel_repr is not None else data.velocity_representation
)
# Get the index of the parent link.
L = idx_of_parent_link(model=model, frame_index=frame_index)
# Compute only the parent-link body Jacobian.
L_J_WL = js.link.jacobian(
model=model,
data=data,
link_index=L,
output_vel_repr=VelRepr.Body,
)
W_H_L = data._link_transforms[L]
L_H_F = model.kin_dyn_parameters.frame_parameters.transform[
frame_index - model.number_of_links()
]
L_p_F = L_H_F[0:3, 3]
# Adjust the output representation.
match output_vel_repr:
case VelRepr.Inertial:
W_X_L = Adjoint.from_rotation_and_translation(
rotation=W_H_L[0:3, 0:3],
translation=W_H_L[0:3, 3],
)
W_J_WL = W_X_L @ L_J_WL
O_J_WL_I = W_J_WL
case VelRepr.Body:
F_X_L = Adjoint.from_rotation_and_translation(
rotation=L_H_F[0:3, 0:3],
translation=L_p_F,
inverse=True,
)
F_J_WL = F_X_L @ L_J_WL
O_J_WL_I = F_J_WL
case VelRepr.Mixed:
W_R_L = W_H_L[0:3, 0:3]
FW_X_L = Adjoint.from_rotation_and_translation(
rotation=W_R_L,
translation=-W_R_L @ L_p_F,
)
FW_J_WL = FW_X_L @ L_J_WL
O_J_WL_I = FW_J_WL
case _:
raise ValueError(output_vel_repr)
return O_J_WL_I
@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
@js.common.named_scope
def jacobian_derivative(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
*,
frame_index: jtp.IntLike,
output_vel_repr: VelRepr | None = None,
) -> jtp.Matrix:
r"""
Compute the derivative of the free-floating jacobian of the frame.
Args:
model: The model to consider.
data: The data of the considered model.
frame_index: The index of the frame.
output_vel_repr:
The output velocity representation of the free-floating jacobian derivative.
Returns:
The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the frame.
Note:
The input representation of the free-floating jacobian derivative is the active
velocity representation.
"""
n_l = model.number_of_links()
n_f = len(model.frame_names())
exceptions.raise_value_error_if(
condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
msg="Invalid frame index '{idx}'",
idx=frame_index,
)
output_vel_repr = (
output_vel_repr if output_vel_repr is not None else data.velocity_representation
)
# Get the index of the parent link.
L = idx_of_parent_link(model=model, frame_index=frame_index)
W_J_WL_I = js.link.jacobian(
model=model,
data=data,
link_index=L,
output_vel_repr=VelRepr.Inertial,
)
W_J̇_WL_I = js.link.jacobian_derivative(
model=model,
data=data,
link_index=L,
output_vel_repr=VelRepr.Inertial,
)
W_H_L = data._link_transforms[L]
L_H_F = model.kin_dyn_parameters.frame_parameters.transform[
frame_index - model.number_of_links()
]
W_H_F = W_H_L @ L_H_F
# =====================================================
# Compute quantities to adjust the output representation
# =====================================================
W_v_WF = W_J_WL_I @ data.generalized_velocity
match output_vel_repr:
case VelRepr.Inertial:
O_X_W = jnp.eye(6, dtype=W_H_F.dtype)
O_Ẋ_W = jnp.zeros((6, 6), dtype=W_H_F.dtype)
case VelRepr.Body:
O_X_W = Adjoint.from_rotation_and_translation(
rotation=W_H_F[0:3, 0:3],
translation=W_H_F[0:3, 3],
inverse=True,
)
O_Ẋ_W = -O_X_W @ Cross.vx(W_v_WF)
case VelRepr.Mixed:
O_X_W = Adjoint.from_rotation_and_translation(
rotation=jnp.eye(3, dtype=W_H_F.dtype),
translation=W_H_F[0:3, 3],
inverse=True,
)
FW_v_WF = O_X_W @ W_v_WF
W_v_W_FW = FW_v_WF.at[3:6].set(jnp.zeros_like(FW_v_WF[3:6]))
O_Ẋ_W = -O_X_W @ Cross.vx(W_v_W_FW)
case _:
raise ValueError(output_vel_repr)
O_J̇_WF_I = O_Ẋ_W @ W_J_WL_I
O_J̇_WF_I += O_X_W @ W_J̇_WL_I
return O_J̇_WF_I
================================================
FILE: src/jaxsim/api/integrators.py
================================================
import dataclasses
from collections.abc import Callable
import jax
import jax.numpy as jnp
import jaxsim
import jaxsim.api as js
import jaxsim.typing as jtp
from jaxsim.api.data import JaxSimModelData
from jaxsim.math import Skew
def semi_implicit_euler_integration(
model: js.model.JaxSimModel,
data: js.data.JaxSimModelData,
link_forces: jtp.Vector,
joint_torques: jtp.Vector,
) -> JaxSimModelData:
"""Integrate the system state using the semi-implicit Euler method."""
with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
# Compute the system acceleration
W_v̇_WB, s̈, contact_state_derivative = js.ode.system_acceleration(
model=model,
data=data,
link_forces=link_forces,
joint_torques=joint_torques,
)
dt = model.time_step
# Compute the new generalized velocity.
new_generalized_acceleration = jnp.hstack([W_v̇_WB, s̈])
new_generalized_velocity = (
data.generalized_velocity + dt * new_generalized_acceleration
)
# Extract the new base and joint velocities.
W_v_B = new_generalized_velocity[0:6]
ṡ = new_generalized_velocity[6:]
# Compute the new base position and orientation.
W_ω_WB = new_generalized_velocity[3:6]
# To obtain the derivative of the base position, we need to subtract
# the skew-symmetric matrix of the base angular velocity times the base position.
# See: S. Traversaro and A. Saccon, “Multibody Dynamics Notation (Version 2), pg.9
W_ṗ_B = new_generalized_velocity[0:3] + Skew.wedge(W_ω_WB) @ data.base_position
W_Q̇_B = jaxsim.math.Quaternion.derivative(
quaternion=data.base_orientation,
omega=W_ω_WB,
omega_in_body_fixed=False,
).squeeze()
W_p_B = data.base_position + dt * W_ṗ_B
W_Q_B = data.base_orientation + dt * W_Q̇_B
base_quaternion_norm = jaxsim.math.safe_norm(W_Q_B, axis=-1)
W_Q_B = W_Q_B / jnp.where(base_quaternion_norm == 0, 1.0, base_quaternion_norm)
s = data.joint_positions + dt * ṡ
integrated_contact_state = jax.tree.map(
lambda x, x_dot: x + dt * x_dot,
data.contact_state,
contact_state_derivative,
)
data = dataclasses.replace(
data,
_base_quaternion=W_Q_B,
_base_position=W_p_B,
_joint_positions=s,
_joint_velocities=ṡ,
_base_linear_velocity=W_v_B[0:3],
_base_angular_velocity=W_ω_WB,
contact_state=integrated_contact_state,
)
# Recompute kinematic caches for the new state.
data = data.replace(model=model)
return data
def rk4_integration(
model: js.model.JaxSimModel,
data: JaxSimModelData,
link_forces: jtp.Vector,
joint_torques: jtp.Vector,
) -> JaxSimModelData:
"""Integrate the system state using the Runge-Kutta 4 method."""
dt = model.time_step
def f(x) -> dict[str, jtp.Matrix]:
with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
data_ti = data.replace(model=model, **x)
return js.ode.system_dynamics(
model=model,
data=data_ti,
link_forces=link_forces,
joint_torques=joint_torques,
)
base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion, axis=-1)
base_quaternion = data._base_quaternion / jnp.where(
base_quaternion_norm == 0, 1.0, base_quaternion_norm
)
x_t0 = dict(
base_position=data._base_position,
base_quaternion=base_quaternion,
joint_positions=data._joint_positions,
base_linear_velocity=data._base_linear_velocity,
base_angular_velocity=data._base_angular_velocity,
joint_velocities=data._joint_velocities,
contact_state=data.contact_state,
)
euler_mid = lambda x, dxdt: x + (0.5 * dt) * dxdt
euler_fin = lambda x, dxdt: x + dt * dxdt
k1 = f(x_t0)
k2 = f(jax.tree.map(euler_mid, x_t0, k1))
k3 = f(jax.tree.map(euler_mid, x_t0, k2))
k4 = f(jax.tree.map(euler_fin, x_t0, k3))
# Average the slopes and compute the RK4 state derivative.
average = lambda k1, k2, k3, k4: (k1 + 2 * k2 + 2 * k3 + k4) / 6
dxdt = jax.tree.map(average, k1, k2, k3, k4)
# Integrate the dynamics
x_tf = jax.tree.map(euler_fin, x_t0, dxdt)
data_tf = dataclasses.replace(
data,
_base_position=x_tf["base_position"],
_base_quaternion=x_tf["base_quaternion"],
_joint_positions=x_tf["joint_positions"],
_base_linear_velocity=x_tf["base_linear_velocity"],
_base_angular_velocity=x_tf["base_angular_velocity"],
_joint_velocities=x_tf["joint_velocities"],
contact_state=x_tf["contact_state"],
)
return data_tf.replace(model=model)
def rk4fast_integration(
model: js.model.JaxSimModel,
data: JaxSimModelData,
link_forces: jtp.Vector,
joint_torques: jtp.Vector,
) -> JaxSimModelData:
"""
Integrate the system state using the Runge-Kutta 4 fast method.
Note:
This method is a faster version of the RK4 method, but it may not be as accurate.
It computes the contact forces only once at the beginning of the integration step.
"""
dt = model.time_step
if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
# Compute the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
# with the terrain.
W_f_L_terrain, contact_state_derivative = js.contact.link_contact_forces(
model=model,
data=data,
link_forces=link_forces,
joint_torques=joint_torques,
)
W_f_L_total = link_forces + W_f_L_terrain
# Update the contact state data. This is necessary only for the contact models
# that require propagation and integration of contact state.
contact_state = model.contact_model.update_contact_state(contact_state_derivative)
def f(x) -> dict[str, jtp.Matrix]:
with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
data_ti = data.replace(model=model, **x)
W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
model=model,
data=data_ti,
joint_forces=joint_torques,
link_forces=W_f_L_total,
)
W_ṗ_B, W_Q̇_B, ṡ = js.ode.system_position_dynamics(
data=data,
baumgarte_quaternion_regularization=1.0,
)
return dict(
base_position=W_ṗ_B,
base_quaternion=W_Q̇_B,
joint_positions=ṡ,
base_linear_velocity=W_v̇_WB[0:3],
base_angular_velocity=W_v̇_WB[3:6],
joint_velocities=s̈,
# The contact state is not updated here, as it is assumed to be constant.
contact_state=data_ti.contact_state,
)
base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion, axis=-1)
base_quaternion = data._base_quaternion / jnp.where(
base_quaternion_norm == 0, 1.0, base_quaternion_norm
)
x_t0 = dict(
base_position=data._base_position,
base_quaternion=base_quaternion,
joint_positions=data._joint_positions,
base_linear_velocity=data._base_linear_velocity,
base_angular_velocity=data._base_angular_velocity,
joint_velocities=data._joint_velocities,
contact_state=contact_state,
)
euler_mid = lambda x, dxdt: x + (0.5 * dt) * dxdt
euler_fin = lambda x, dxdt: x + dt * dxdt
k1 = f(x_t0)
k2 = f(jax.tree.map(euler_mid, x_t0, k1))
k3 = f(jax.tree.map(euler_mid, x_t0, k2))
k4 = f(jax.tree.map(euler_fin, x_t0, k3))
# Average the slopes and compute the RK4 state derivative.
average = lambda k1, k2, k3, k4: (k1 + 2 * k2 + 2 * k3 + k4) / 6
dxdt = jax.tree.map(average, k1, k2, k3, k4)
# Integrate the dynamics
x_tf = jax.tree.map(euler_fin, x_t0, dxdt)
data_tf = dataclasses.replace(
data,
_base_position=x_tf["base_position"],
_base_quaternion=x_tf["base_quaternion"],
_joint_positions=x_tf["joint_positions"],
_base_linear_velocity=x_tf["base_linear_velocity"],
_base_angular_velocity=x_tf["base_angular_velocity"],
_joint_velocities=x_tf["joint_velocities"],
contact_state=x_tf["contact_state"],
)
return data_tf.replace(model=model)
_INTEGRATORS_MAP: dict[
js.model.IntegratorType, Callable[..., js.data.JaxSimModelData]
] = {
js.model.IntegratorType.SemiImplicitEuler: semi_implicit_euler_integration,
js.model.IntegratorType.RungeKutta4: rk4_integration,
js.model.IntegratorType.RungeKutta4Fast: rk4fast_integration,
}
================================================
FILE: src/jaxsim/api/joint.py
================================================
import functools
from collections.abc import Sequence
import jax
import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.typing as jtp
from jaxsim import exceptions
# =======================
# Index-related functions
# =======================
@functools.partial(jax.jit, static_argnames="joint_name")
@js.common.named_scope
def name_to_idx(model: js.model.JaxSimModel, *, joint_name: str) -> jtp.Int:
"""
Convert the name of a joint to its index.
Args:
model: The model to consider.
joint_name: The name of the joint.
Returns:
The index of the joint.
"""
if joint_name not in model.joint_names():
raise ValueError(f"Joint '{joint_name}' not found in the model.")
# Note: the index of the joint for RBDAs starts from 1, but the index for
# accessing the right element starts from 0. Therefore, there is a -1.
return (
jnp.array(
model.kin_dyn_parameters.joint_model.joint_names.index(joint_name) - 1
)
.astype(int)
.squeeze()
)
def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str:
"""
Convert the index of a joint to its name.
Args:
model: The model to consider.
joint_index: The index of the joint.
Returns:
The name of the joint.
"""
exceptions.raise_value_error_if(
condition=joint_index < 0,
msg="Invalid joint index '{idx}'",
idx=joint_index,
)
return model.kin_dyn_parameters.joint_model.joint_names[joint_index + 1]
@functools.partial(jax.jit, static_argnames="joint_names")
@js.common.named_scope
def names_to_idxs(
model: js.model.JaxSimModel, *, joint_names: Sequence[str]
) -> jax.Array:
"""
Convert a sequence of joint names to their corresponding indices.
Args:
model: The model to consider.
joint_names: The names of the joints.
Returns:
The indices of the joints.
"""
return jnp.array(
[name_to_idx(model=model, joint_name=name) for name in joint_names],
).astype(int)
def idxs_to_names(
model: js.model.JaxSimModel,
*,
joint_indices: Sequence[jtp.IntLike] | jtp.VectorLike,
) -> tuple[str, ...]:
"""
Convert a sequence of joint indices to their corresponding names.
Args:
model: The model to consider.
joint_indices: The indices of the joints.
Returns:
The names of the joints.
"""
return tuple(idx_to_name(model=model, joint_index=idx) for idx in joint_indices)
# ============
# Joint limits
# ============
@jax.jit
def position_limit(
model: js.model.JaxSimModel, *, joint_index: jtp.IntLike
) -> tuple[jtp.Float, jtp.Float]:
"""
Get the position limits of a joint.
Args:
model: The model to consider.
joint_index: The index of the joint.
Returns:
The position limits of the joint.
"""
if model.number_of_joints() == 0:
return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
exceptions.raise_value_error_if(
condition=jnp.array(
[joint_index < 0, joint_index >= model.number_of_joints()]
).any(),
msg="Invalid joint index '{idx}'",
idx=joint_index,
)
s_min = jnp.atleast_1d(
model.kin_dyn_parameters.joint_parameters.position_limits_min
)[joint_index]
s_max = jnp.atleast_1d(
model.kin_dyn_parameters.joint_parameters.position_limits_max
)[joint_index]
return s_min.astype(float), s_max.astype(float)
@functools.partial(jax.jit, static_argnames=["joint_names"])
@js.common.named_scope
def position_limits(
model: js.model.JaxSimModel, *, joint_names: Sequence[str] | None = None
) -> tuple[jtp.Vector, jtp.Vector]:
"""
Get the position limits of a list of joint.
Args:
model: The model to consider.
joint_names: The names of the joints.
Returns:
The position limits of the joints.
"""
joint_idxs = (
names_to_idxs(joint_names=joint_names, model=model)
if joint_names is not None
else jnp.arange(model.number_of_joints())
)
if len(joint_idxs) == 0:
return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min[joint_idxs]
s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max[joint_idxs]
return s_min.astype(float), s_max.astype(float)
# ======================
# Random data generation
# ======================
@functools.partial(jax.jit, static_argnames=["joint_names"])
@js.common.named_scope
def random_joint_positions(
model: js.model.JaxSimModel,
*,
joint_names: Sequence[str] | None = None,
key: jax.Array | None = None,
) -> jtp.Vector:
"""
Generate random joint positions.
Args:
model: The model to consider.
joint_names: The names of the considered joints (all if None).
key: The random key (initialized from seed 0 if None).
Note:
If the joint range or revolute joints is larger than 2π, their joint positions
will be sampled from an interval of size 2π.
Returns:
The random joint positions.
"""
# Consider the key corresponding to a zero seed if it was not passed.
key = key if key is not None else jax.random.PRNGKey(seed=0)
# Get the joint limits parsed from the model description.
s_min, s_max = position_limits(model=model, joint_names=joint_names)
# Get the joint indices.
# Note that it will trigger an exception if the given `joint_names` are not valid.
joint_names = joint_names if joint_names is not None else model.joint_names()
joint_indices = (
names_to_idxs(model=model, joint_names=joint_names)
if joint_names is not None
else jnp.arange(model.number_of_joints())
)
from jaxsim.parsers.descriptions.joint import JointType
# Filter for revolute joints.
is_revolute = jnp.where(
jnp.array(model.kin_dyn_parameters.joint_model.joint_types[1:])[joint_indices]
== JointType.Revolute,
True,
False,
)
# Shorthand for π.
π = jnp.pi
# Filter for revolute with full range (or continuous).
is_revolute_full_range = jnp.logical_and(is_revolute, s_max - s_min >= 2 * π)
# Clip the lower limit to -π if the joint range is larger than [-π, π].
s_min = jnp.where(
jnp.logical_and(
is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
),
-π,
s_min,
)
# Clip the upper limit to +π if the joint range is larger than [-π, π].
s_max = jnp.where(
jnp.logical_and(
is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
),
π,
s_max,
)
# Shift the lower limit if the upper limit is smaller than +π.
s_min = jnp.where(
jnp.logical_and(is_revolute_full_range, s_max < π),
s_max - 2 * π,
s_min,
)
# Shift the upper limit if the lower limit is larger than -π.
s_max = jnp.where(
jnp.logical_and(is_revolute_full_range, s_min > -π),
s_min + 2 * π,
s_max,
)
# Sample the joint positions.
s_random = jax.random.uniform(
minval=s_min,
maxval=s_max,
key=key,
shape=s_min.shape,
)
return s_random
================================================
FILE: src/jaxsim/api/kin_dyn_parameters.py
================================================
from __future__ import annotations
import dataclasses
from itertools import starmap
from typing import ClassVar
import jax.lax
import jax.numpy as jnp
import jax_dataclasses
import numpy as np
import numpy.typing as npt
from jax_dataclasses import Static
import jaxsim
import jaxsim.typing as jtp
from jaxsim.math import Inertia, JointModel, supported_joint_motion
from jaxsim.math.adjoint import Adjoint
from jaxsim.parsers.descriptions import JointDescription, JointType, ModelDescription
from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
class KinDynParameters(JaxsimDataclass):
r"""
Class storing the kinematic and dynamic parameters of a model.
Attributes:
link_names: The names of the links.
parent_array: The parent array :math:`\lambda(i)` of the model.
support_body_array_bool:
The boolean support parent array :math:`\kappa_{b}(i)` of the model.
link_parameters: The parameters of the links.
frame_parameters: The parameters of the frames.
contact_parameters: The parameters of the collidable points.
joint_model: The joint model of the model.
joint_parameters: The parameters of the joints.
hw_link_metadata: The hardware parameters of the model links.
constraints: The kinematic constraints of the model. They can be used only with Relaxed-Rigid contact model.
"""
# Static
link_names: Static[tuple[str]]
_parent_array: Static[HashedNumpyArray]
_support_body_array_bool: Static[HashedNumpyArray]
_motion_subspaces: Static[HashedNumpyArray]
# Tree level structure for parallel algorithms.
# level_nodes: (n_levels, max_width) array of link indices at each depth level,
# padded with 0 for levels with fewer nodes than max_width.
# level_mask: (n_levels, max_width) boolean mask, True for real nodes.
_level_nodes: Static[HashedNumpyArray]
_level_mask: Static[HashedNumpyArray]
# Links
link_parameters: LinkParameters
# Contacts
contact_parameters: ContactParameters
# Frames
frame_parameters: FrameParameters
# Joints
joint_model: JointModel
joint_parameters: JointParameters | None
# Model hardware parameters
hw_link_metadata: HwLinkMetadata | None = dataclasses.field(default=None)
# Kinematic constraints
constraints: ConstraintMap | None = dataclasses.field(default=None)
@property
def motion_subspaces(self) -> jtp.Matrix:
r"""
Return the motion subspaces :math:`\mathbf{S}(s)` of the joints.
"""
return self._motion_subspaces.get()
@property
def parent_array(self) -> jtp.Vector:
r"""
Return the parent array :math:`\lambda(i)` of the model.
"""
return self._parent_array.get()
@property
def support_body_array_bool(self) -> jtp.Matrix:
r"""
Return the boolean support parent array :math:`\kappa_{b}(i)` of the model.
"""
return self._support_body_array_bool.get()
@property
def level_nodes(self) -> jtp.Matrix:
r"""
Return the tree level nodes array of shape ``(n_levels, max_width)``.
Each row contains the link indices at the corresponding depth level,
padded with 0 for levels with fewer nodes than ``max_width``.
"""
return self._level_nodes.get()
@property
def level_mask(self) -> jtp.Matrix:
r"""
Return the tree level mask of shape ``(n_levels, max_width)``.
Each entry is ``True`` for real nodes and ``False`` for padding.
"""
return self._level_mask.get()
@staticmethod
def _compute_tree_levels(
parent_array: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""
Compute the tree level decomposition from a parent array.
Args:
parent_array: Array of shape ``(n,)`` where ``parent_array[i]``
is the parent of link ``i``. ``parent_array[0] == -1`` for the root.
Returns:
A tuple ``(level_nodes, level_mask)`` where:
- ``level_nodes`` has shape ``(n_levels, max_width)`` with link
indices at each depth level (padded with 0).
- ``level_mask`` has shape ``(n_levels, max_width)`` with ``True``
for real nodes.
"""
import numpy as np
n = len(parent_array)
# Compute depth of each node.
depth = np.zeros(n, dtype=int)
for i in range(1, n):
depth[i] = depth[parent_array[i]] + 1
max_depth = int(depth.max()) if n > 0 else 0
n_levels = max_depth + 1
# Group nodes by depth level.
levels: list[list[int]] = [[] for _ in range(n_levels)]
for i in range(n):
levels[depth[i]].append(i)
max_width = max(len(lev) for lev in levels) if levels else 1
# Build padded arrays.
level_nodes = np.zeros((n_levels, max_width), dtype=int)
level_mask = np.zeros((n_levels, max_width), dtype=bool)
for d, lev in enumerate(levels):
for j, node_idx in enumerate(lev):
level_nodes[d, j] = node_idx
level_mask[d, j] = True
return level_nodes, level_mask
@staticmethod
def build(
model_description: ModelDescription, constraints: ConstraintMap | None
) -> KinDynParameters:
"""
Construct the kinematic and dynamic parameters of the model.
Args:
model_description: The parsed model description to consider.
constraints: An object of type ConstraintMap specifying the kinematic constraint of the model.
Returns:
The kinematic and dynamic parameters of the model.
Note:
This class is meant to ease the management of parametric models in
an automatic differentiation context.
"""
# Extract the links ordered by their index.
# The link index corresponds to the body index ∈ [0, num_bodies - 1].
ordered_links = sorted(
list(model_description.links_dict.values()),
key=lambda l: l.index,
)
# Extract the joints ordered by their index.
# The joint index matches the index of its child link, therefore it starts
# from 1. Keep this in mind since this 1-indexing might introduce bugs.
ordered_joints = sorted(
list(model_description.joints_dict.values()),
key=lambda j: j.index,
)
# ================
# Links properties
# ================
# Create a list of link parameters objects.
link_parameters_list = [
LinkParameters.build_from_spatial_inertia(index=link.index, M=link.inertia)
for link in ordered_links
]
# Create a vectorized object of link parameters.
link_parameters = jax.tree.map(lambda *l: jnp.stack(l), *link_parameters_list)
# =================
# Joints properties
# =================
# Create a list of joint parameters objects.
joint_parameters_list = [
JointParameters.build_from_joint_description(joint_description=joint)
for joint in ordered_joints
]
# Create a vectorized object of joint parameters.
joint_parameters = (
jax.tree.map(lambda *l: jnp.stack(l), *joint_parameters_list)
if ordered_joints
else JointParameters(
index=jnp.array([], dtype=int),
friction_static=jnp.array([], dtype=float),
friction_viscous=jnp.array([], dtype=float),
position_limits_min=jnp.array([], dtype=float),
position_limits_max=jnp.array([], dtype=float),
position_limit_spring=jnp.array([], dtype=float),
position_limit_damper=jnp.array([], dtype=float),
)
)
# Create an object that defines the joint model (parent-to-child transforms).
joint_model = JointModel.build(description=model_description)
# ===================
# Contacts properties
# ===================
# Create the object storing the parameters of collidable points.
# Note that, contrarily to LinkParameters and JointsParameters, this object
# is not created with vmap. This is because the "body" attribute of the object
# must be Static for JIT-related reasons, and tree_map would not consider it
# as a leaf.
contact_parameters = ContactParameters.build_from(
model_description=model_description
)
# =================
# Frames properties
# =================
# Create the object storing the parameters of frames.
# Note that, contrarily to LinkParameters and JointsParameters, this object
# is not created with vmap. This is because the "name" attribute of the object
# must be Static for JIT-related reasons, and tree_map would not consider it
# as a leaf.
frame_parameters = FrameParameters.build_from(
model_description=model_description
)
# ===============
# Tree properties
# ===============
# Build the parent array λ(i) of the model.
# Note: the parent of the base link is not set since it's not defined.
parent_array_dict = {
link.index: model_description.links_dict[link.parent_name].index
for link in ordered_links
if link.parent_name is not None
}
parent_array = jnp.array([-1, *list(parent_array_dict.values())], dtype=int)
# Instead of building the support parent array κ(i) for each link of the model,
# that has a variable length depending on the number of links connecting the
# root to the i-th link, we build the corresponding boolean version.
# Given a link index i, the boolean support parent array κb(i) is an array
# with the same number of elements of λ(i) having the i-th element set to True
# if the i-th link is in the support parent array κ(i), False otherwise.
# We store the boolean κb(i) as static attribute of the PyTree so that
# algorithms that need to access it can be jit-compiled.
def κb(link_index: jtp.IntLike) -> jtp.Vector:
κb = jnp.zeros(len(ordered_links), dtype=bool)
carry0 = κb, link_index
def scan_body(carry: tuple, i: jtp.Int) -> tuple[tuple, None]:
κb, active_link_index = carry
κb, active_link_index = jax.lax.cond(
pred=(i == active_link_index),
false_fun=lambda: (κb, active_link_index),
true_fun=lambda: (
κb.at[active_link_index].set(True),
parent_array[active_link_index],
),
)
return (κb, active_link_index), None
(κb, _), _ = jax.lax.scan(
f=scan_body,
init=carry0,
xs=jnp.flip(jnp.arange(start=0, stop=len(ordered_links))),
)
return κb
support_body_array_bool = jax.vmap(κb)(
jnp.arange(start=0, stop=len(ordered_links))
)
def motion_subspace(joint_type: int, axis: npt.ArrayLike) -> npt.ArrayLike:
S = {
JointType.Fixed: np.zeros(shape=(6, 1)),
JointType.Revolute: np.vstack(np.hstack([np.zeros(3), axis.axis])),
JointType.Prismatic: np.vstack(np.hstack([axis.axis, np.zeros(3)])),
}
return S[joint_type]
S_J = jnp.array(
list(
starmap(
motion_subspace,
zip(
joint_model.joint_types[1:], joint_model.joint_axis, strict=True
),
)
)
if len(joint_model.joint_axis) != 0
else jnp.empty((0, 6, 1))
)
motion_subspaces = jnp.vstack([jnp.zeros((6, 1))[jnp.newaxis, ...], S_J])
# ====================
# Tree level structure
# ====================
parent_array_np = np.array([-1, *list(parent_array_dict.values())], dtype=int)
level_nodes, level_mask = KinDynParameters._compute_tree_levels(parent_array_np)
# ===========
# Constraints
# ===========
constraints = ConstraintMap() if constraints is None else constraints
# =================================
# Build and return KinDynParameters
# =================================
return KinDynParameters(
link_names=tuple(l.name for l in ordered_links),
_parent_array=HashedNumpyArray(array=parent_array),
_support_body_array_bool=HashedNumpyArray(array=support_body_array_bool),
_motion_subspaces=HashedNumpyArray(array=motion_subspaces),
_level_nodes=HashedNumpyArray(array=level_nodes),
_level_mask=HashedNumpyArray(array=level_mask),
link_parameters=link_parameters,
joint_model=joint_model,
joint_parameters=joint_parameters,
contact_parameters=contact_parameters,
frame_parameters=frame_parameters,
constraints=constraints,
)
def __eq__(self, other: KinDynParameters) -> bool:
if not isinstance(other, KinDynParameters):
return False
return hash(self) == hash(other)
def __hash__(self) -> int:
return hash(
(
hash(self.number_of_links()),
hash(self.number_of_joints()),
hash(self.frame_parameters.name),
hash(self.frame_parameters.body),
hash(self._parent_array),
hash(self._support_body_array_bool),
)
)
# =============================
# Helpers to extract parameters
# =============================
def number_of_links(self) -> int:
"""
Return the number of links of the model.
Returns:
The number of links of the model.
"""
return len(self.link_names)
def number_of_joints(self) -> int:
"""
Return the number of joints of the model.
Returns:
The number of joints of the model.
"""
return len(self.joint_model.joint_names) - 1
def number_of_frames(self) -> int:
"""
Return the number of frames of the model.
Returns:
The number of frames of the model.
"""
return len(self.frame_parameters.name)
def support_body_array(self, link_index: jtp.IntLike) -> jtp.Vector:
r"""
Return the support parent array :math:`\kappa(i)` of a link.
Args:
link_index: The index of the link.
Returns:
The support parent array :math:`\kappa(i)` of the link.
Note:
This method returns a variable-length vector. In jit-compiled functions,
it's better to use the (static) boolean version `support_body_array_bool`.
"""
return jnp.array(
jnp.where(self.support_body_array_bool[link_index])[0], dtype=int
)
# ========================
# Quantities used by RBDAs
# ========================
@jax.jit
def links_spatial_inertia(self) -> jtp.Array:
"""
Return the spatial inertia of all links of the model.
Returns:
The spatial inertia of all links of the model.
"""
return jax.vmap(LinkParameters.spatial_inertia)(self.link_parameters)
@jax.jit
def tree_transforms(self) -> jtp.Array:
r"""
Return the tree transforms of the model.
Returns:
The transforms
:math:`{}^{\text{pre}(i)} H_{\lambda(i)}`
of all joints of the model.
"""
pre_Xi_λ = jax.vmap(
lambda i: self.joint_model.parent_H_predecessor(joint_index=i)
.inverse()
.adjoint()
)(jnp.arange(1, self.number_of_joints() + 1))
return jnp.vstack(
[
jnp.zeros(shape=(1, 6, 6), dtype=float),
pre_Xi_λ,
]
)
@jax.jit
def joint_transforms(
self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
) -> jtp.Array:
r"""
Return the transforms of the joints.
Args:
joint_positions: The joint positions.
base_transform: The homogeneous matrix defining the base pose.
Returns:
The stacked transforms
:math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
of each joint.
"""
# Rename the base transform.
W_H_B = base_transform
# Extract the parent-to-predecessor fixed transforms of the joints.
λ_H_pre = jnp.vstack(
[
jnp.eye(4)[jnp.newaxis],
self.joint_model.λ_H_pre[1 : 1 + self.number_of_joints()],
]
)
if self.number_of_joints() == 0:
pre_H_suc_J = jnp.empty((0, 4, 4))
else:
pre_H_suc_J = jax.vmap(supported_joint_motion)(
joint_types=jnp.array(self.joint_model.joint_types[1:]).astype(int),
joint_positions=jnp.array(joint_positions),
joint_axes=jnp.array([j.axis for j in self.joint_model.joint_axis]),
)
# Extract the transforms and motion subspaces of the joints.
# We stack the base transform W_H_B at index 0, and a dummy motion subspace
# for either the fixed or free-floating joint connecting the world to the base.
pre_H_suc = jnp.vstack([W_H_B[jnp.newaxis, ...], pre_H_suc_J])
# Extract the successor-to-child fixed transforms.
# Note that here we include also the index 0 since suc_H_child[0] stores the
# optional pose of the base link w.r.t. the root frame of the model.
# This is supported by SDF when the base link