chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Publicize the rollout notebook, small improvements to plotting.

Browse Source 
PiperOrigin-RevId: 730841667
Change-Id: Id365f21cee078f28cb3fa797701c3c5a409607cd
This commit is contained in:
Yuval Tassa
2025-02-25 05:19:46 -08:00
committed by Copybara-Service
parent ce82b63155
commit dc98ae978d
4 changed files with 46 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2
View File
@@ -52,6 +52,8 @@ running on Google Colab:
- The **introductory** tutorial teaches MuJoCo basics: - The **introductory** tutorial teaches MuJoCo basics:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)
- The **rollout** tutorial shows how to use the multithreaded `rollout` module:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb)
- The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a humanoid on one leg: - The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a humanoid on one leg:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb)
- The **least-squares** tutorial explains how to use the Python-based nonlinear least-squares solver: - The **least-squares** tutorial explains how to use the Python-based nonlinear least-squares solver:
doc/images/python/rollout.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 45 KiB

doc/python.rst
+13 -4
View File
@@ -736,15 +736,24 @@ The ``mujoco`` package contains two sub-modules: ``mujoco.rollout`` and ``mujoco
rollout rollout
------- -------
``mujoco.rollout`` and ``mujoco.rollout.Rollout`` shows how to add additional C/C++ functionality, exposed as a Python ``mujoco.rollout`` and ``mujoco.rollout.Rollout`` shows how to add additional C/C++ functionality, exposed as a Python
module via pybind11. It is implemented in `rollout.cc module via pybind11. It is implemented in `rollout.cc
<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__ and wrapped in `rollout.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__ and wrapped in `rollout.py
<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__. The module performs a common <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__. The module addresses a common
functionality where tight loops implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling use-case where tight loops implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling
:ref:`mj_step` in a loop), given an initial state and sequence of controls, and returning subsequent states and sensor :ref:`mj_step` in a loop), given an initial state and sequence of controls, and returning subsequent states and sensor
values. The rollouts are run in parallel with an internally managed thread pool if multiple MjData instances (one per values. The rollouts are run in parallel with an internally managed thread pool if multiple MjData instances (one per
thread) are passed as an argument. The basic usage form is thread) are passed as an argument. This notebook shows how to use ``rollout`` |rollout_colab|, along with some
benchmarks e.g., the figure below.
.. |rollout_colab| image:: https://colab.research.google.com/assets/colab-badge.svg
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb
.. image:: images/python/rollout.png
:align: right
:width: 97%
The basic usage form is
.. code-block:: python .. code-block:: python
python/rollout.ipynb
+31 -29
View File
@@ -51,7 +51,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f", "id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
"metadata": { "metadata": {
"editable": true, "editable": true,
@@ -166,7 +166,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "082482c7", "id": "082482c7",
"metadata": { "metadata": {
"editable": true, "editable": true,
@@ -288,7 +288,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "849b93e5", "id": "849b93e5",
"metadata": { "metadata": {
"id": "849b93e5" "id": "849b93e5"
@@ -392,7 +392,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5", "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
"metadata": { "metadata": {
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5" "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5"
@@ -418,7 +418,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe", "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
"metadata": { "metadata": {
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe" "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe"
@@ -463,7 +463,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "957b8566-da31-410b-b385-e78241c5247a", "id": "957b8566-da31-410b-b385-e78241c5247a",
"metadata": { "metadata": {
"id": "957b8566-da31-410b-b385-e78241c5247a" "id": "957b8566-da31-410b-b385-e78241c5247a"
@@ -497,7 +497,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2", "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
"metadata": { "metadata": {
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2" "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2"
@@ -578,7 +578,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03", "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
"metadata": { "metadata": {
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03" "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03"
@@ -610,7 +610,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "2a184873-8d24-45da-b444-8d21f5dcd733", "id": "2a184873-8d24-45da-b444-8d21f5dcd733",
"metadata": { "metadata": {
"id": "2a184873-8d24-45da-b444-8d21f5dcd733" "id": "2a184873-8d24-45da-b444-8d21f5dcd733"
@@ -676,7 +676,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5", "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
"metadata": { "metadata": {
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5" "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5"
@@ -747,7 +747,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95", "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
"metadata": { "metadata": {
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95" "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95"
@@ -820,7 +820,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49", "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
"metadata": { "metadata": {
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49" "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49"
@@ -875,7 +875,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "b6aa6801", "id": "b6aa6801",
"metadata": { "metadata": {
"id": "b6aa6801" "id": "b6aa6801"
@@ -908,7 +908,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "7f46a6d8", "id": "7f46a6d8",
"metadata": { "metadata": {
"id": "7f46a6d8" "id": "7f46a6d8"
@@ -941,7 +941,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "a1be8f93", "id": "a1be8f93",
"metadata": { "metadata": {
"id": "a1be8f93" "id": "a1be8f93"
@@ -1004,7 +1004,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb", "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
"metadata": { "metadata": {
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb" "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb"
@@ -1098,7 +1098,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "cb6355dd", "id": "cb6355dd",
"metadata": { "metadata": {
"id": "cb6355dd" "id": "cb6355dd"
@@ -1128,7 +1128,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "74f143e2", "id": "74f143e2",
"metadata": { "metadata": {
"id": "74f143e2" "id": "74f143e2"
@@ -1155,7 +1155,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "0301e3ee", "id": "0301e3ee",
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
@@ -1253,6 +1253,7 @@
" ax1.set_xticks(x + width, nbatch)\n", " ax1.set_xticks(x + width, nbatch)\n",
" ax1.yaxis.set_major_formatter(ticker)\n", " ax1.yaxis.set_major_formatter(ticker)\n",
" ax1.grid()\n", " ax1.grid()\n",
" ax1.set_axisbelow(True)\n",
" ax1.set_xlabel('nbatch')\n", " ax1.set_xlabel('nbatch')\n",
" ax1.set_ylabel('steps per second')\n", " ax1.set_ylabel('steps per second')\n",
" ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n", " ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n",
@@ -1268,11 +1269,12 @@
" ax2.set_xticks(x + width, nstep)\n", " ax2.set_xticks(x + width, nstep)\n",
" ax2.yaxis.set_major_formatter(ticker)\n", " ax2.yaxis.set_major_formatter(ticker)\n",
" ax2.grid()\n", " ax2.grid()\n",
" ax2.set_axisbelow(True)\n",
" ax2.set_xlabel('nstep')\n", " ax2.set_xlabel('nstep')\n",
" ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n", " ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n",
"\n", "\n",
" ax2.legend(loc=(1.04, 0.0))\n", " ax1.legend(loc=(0.03, 0.8))\n",
" fig.set_size_inches(10, 4)\n", " fig.set_size_inches(10, 5)\n",
" plt.suptitle(title)\n", " plt.suptitle(title)\n",
" plt.tight_layout()" " plt.tight_layout()"
] ]
@@ -1289,7 +1291,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "f7e54830", "id": "f7e54830",
"metadata": { "metadata": {
"id": "f7e54830" "id": "f7e54830"
@@ -1321,7 +1323,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "c9e58c6c", "id": "c9e58c6c",
"metadata": { "metadata": {
"id": "c9e58c6c" "id": "c9e58c6c"
@@ -1353,7 +1355,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "83d775d4", "id": "83d775d4",
"metadata": { "metadata": {
"id": "83d775d4" "id": "83d775d4"
@@ -1404,7 +1406,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "7c86d157", "id": "7c86d157",
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
@@ -1538,7 +1540,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "98c580b0", "id": "98c580b0",
"metadata": { "metadata": {
"id": "98c580b0" "id": "98c580b0"
@@ -1567,7 +1569,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "53166ae1", "id": "53166ae1",
"metadata": { "metadata": {
"id": "53166ae1" "id": "53166ae1"
@@ -1602,7 +1604,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "3d6be608", "id": "3d6be608",
"metadata": { "metadata": {
"id": "3d6be608" "id": "3d6be608"
@@ -1654,7 +1656,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"id": "b6c5fc2e", "id": "b6c5fc2e",
"metadata": { "metadata": {
"id": "b6c5fc2e" "id": "b6c5fc2e"