Skip to content

Commit dc98ae9

Browse files
yuvaltassacopybara-github
authored andcommitted
Publicize the rollout notebook, small improvements to plotting.
PiperOrigin-RevId: 730841667 Change-Id: Id365f21cee078f28cb3fa797701c3c5a409607cd
1 parent ce82b63 commit dc98ae9

File tree

4 files changed

+46
-33
lines changed

4 files changed

+46
-33
lines changed

README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ running on Google Colab:
5252

5353
- The **introductory** tutorial teaches MuJoCo basics:
5454
[![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)
55+
- The **rollout** tutorial shows how to use the multithreaded `rollout` module:
56+
[![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)
5557
- The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a humanoid on one leg:
5658
[![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)
5759
- The **least-squares** tutorial explains how to use the Python-based nonlinear least-squares solver:

doc/images/python/rollout.png

45.3 KB
Loading

doc/python.rst

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -736,15 +736,24 @@ The ``mujoco`` package contains two sub-modules: ``mujoco.rollout`` and ``mujoco
736736

737737
rollout
738738
-------
739-
740739
``mujoco.rollout`` and ``mujoco.rollout.Rollout`` shows how to add additional C/C++ functionality, exposed as a Python
741740
module via pybind11. It is implemented in `rollout.cc
742741
<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__ and wrapped in `rollout.py
743-
<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__. The module performs a common
744-
functionality where tight loops implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling
742+
<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__. The module addresses a common
743+
use-case where tight loops implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling
745744
:ref:`mj_step` in a loop), given an initial state and sequence of controls, and returning subsequent states and sensor
746745
values. The rollouts are run in parallel with an internally managed thread pool if multiple MjData instances (one per
747-
thread) are passed as an argument. The basic usage form is
746+
thread) are passed as an argument. This notebook shows how to use ``rollout`` |rollout_colab|, along with some
747+
benchmarks e.g., the figure below.
748+
749+
.. |rollout_colab| image:: https://colab.research.google.com/assets/colab-badge.svg
750+
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb
751+
752+
.. image:: images/python/rollout.png
753+
:align: right
754+
:width: 97%
755+
756+
The basic usage form is
748757

749758
.. code-block:: python
750759

python/rollout.ipynb

Lines changed: 31 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
},
5252
{
5353
"cell_type": "code",
54-
"execution_count": null,
54+
"execution_count": 0,
5555
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
5656
"metadata": {
5757
"editable": true,
@@ -166,7 +166,7 @@
166166
},
167167
{
168168
"cell_type": "code",
169-
"execution_count": null,
169+
"execution_count": 0,
170170
"id": "082482c7",
171171
"metadata": {
172172
"editable": true,
@@ -288,7 +288,7 @@
288288
},
289289
{
290290
"cell_type": "code",
291-
"execution_count": null,
291+
"execution_count": 0,
292292
"id": "849b93e5",
293293
"metadata": {
294294
"id": "849b93e5"
@@ -392,7 +392,7 @@
392392
},
393393
{
394394
"cell_type": "code",
395-
"execution_count": null,
395+
"execution_count": 0,
396396
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
397397
"metadata": {
398398
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5"
@@ -418,7 +418,7 @@
418418
},
419419
{
420420
"cell_type": "code",
421-
"execution_count": null,
421+
"execution_count": 0,
422422
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
423423
"metadata": {
424424
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe"
@@ -463,7 +463,7 @@
463463
},
464464
{
465465
"cell_type": "code",
466-
"execution_count": null,
466+
"execution_count": 0,
467467
"id": "957b8566-da31-410b-b385-e78241c5247a",
468468
"metadata": {
469469
"id": "957b8566-da31-410b-b385-e78241c5247a"
@@ -497,7 +497,7 @@
497497
},
498498
{
499499
"cell_type": "code",
500-
"execution_count": null,
500+
"execution_count": 0,
501501
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
502502
"metadata": {
503503
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2"
@@ -578,7 +578,7 @@
578578
},
579579
{
580580
"cell_type": "code",
581-
"execution_count": null,
581+
"execution_count": 0,
582582
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
583583
"metadata": {
584584
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03"
@@ -610,7 +610,7 @@
610610
},
611611
{
612612
"cell_type": "code",
613-
"execution_count": null,
613+
"execution_count": 0,
614614
"id": "2a184873-8d24-45da-b444-8d21f5dcd733",
615615
"metadata": {
616616
"id": "2a184873-8d24-45da-b444-8d21f5dcd733"
@@ -676,7 +676,7 @@
676676
},
677677
{
678678
"cell_type": "code",
679-
"execution_count": null,
679+
"execution_count": 0,
680680
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
681681
"metadata": {
682682
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5"
@@ -747,7 +747,7 @@
747747
},
748748
{
749749
"cell_type": "code",
750-
"execution_count": null,
750+
"execution_count": 0,
751751
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
752752
"metadata": {
753753
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95"
@@ -820,7 +820,7 @@
820820
},
821821
{
822822
"cell_type": "code",
823-
"execution_count": null,
823+
"execution_count": 0,
824824
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
825825
"metadata": {
826826
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49"
@@ -875,7 +875,7 @@
875875
},
876876
{
877877
"cell_type": "code",
878-
"execution_count": null,
878+
"execution_count": 0,
879879
"id": "b6aa6801",
880880
"metadata": {
881881
"id": "b6aa6801"
@@ -908,7 +908,7 @@
908908
},
909909
{
910910
"cell_type": "code",
911-
"execution_count": null,
911+
"execution_count": 0,
912912
"id": "7f46a6d8",
913913
"metadata": {
914914
"id": "7f46a6d8"
@@ -941,7 +941,7 @@
941941
},
942942
{
943943
"cell_type": "code",
944-
"execution_count": null,
944+
"execution_count": 0,
945945
"id": "a1be8f93",
946946
"metadata": {
947947
"id": "a1be8f93"
@@ -1004,7 +1004,7 @@
10041004
},
10051005
{
10061006
"cell_type": "code",
1007-
"execution_count": null,
1007+
"execution_count": 0,
10081008
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
10091009
"metadata": {
10101010
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb"
@@ -1098,7 +1098,7 @@
10981098
},
10991099
{
11001100
"cell_type": "code",
1101-
"execution_count": null,
1101+
"execution_count": 0,
11021102
"id": "cb6355dd",
11031103
"metadata": {
11041104
"id": "cb6355dd"
@@ -1128,7 +1128,7 @@
11281128
},
11291129
{
11301130
"cell_type": "code",
1131-
"execution_count": null,
1131+
"execution_count": 0,
11321132
"id": "74f143e2",
11331133
"metadata": {
11341134
"id": "74f143e2"
@@ -1155,7 +1155,7 @@
11551155
},
11561156
{
11571157
"cell_type": "code",
1158-
"execution_count": null,
1158+
"execution_count": 0,
11591159
"id": "0301e3ee",
11601160
"metadata": {
11611161
"cellView": "form",
@@ -1253,6 +1253,7 @@
12531253
" ax1.set_xticks(x + width, nbatch)\n",
12541254
" ax1.yaxis.set_major_formatter(ticker)\n",
12551255
" ax1.grid()\n",
1256+
" ax1.set_axisbelow(True)\n",
12561257
" ax1.set_xlabel('nbatch')\n",
12571258
" ax1.set_ylabel('steps per second')\n",
12581259
" ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n",
@@ -1268,11 +1269,12 @@
12681269
" ax2.set_xticks(x + width, nstep)\n",
12691270
" ax2.yaxis.set_major_formatter(ticker)\n",
12701271
" ax2.grid()\n",
1272+
" ax2.set_axisbelow(True)\n",
12711273
" ax2.set_xlabel('nstep')\n",
12721274
" ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n",
12731275
"\n",
1274-
" ax2.legend(loc=(1.04, 0.0))\n",
1275-
" fig.set_size_inches(10, 4)\n",
1276+
" ax1.legend(loc=(0.03, 0.8))\n",
1277+
" fig.set_size_inches(10, 5)\n",
12761278
" plt.suptitle(title)\n",
12771279
" plt.tight_layout()"
12781280
]
@@ -1289,7 +1291,7 @@
12891291
},
12901292
{
12911293
"cell_type": "code",
1292-
"execution_count": null,
1294+
"execution_count": 0,
12931295
"id": "f7e54830",
12941296
"metadata": {
12951297
"id": "f7e54830"
@@ -1321,7 +1323,7 @@
13211323
},
13221324
{
13231325
"cell_type": "code",
1324-
"execution_count": null,
1326+
"execution_count": 0,
13251327
"id": "c9e58c6c",
13261328
"metadata": {
13271329
"id": "c9e58c6c"
@@ -1353,7 +1355,7 @@
13531355
},
13541356
{
13551357
"cell_type": "code",
1356-
"execution_count": null,
1358+
"execution_count": 0,
13571359
"id": "83d775d4",
13581360
"metadata": {
13591361
"id": "83d775d4"
@@ -1404,7 +1406,7 @@
14041406
},
14051407
{
14061408
"cell_type": "code",
1407-
"execution_count": null,
1409+
"execution_count": 0,
14081410
"id": "7c86d157",
14091411
"metadata": {
14101412
"cellView": "form",
@@ -1538,7 +1540,7 @@
15381540
},
15391541
{
15401542
"cell_type": "code",
1541-
"execution_count": null,
1543+
"execution_count": 0,
15421544
"id": "98c580b0",
15431545
"metadata": {
15441546
"id": "98c580b0"
@@ -1567,7 +1569,7 @@
15671569
},
15681570
{
15691571
"cell_type": "code",
1570-
"execution_count": null,
1572+
"execution_count": 0,
15711573
"id": "53166ae1",
15721574
"metadata": {
15731575
"id": "53166ae1"
@@ -1602,7 +1604,7 @@
16021604
},
16031605
{
16041606
"cell_type": "code",
1605-
"execution_count": null,
1607+
"execution_count": 0,
16061608
"id": "3d6be608",
16071609
"metadata": {
16081610
"id": "3d6be608"
@@ -1654,7 +1656,7 @@
16541656
},
16551657
{
16561658
"cell_type": "code",
1657-
"execution_count": null,
1659+
"execution_count": 0,
16581660
"id": "b6c5fc2e",
16591661
"metadata": {
16601662
"id": "b6c5fc2e"

0 commit comments

Comments
 (0)