From a0eb98553216375c8e5b6bb27b0a1fa9086d427d Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 21 Feb 2025 10:04:26 -0800 Subject: [PATCH] rollout notebook: fix bugs due to cell reordering PiperOrigin-RevId: 729567659 Change-Id: I439da667317ed51370d7c8a9fbfc31d11140752c --- python/rollout.ipynb | 88 +++++++++++++++++--------------------------- 1 file changed, 33 insertions(+), 55 deletions(-) diff --git a/python/rollout.ipynb b/python/rollout.ipynb index 48bf4647..e87a5067 100644 --- a/python/rollout.ipynb +++ b/python/rollout.ipynb @@ -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",