diff --git a/README.md b/README.md index d5f5a27f..db022581 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,8 @@ running on Google Colab: - 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) + - 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: [![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: diff --git a/doc/images/python/rollout.png b/doc/images/python/rollout.png new file mode 100644 index 00000000..9901e687 Binary files /dev/null and b/doc/images/python/rollout.png differ diff --git a/doc/python.rst b/doc/python.rst index e2fb7dbc..163d7896 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -736,15 +736,24 @@ The ``mujoco`` package contains two sub-modules: ``mujoco.rollout`` and ``mujoco rollout ------- - ``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 `__ and wrapped in `rollout.py -`__. The module performs a common -functionality where tight loops implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling +`__. The module addresses a common +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 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 diff --git a/python/rollout.ipynb b/python/rollout.ipynb index fb559c3d..61b21ffc 100644 --- a/python/rollout.ipynb +++ b/python/rollout.ipynb @@ -51,7 +51,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f", "metadata": { "editable": true, @@ -166,7 +166,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "082482c7", "metadata": { "editable": true, @@ -288,7 +288,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "849b93e5", "metadata": { "id": "849b93e5" @@ -392,7 +392,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5", "metadata": { "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5" @@ -418,7 +418,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe", "metadata": { "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe" @@ -463,7 +463,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "957b8566-da31-410b-b385-e78241c5247a", "metadata": { "id": "957b8566-da31-410b-b385-e78241c5247a" @@ -497,7 +497,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2", "metadata": { "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2" @@ -578,7 +578,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03", "metadata": { "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03" @@ -610,7 +610,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "2a184873-8d24-45da-b444-8d21f5dcd733", "metadata": { "id": "2a184873-8d24-45da-b444-8d21f5dcd733" @@ -676,7 +676,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5", "metadata": { "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5" @@ -747,7 +747,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95", "metadata": { "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95" @@ -820,7 +820,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49", "metadata": { "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49" @@ -875,7 +875,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "b6aa6801", "metadata": { "id": "b6aa6801" @@ -908,7 +908,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "7f46a6d8", "metadata": { "id": "7f46a6d8" @@ -941,7 +941,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "a1be8f93", "metadata": { "id": "a1be8f93" @@ -1004,7 +1004,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb", "metadata": { "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb" @@ -1098,7 +1098,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "cb6355dd", "metadata": { "id": "cb6355dd" @@ -1128,7 +1128,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "74f143e2", "metadata": { "id": "74f143e2" @@ -1155,7 +1155,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "0301e3ee", "metadata": { "cellView": "form", @@ -1253,6 +1253,7 @@ " ax1.set_xticks(x + width, nbatch)\n", " ax1.yaxis.set_major_formatter(ticker)\n", " ax1.grid()\n", + " ax1.set_axisbelow(True)\n", " ax1.set_xlabel('nbatch')\n", " ax1.set_ylabel('steps per second')\n", " ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n", @@ -1268,11 +1269,12 @@ " ax2.set_xticks(x + width, nstep)\n", " ax2.yaxis.set_major_formatter(ticker)\n", " ax2.grid()\n", + " ax2.set_axisbelow(True)\n", " ax2.set_xlabel('nstep')\n", " ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n", "\n", - " ax2.legend(loc=(1.04, 0.0))\n", - " fig.set_size_inches(10, 4)\n", + " ax1.legend(loc=(0.03, 0.8))\n", + " fig.set_size_inches(10, 5)\n", " plt.suptitle(title)\n", " plt.tight_layout()" ] @@ -1289,7 +1291,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "f7e54830", "metadata": { "id": "f7e54830" @@ -1321,7 +1323,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "c9e58c6c", "metadata": { "id": "c9e58c6c" @@ -1353,7 +1355,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "83d775d4", "metadata": { "id": "83d775d4" @@ -1404,7 +1406,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "7c86d157", "metadata": { "cellView": "form", @@ -1538,7 +1540,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "98c580b0", "metadata": { "id": "98c580b0" @@ -1567,7 +1569,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "53166ae1", "metadata": { "id": "53166ae1" @@ -1602,7 +1604,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "3d6be608", "metadata": { "id": "3d6be608" @@ -1654,7 +1656,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "id": "b6c5fc2e", "metadata": { "id": "b6c5fc2e"