Skip to content

Commit

Permalink
Update notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
@diegoferigo
diegoferigo committed Oct 17, 2024
1 parent 1a4bf4e commit 056073a
Showing 1 changed file with 40 additions and 57 deletions.
97 changes: 40 additions & 57 deletions examples/jaxsim_as_multibody_dynamics_library.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,15 @@
"id": "DpLq0-lltwZ1"
},
"source": [
"# `JaxSim` as a Multibody Dynamics Library\n",
"\n",
"# `JaxSim` as a multibody dynamics library\n",
"\n",
"JaxSim was initially developed as a **hardware-accelerated physics engine**. Over time, it has evolved, adding new features to become a comprehensive **JAX-based multibody dynamics library**.\n",
"\n",
"In this notebook, you'll explore the main APIs for loading robot models and computing key quantities for applications such as control, planning, and more.\n",
"\n",
"A key advantage of JaxSim is its ability to create fully differentiable closed-loop systems, enabling end-to-end optimization. Combined with the flexibility to parameterize model kinematics and dynamics, JaxSim can serve as an excellent playground for robot learning applications.\n",
"\n",
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/Rigid_body_dynamics.ipynb\">\n",
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/jaxsim_as_multibody_dynamics_library.ipynb\">\n",
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
"</a>\n"
]
Expand All @@ -35,26 +34,23 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "u4xL7dbBtwZ3",
"outputId": "b30c4790-f0f4-4ce6-f439-db662176640e"
"outputId": "1a088e28-e005-4910-928c-cb641e589ab5"
},
"outputs": [],
"source": [
"# @title Imports and setup\n",
"from IPython.display import clear_output\n",
"import os\n",
"import pathlib\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 -qU jaxsim\n",
" !{sys.executable} -m pip install --pre -qU jaxsim\n",
" !{sys.executable} -m pip install robot_descriptions\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",
Expand All @@ -64,18 +60,19 @@
"\n",
" clear_output()\n",
"\n",
"import os\n",
"import pathlib\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"from jaxsim import logging\n",
"import robot_descriptions\n",
"\n",
"import jaxsim.api as js\n",
"import jaxsim.math\n",
"import robot_descriptions\n",
"from jaxsim import logging\n",
"from jaxsim import VelRepr\n",
"\n",
"logging.set_logging_level(logging.LoggingLevel.INFO)\n",
"logging.info(f\"Running on {jax.devices()}\")"
"logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
"print(f\"Running on {jax.devices()}\")"
]
},
{
Expand Down Expand Up @@ -104,32 +101,27 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "rB0BFxyPtwZ5",
"outputId": "78d99087-7dd3-44dc-f9f4-875472946942"
"id": "rB0BFxyPtwZ5"
},
"outputs": [],
"source": [
"# @title Fetch the URDF file\n",
"from IPython.utils import io\n",
"\n",
"with io.capture_output() as captured:\n",
" try:\n",
" os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.1\"\n",
"try:\n",
" os.environ[\"ROBOT_DESCRIPTION_COMMIT\"] = \"v0.7.1\"\n",
"\n",
" import robot_descriptions.ergocub_description\n",
" import robot_descriptions.ergocub_description\n",
"\n",
" finally:\n",
" _ = os.environ.pop(\"ROBOT_DESCRIPTION_COMMIT\", None)\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",
" \"ergoCubSN000\", \"ergoCubSN001\"\n",
" )\n",
" )"
"model_description_path = pathlib.Path(\n",
" robot_descriptions.ergocub_description.URDF_PATH.replace(\n",
" \"ergoCubSN000\", \"ergoCubSN001\"\n",
" )\n",
")\n",
"\n",
"clear_output()"
]
},
{
Expand All @@ -143,7 +135,7 @@
"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 \\in \\mathcal{L}} J_{W,L}^\\top(\\mathbf{q}) \\: \\mathbf{f}_L\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",
Expand All @@ -155,7 +147,7 @@
"The inputs to the system are:\n",
"\n",
"- $\\boldsymbol{\\tau} \\in \\mathbb{R}^n$ are the joint torques.\n",
"- $\\mathbf{f}_L \\in \\mathbb{R}^6$ is the 6D force applied to the frame associated to link $L$.\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",
Expand All @@ -175,11 +167,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "WYgBAxU0twZ6",
"outputId": "994b1cd5-d742-4735-a713-ed64a2ad1afa"
"id": "WYgBAxU0twZ6"
},
"outputs": [],
"source": [
Expand All @@ -205,11 +193,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "QuhG7Zv5twZ7",
"outputId": "0a7fd714-218c-4a35-ce78-5d5a5c2bbbd0"
"id": "QuhG7Zv5twZ7"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -242,7 +226,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "RLvAit_i2ZiA",
"outputId": "f6a77f31-e341-4350-c057-b1edabfe8f0e"
"outputId": "ea3954af-b9b9-46ac-d9cb-20b99b1eac94"
},
"outputs": [],
"source": [
Expand All @@ -269,7 +253,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "Xp8V5on5twZ8",
"outputId": "5e1c50e3-267d-42aa-fd1c-6cac3e246790"
"outputId": "cc1564db-ae91-4dba-92c9-b8b87bd65f10"
},
"outputs": [],
"source": [
Expand All @@ -293,7 +277,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "XLx3sv9VtwZ9",
"outputId": "7f83bcd0-1099-44b1-e407-2c719aca6895"
"outputId": "28f5f070-e37e-464e-d84e-2944cfdc28dc"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -366,7 +350,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "C22Iqu2i4G-I",
"outputId": "cf1614dc-6c3c-442e-adf8-d3344a331d4a"
"outputId": "94376151-177d-410f-f375-b7b8bd080992"
},
"outputs": [],
"source": [
Expand All @@ -392,7 +376,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "DnSpE_f97RkX",
"outputId": "47765e57-3fc1-40b6-ac20-04392aba85ef"
"outputId": "a3f6b535-4ae5-49f4-8921-7fe4dda5debb"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -439,7 +423,7 @@
"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 \\in \\mathcal{L}} J_{W,L}^\\top(\\mathbf{q}) \\: \\mathbf{f}_L\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",
Expand Down Expand Up @@ -480,7 +464,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "fVp_xP_1twZ9",
"outputId": "b06f20a0-6c01-435c-e353-502045234487"
"outputId": "cfaa0569-d768-4708-c98c-a5867c056d04"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -550,7 +534,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "d_vp6D74GoVZ",
"outputId": "ae1c1b59-f5ce-4539-bca0-5d34d62995c8"
"outputId": "798b9283-792e-4339-b56c-df2595fac974"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -595,7 +579,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "FlNo8dNWKKtu",
"outputId": "867e7568-016d-4418-a2bd-ab7d75fa69a7"
"outputId": "313e939b-f88f-4407-c9ee-b5b3b7443061"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -676,7 +660,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "g5GOYXDnLySU",
"outputId": "b6e5df8f-1d56-4208-a5b0-87dfa1737ca4"
"outputId": "ad4ce77d-d06f-473a-9c32-040680d76aa5"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -766,7 +750,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "gYZ1jK1Neg1H",
"outputId": "274a628d-efad-4738-b500-bb2927f3e826"
"outputId": "0de79770-1e18-4027-bb47-5713bc1b4a72"
},
"outputs": [],
"source": [
Expand Down Expand Up @@ -834,11 +818,10 @@
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 422
"base_uri": "https://localhost:8080/"
},
"id": "Ot6HePB_twaE",
"outputId": "5d1d717c-4523-402d-9763-1d485720aa3b"
"outputId": "02a6abae-257e-45ee-e9de-6a607cdbeb9a"
},
"outputs": [],
"source": [
Expand Down

0 comments on commit 056073a

Please sign in to comment.