Merge pull request #2589 from Rahul-Lashkari:fix-py-rollout-typos
PiperOrigin-RevId: 763686962 Change-Id: I6a9c419200f410f60a1801d74ae05862464ad9a1
This commit is contained in:
Copybara-Service
+3
-2
@@ -10,7 +10,8 @@
|
||||
"\n",
|
||||
"# <h1><center>Model Editing <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
|
||||
"\n",
|
||||
"This notebook provides an introductory tutorial for model editing in MuJoCo using the `mjSpec` API. This notebook assumes that the reader is already familiar with MuJoCo basic concepts, as demostrated in the [introductory tutorial](https://github.com/google-deepmind/mujoco?tab=readme-ov-file#getting-started). Documentation for this API can be found in the [Model Editing](https://mujoco.readthedocs.io/en/latest/programming/modeledit.html) chapter in the documentation (C API) and in the [Python chapter](https://mujoco.readthedocs.io/en/latest/python.html#model-editing). Here we use the Python API.\n",
|
||||
"This notebook provides an introductory tutorial for model editing in MuJoCo using the `mjSpec` API. This notebook assumes that the reader is already familiar with MuJoCo basic concepts, as\n",
|
||||
"demonstrated in the [introductory tutorial](https://github.com/google-deepmind/mujoco?tab=readme-ov-file#getting-started). Documentation for this API can be found in the [Model Editing](https://mujoco.readthedocs.io/en/latest/programming/modeledit.html) chapter in the documentation (C API) and in the [Python chapter](https://mujoco.readthedocs.io/en/latest/python.html#model-editing). Here we use the Python API.\n",
|
||||
"\n",
|
||||
"The goal of the API is to allow users to easily interact with and modify MuJoCo\n",
|
||||
"models in Python, similarly to what the JavaScript DOM does for HTML.\n",
|
||||
@@ -636,7 +637,7 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def add_hfield(spec=None, hsize=10, vsize=4):\n",
|
||||
" \"\"\" Function that adds a heighfield with countours\"\"\"\n",
|
||||
" \"\"\" Function that adds a height field with contours\"\"\"\n",
|
||||
"\n",
|
||||
" # Initialize spec\n",
|
||||
" if spec is None:\n",
|
||||
|
||||
+77
-186
@@ -5,7 +5,6 @@
|
||||
"id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34",
|
||||
"metadata": {
|
||||
"editable": true,
|
||||
"id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34",
|
||||
"tags": []
|
||||
},
|
||||
"source": [
|
||||
@@ -42,7 +41,6 @@
|
||||
"id": "5d8a6604-0948-4a42-a48d-249c7f0c462b",
|
||||
"metadata": {
|
||||
"editable": true,
|
||||
"id": "5d8a6604-0948-4a42-a48d-249c7f0c462b",
|
||||
"tags": []
|
||||
},
|
||||
"source": [
|
||||
@@ -55,7 +53,6 @@
|
||||
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
|
||||
"metadata": {
|
||||
"editable": true,
|
||||
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
|
||||
"tags": [
|
||||
"hide-input"
|
||||
]
|
||||
@@ -157,9 +154,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fc69d0f4",
|
||||
"metadata": {
|
||||
"id": "fc69d0f4"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Helper Functions"
|
||||
]
|
||||
@@ -170,7 +165,6 @@
|
||||
"id": "082482c7",
|
||||
"metadata": {
|
||||
"editable": true,
|
||||
"id": "082482c7",
|
||||
"tags": [
|
||||
"hide-input"
|
||||
]
|
||||
@@ -273,9 +267,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c0570c2c",
|
||||
"metadata": {
|
||||
"id": "c0570c2c"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Using `rollout`\n",
|
||||
"\n",
|
||||
@@ -288,11 +280,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "849b93e5",
|
||||
"metadata": {
|
||||
"id": "849b93e5"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Benchmarked models\n",
|
||||
@@ -378,9 +368,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "55d171f7-541b-4441-aa18-da86d6716410",
|
||||
"metadata": {
|
||||
"id": "55d171f7-541b-4441-aa18-da86d6716410"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Usage\n",
|
||||
"\n",
|
||||
@@ -392,11 +380,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
|
||||
"metadata": {
|
||||
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(rollout.rollout.__doc__)"
|
||||
@@ -406,7 +392,6 @@
|
||||
"cell_type": "markdown",
|
||||
"id": "b6f7a094-8352-4b07-99ee-5278e3036cd5",
|
||||
"metadata": {
|
||||
"id": "b6f7a094-8352-4b07-99ee-5278e3036cd5",
|
||||
"tags": []
|
||||
},
|
||||
"source": [
|
||||
@@ -418,11 +403,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
|
||||
"metadata": {
|
||||
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nbatch = 100 # Simulate this many tops\n",
|
||||
@@ -454,20 +437,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f",
|
||||
"metadata": {
|
||||
"id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Our model has an angular velocity sensor the middle of the top. Let's plot the response using the `sensordata` array that rollout returns."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "957b8566-da31-410b-b385-e78241c5247a",
|
||||
"metadata": {
|
||||
"id": "957b8566-da31-410b-b385-e78241c5247a"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(12, 8))\n",
|
||||
@@ -483,9 +462,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "58044bc1-f98c-4bbf-a703-40ba075552a0",
|
||||
"metadata": {
|
||||
"id": "58044bc1-f98c-4bbf-a703-40ba075552a0"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Example: different models\n",
|
||||
"100 gray tops is kind of boring. It would be better if they were colorful and different sizes!\n",
|
||||
@@ -497,11 +474,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
|
||||
"metadata": {
|
||||
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Make 100 tippe tops with different colors and sizes\n",
|
||||
@@ -569,20 +544,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cf485c08-72be-4169-89b6-9d93df8ebbe3",
|
||||
"metadata": {
|
||||
"id": "cf485c08-72be-4169-89b6-9d93df8ebbe3"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Because the models are now different, the measurements of the gyro sensor are not consistent even though the initial state for each rollout was the same."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
|
||||
"metadata": {
|
||||
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(12, 8))\n",
|
||||
@@ -598,9 +569,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3841a669-6cd1-427e-a629-20a10a6e3a34",
|
||||
"metadata": {
|
||||
"id": "3841a669-6cd1-427e-a629-20a10a6e3a34"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Example: control inputs\n",
|
||||
"Open loop controls can be passed to `rollout` via the `control` argument. If passed, `nstep` no longer needs to be specified as it can be inferred from the size of `control`.\n",
|
||||
@@ -610,11 +579,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "2a184873-8d24-45da-b444-8d21f5dcd733",
|
||||
"metadata": {
|
||||
"id": "2a184873-8d24-45da-b444-8d21f5dcd733"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Episode parameters.\n",
|
||||
@@ -665,9 +632,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7",
|
||||
"metadata": {
|
||||
"id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"`rollout`'s `control_spec` argument can be used to indicate `control` contains values for actuators, generalized forces, cartesian forces, mocap poses, and/or the activation/deactivation of equality constraints. Internally, this is managed through [mj_setState](https://mujoco.readthedocs.io/en/stable/APIreference/APIfunctions.html#mj-setstate) and `control_spec` corresponds to `mj_setState`'s `spec` argument.\n",
|
||||
"\n",
|
||||
@@ -676,11 +641,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
|
||||
"metadata": {
|
||||
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"xfrc_size = mujoco.mj_stateSize(humanoid_model, mujoco.mjtState.mjSTATE_XFRC_APPLIED)\n",
|
||||
@@ -713,9 +676,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "0961c3ec-a691-4875-9a55-227a3d29c472",
|
||||
"metadata": {
|
||||
"id": "0961c3ec-a691-4875-9a55-227a3d29c472"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Advanced usage"
|
||||
]
|
||||
@@ -723,9 +684,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "VfYIyXWcLKfg",
|
||||
"metadata": {
|
||||
"id": "VfYIyXWcLKfg"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## skip_checks\n",
|
||||
"\n",
|
||||
@@ -747,11 +706,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
|
||||
"metadata": {
|
||||
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nbatch = 1000\n",
|
||||
@@ -795,9 +752,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "92627030-4726-4689-be8b-f1ba75905104",
|
||||
"metadata": {
|
||||
"id": "92627030-4726-4689-be8b-f1ba75905104"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As expected, as `nstep` increases, the benefits of using skip checks fades quickly. However, at low nstep and high batch sizes, it can make a significant difference.\n",
|
||||
"\n",
|
||||
@@ -807,9 +762,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c",
|
||||
"metadata": {
|
||||
"id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Reusing threadpools (`Rollout` class)\n",
|
||||
"\n",
|
||||
@@ -820,11 +773,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
|
||||
"metadata": {
|
||||
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nbatch = 100\n",
|
||||
@@ -862,9 +813,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671",
|
||||
"metadata": {
|
||||
"id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Reusing threadpools (`rollout` method)\n",
|
||||
"\n",
|
||||
@@ -875,11 +824,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "b6aa6801",
|
||||
"metadata": {
|
||||
"id": "b6aa6801"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nbatch = 1000\n",
|
||||
@@ -899,20 +846,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "144378d3",
|
||||
"metadata": {
|
||||
"id": "144378d3"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Second, if `rollout` reuses the same threadpool between calls, it is no longer safe to call `rollout` from multiple threads. For example the following is not allowed (the offending lines are commented out to avoid crashing the interpreter):"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "7f46a6d8",
|
||||
"metadata": {
|
||||
"id": "7f46a6d8"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"thread1 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n",
|
||||
@@ -928,9 +871,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "78c1f864-5238-4e27-a7ec-d03c45484d9a",
|
||||
"metadata": {
|
||||
"id": "78c1f864-5238-4e27-a7ec-d03c45484d9a"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## chunk_size\n",
|
||||
"\n",
|
||||
@@ -941,11 +882,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "a1be8f93",
|
||||
"metadata": {
|
||||
"id": "a1be8f93"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nbatch = 100\n",
|
||||
@@ -989,9 +928,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7",
|
||||
"metadata": {
|
||||
"id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Warmstarting\n",
|
||||
"\n",
|
||||
@@ -1004,11 +941,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
|
||||
"metadata": {
|
||||
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"top_model_cg = copy.copy(top_model)\n",
|
||||
@@ -1069,9 +1004,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7c2cf4fa",
|
||||
"metadata": {
|
||||
"id": "7c2cf4fa"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As expected, the middle animation (with warmstarting) matches the continuous rollout on the left. However, the model that did not use warmstarting diverged."
|
||||
]
|
||||
@@ -1079,9 +1012,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7944637f",
|
||||
"metadata": {
|
||||
"id": "7944637f"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Benchmarks\n",
|
||||
"\n",
|
||||
@@ -1098,11 +1029,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "cb6355dd",
|
||||
"metadata": {
|
||||
"id": "cb6355dd"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def python_rollout(model, data, nbatch, nstep):\n",
|
||||
@@ -1114,9 +1043,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6fe4a78b",
|
||||
"metadata": {
|
||||
"id": "6fe4a78b"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To run nbatch rollouts with `rollout`, we need to make an array of nbatch initial states to start the rollouts from.\n",
|
||||
"\n",
|
||||
@@ -1128,11 +1055,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "74f143e2",
|
||||
"metadata": {
|
||||
"id": "74f143e2"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def nthread_rollout(model, data, nbatch, nstep, nthread, rollout_):\n",
|
||||
@@ -1146,9 +1071,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b75dc44c",
|
||||
"metadata": {
|
||||
"id": "b75dc44c"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next, we benchmark the Python loop and `rollout` in both single threaded and multithreaded modes. The three benchmarks take about 2.5 minutes in total to run in total on an AMD 5800X3D."
|
||||
]
|
||||
@@ -1158,8 +1081,7 @@
|
||||
"execution_count": 0,
|
||||
"id": "0301e3ee",
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "0301e3ee"
|
||||
"cellView": "form"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1282,20 +1204,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "08fb0c12",
|
||||
"metadata": {
|
||||
"id": "08fb0c12"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Tippe Top Benchmark"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "f7e54830",
|
||||
"metadata": {
|
||||
"id": "f7e54830"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nominal_nbatch = 256 # Batch size to use when testing different nstep\n",
|
||||
@@ -1314,20 +1232,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "edefb26e",
|
||||
"metadata": {
|
||||
"id": "edefb26e"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Humanoid Benchmark"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "c9e58c6c",
|
||||
"metadata": {
|
||||
"id": "c9e58c6c"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nominal_nbatch = 256 # Batch size to use when testing different nstep\n",
|
||||
@@ -1346,20 +1260,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "468903bb",
|
||||
"metadata": {
|
||||
"id": "468903bb"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Humanoid100 Benchmark"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "83d775d4",
|
||||
"metadata": {
|
||||
"id": "83d775d4"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nominal_nbatch = 128 # Batch size to use when testing different nstep\n",
|
||||
@@ -1383,9 +1293,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d1133084",
|
||||
"metadata": {
|
||||
"id": "d1133084"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# MJX versus `rollout`"
|
||||
]
|
||||
@@ -1393,9 +1301,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c1638f2d",
|
||||
"metadata": {
|
||||
"id": "c1638f2d"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next we will benchmark `rollout` and MJX using the tippe top and humanoid models (humanoid100 is not supported by MJX).\n",
|
||||
"\n",
|
||||
@@ -1409,8 +1315,7 @@
|
||||
"execution_count": 0,
|
||||
"id": "7c86d157",
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "7c86d157"
|
||||
"cellView": "form"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1531,20 +1436,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a2dafd2e",
|
||||
"metadata": {
|
||||
"id": "a2dafd2e"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### MJX Tippe Top Benchmark"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "98c580b0",
|
||||
"metadata": {
|
||||
"id": "98c580b0"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nominal_nbatch = 16384 # Batch size to use when testing different nstep\n",
|
||||
@@ -1560,20 +1461,16 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "205da5da",
|
||||
"metadata": {
|
||||
"id": "205da5da"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### MJX Humanoid Benchmark"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "53166ae1",
|
||||
"metadata": {
|
||||
"id": "53166ae1"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"nominal_nbatch = 4096 # Batch size to use when testing different nstep\n",
|
||||
@@ -1589,9 +1486,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fb2caa72",
|
||||
"metadata": {
|
||||
"id": "fb2caa72"
|
||||
},
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### MJX Multiple Humanoids in one model\n",
|
||||
"\n",
|
||||
@@ -1599,16 +1494,14 @@
|
||||
"\n",
|
||||
"Here we will produce a similar plot to compare MJX and with `rollout`. On a 5800X3D and 4090 the benchmark takes about 16.5 minutes to run.\n",
|
||||
"\n",
|
||||
"**Note:** These results are not directly comparable since with the plot in the documentation because, in particular, the batch size was redued from 8192 to 4096 in order to fit the batch on a 4090."
|
||||
"**Note:** These results are not directly comparable since with the plot in the documentation because, in particular, the batch size was reduced from 8192 to 4096 in order to fit the batch on a 4090."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "3d6be608",
|
||||
"metadata": {
|
||||
"id": "3d6be608"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"max_humanoids = 10\n",
|
||||
@@ -1656,11 +1549,9 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"execution_count": null,
|
||||
"id": "b6c5fc2e",
|
||||
"metadata": {
|
||||
"id": "b6c5fc2e"
|
||||
},
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Plot MJX nhumanoid benchmark\n",
|
||||
|
||||
Reference in New Issue
Block a user