rollout notebook: fix bugs due to cell reordering
PiperOrigin-RevId: 729567659 Change-Id: I439da667317ed51370d7c8a9fbfc31d11140752c
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
1532cffdf0
commit
a0eb985532
+33
-55
@@ -139,9 +139,6 @@
|
||||
"# More legible printing from numpy.\n",
|
||||
"np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
|
||||
"\n",
|
||||
"from IPython.display import clear_output\n",
|
||||
"clear_output()\n",
|
||||
"\n",
|
||||
"# Set the number of threads to the number of cpu's that the multiprocessing module reports\n",
|
||||
"nthread = cpu_count()\n",
|
||||
"\n",
|
||||
@@ -152,7 +149,11 @@
|
||||
"humanoid100_path = 'mujoco/model/humanoid/humanoid100.xml'\n",
|
||||
"print('Getting hopper XML description from GitHub:')\n",
|
||||
"!git clone https://github.com/google-deepmind/dm_control\n",
|
||||
"hopper_path ='dm_control/dm_control/suite/hopper.xml'"
|
||||
"hopper_path ='dm_control/dm_control/suite/hopper.xml'\n",
|
||||
"\n",
|
||||
"# clear installation printouts\n",
|
||||
"from IPython.display import clear_output\n",
|
||||
"clear_output()"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -374,7 +375,7 @@
|
||||
"id": "55d171f7-541b-4441-aa18-da86d6716410"
|
||||
},
|
||||
"source": [
|
||||
"## Detailed Usage\n",
|
||||
"## Usage\n",
|
||||
"\n",
|
||||
"It is helpful to read `rollout`'s docstring before beginning. The main takeaways are that `rollout` runs `nbatch` rollouts for `nstep` steps. Each `MjModel` can be different but should be the same up to parameter values. Passing multiple `MjData` enables multithreading, one thread per `MjData`.\n",
|
||||
"Further documentation can be found [here](https://mujoco.readthedocs.io/en/latest/python.html#rollout).\n",
|
||||
@@ -748,10 +749,11 @@
|
||||
"nstep = [1, 10, 100, 500]\n",
|
||||
"ntiming = 5\n",
|
||||
"\n",
|
||||
"top_data = init_top(top_model)\n",
|
||||
"top_data = mujoco.MjData(top_model)\n",
|
||||
"mujoco.mj_resetDataKeyframe(top_model, top_data, 0)\n",
|
||||
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
|
||||
"initial_state = get_state(top_model, top_data)\n",
|
||||
"initial_state_tiled = np.tile(initial_state, (nbatch, 1))\n",
|
||||
"initial_state_tiled = get_state(top_model, top_data, nbatch)\n",
|
||||
"\n",
|
||||
"# Note: state, sensordata array automatically allocated and return\n",
|
||||
"def rollout_with_checks(nstep):\n",
|
||||
@@ -818,13 +820,12 @@
|
||||
"nsteps = [2**i for i in [2, 3, 4, 5, 6, 7]]\n",
|
||||
"ntiming = 5\n",
|
||||
"\n",
|
||||
"top_data = init_top(top_model)\n",
|
||||
"\n",
|
||||
"initial_state = get_state(top_model, top_data)\n",
|
||||
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
|
||||
"\n",
|
||||
"top_data = mujoco.MjData(top_model)\n",
|
||||
"mujoco.mj_resetDataKeyframe(top_model, top_data, 0)\n",
|
||||
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
|
||||
"\n",
|
||||
"initial_states = get_state(top_model, top_data, nbatch)\n",
|
||||
"\n",
|
||||
"def rollout_method(nstep):\n",
|
||||
" for i in range(20):\n",
|
||||
" rollout.rollout(top_model, top_datas, initial_states, nstep=nstep)\n",
|
||||
@@ -852,16 +853,8 @@
|
||||
"id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671"
|
||||
},
|
||||
"source": [
|
||||
"## Reusing threadpools (`rollout` method)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "0aba4dd6",
|
||||
"metadata": {
|
||||
"id": "0aba4dd6"
|
||||
},
|
||||
"source": [
|
||||
"## Reusing threadpools (`rollout` method)\n",
|
||||
"\n",
|
||||
"`rollout` will create and reuse a persistent threadpool by passing `persistent_pool=True`. However there are some caveats.\n",
|
||||
"\n",
|
||||
"First, because `rollout` is a function and does not know when the user is done calling it, the threadpool pool needs to be shutdown manually like this:"
|
||||
@@ -879,14 +872,14 @@
|
||||
"nbatch = 1000\n",
|
||||
"nstep = 1\n",
|
||||
"\n",
|
||||
"top_data = init_top(top_model)\n",
|
||||
"top_data = mujoco.MjData(top_model)\n",
|
||||
"mujoco.mj_resetDataKeyframe(top_model, top_data, 0)\n",
|
||||
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
|
||||
"\n",
|
||||
"initial_state = get_state(top_model, top_data)\n",
|
||||
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
|
||||
"initial_states = get_state(top_model, top_data, nbatch)\n",
|
||||
"\n",
|
||||
"rollout.rollout(model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Creates a pool\n",
|
||||
"rollout.rollout(model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Reuses the previously created pool\n",
|
||||
"rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Creates a pool\n",
|
||||
"rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Reuses the previously created pool\n",
|
||||
"rollout.shutdown_persistent_pool() # Shutdown the pool manually when finished"
|
||||
]
|
||||
},
|
||||
@@ -926,16 +919,8 @@
|
||||
"id": "78c1f864-5238-4e27-a7ec-d03c45484d9a"
|
||||
},
|
||||
"source": [
|
||||
"## chunk_size"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2573fff5",
|
||||
"metadata": {
|
||||
"id": "2573fff5"
|
||||
},
|
||||
"source": [
|
||||
"## chunk_size\n",
|
||||
"\n",
|
||||
"To minimize communication overhead, `rollout` distributes rollouts to threads in groups of rollouts called chunks. By default, `max(1, 0.1 * (nbatch / nthread))` rollouts are assigned to each chunk. While this chunking rule works well for most workloads it is not always optimal, especially when doing short rollouts with small models.\n",
|
||||
"\n",
|
||||
"Below we plot the steps per second versus chunk_size when running 1000 hoppers for 1 step each. In his case, the default chunk_size turns out to be quite a bit slower than using an increased chunk size."
|
||||
@@ -957,12 +942,10 @@
|
||||
"# Load model\n",
|
||||
"hopper_model = mujoco.MjModel.from_xml_path(hopper_path)\n",
|
||||
"hopper_data = mujoco.MjData(hopper_model)\n",
|
||||
"hopper_datas = [copy.copy(hopper_data) for _ in range(nthread)]\n",
|
||||
"\n",
|
||||
"# Get initial states\n",
|
||||
"initial_state = get_state(hopper_model, hopper_data)\n",
|
||||
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
|
||||
"\n",
|
||||
"hopper_datas = [copy.copy(hopper_data) for _ in range(nthread)]\n",
|
||||
"initial_states = get_state(hopper_model, hopper_data, nbatch)\n",
|
||||
"\n",
|
||||
"def rollout_chunk_size(chunk_size=None):\n",
|
||||
" rollout.rollout(hopper_model, hopper_datas, initial_states, nstep=nstep, chunk_size=chunk_size)\n",
|
||||
@@ -995,16 +978,8 @@
|
||||
"id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7"
|
||||
},
|
||||
"source": [
|
||||
"## Warmstarting"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a113380b-5cde-4eff-a235-c3a333910047",
|
||||
"metadata": {
|
||||
"id": "a113380b-5cde-4eff-a235-c3a333910047"
|
||||
},
|
||||
"source": [
|
||||
"## Warmstarting\n",
|
||||
"\n",
|
||||
"The `initial_warmstart` parameter can be used to warmstart the constraint solver as described in the [computation chapter](https://mujoco.readthedocs.io/en/stable/computation/index.html#warmstart-acceleration) of the documentation. This can be useful when rolling out models in chunks of steps. Without warmstarting, chaotic systems involving multi-body contact may diverge.\n",
|
||||
"\n",
|
||||
"Below we demonstrate this with the tippe top model where the contact solver was changed to CG. This makes the contact force calculation a less repeatable than if the default, Newton's method, were used and allows demonstrating the benefits of warmstarting.\n",
|
||||
@@ -1022,15 +997,18 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"model = copy.copy(top_model)\n",
|
||||
"model.opt.solver = mujoco.mjtSolver.mjSOL_CG # Change to CG solver\n",
|
||||
"data = init_top(model)\n",
|
||||
"\n",
|
||||
"# Change to CG solver, the Newton solver converges too well for\n",
|
||||
"# warmstarting to have an appreciable effect\n",
|
||||
"model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
|
||||
"\n",
|
||||
"chunks = 100\n",
|
||||
"steps_per_chunk = 60\n",
|
||||
"nstep = steps_per_chunk*chunks\n",
|
||||
"\n",
|
||||
"# Get initial states\n",
|
||||
"data = init_top(model)\n",
|
||||
"data = mujoco.MjData(top_model)\n",
|
||||
"mujoco.mj_resetDataKeyframe(top_model, data, 0)\n",
|
||||
"initial_state = get_state(model, data)\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
@@ -1655,7 +1633,7 @@
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [
|
||||
"VfYIyXWcLKfg",
|
||||
"fc69d0f4",
|
||||
"d32a77b5-24bd-4d17-80ac-15cc4d03731c",
|
||||
"9b3e14a1-71f3-430d-a3c2-1aadcf6c2671",
|
||||
"78c1f864-5238-4e27-a7ec-d03c45484d9a",
|
||||
|
||||
Reference in New Issue
Block a user