diff --git a/python/rollout.ipynb b/python/rollout.ipynb index 08a1a667..48bf4647 100644 --- a/python/rollout.ipynb +++ b/python/rollout.ipynb @@ -700,6 +700,388 @@ "print(f'Render time {end_render-start_render:.1f} seconds')" ] }, + { + "cell_type": "markdown", + "id": "0961c3ec-a691-4875-9a55-227a3d29c472", + "metadata": { + "id": "0961c3ec-a691-4875-9a55-227a3d29c472" + }, + "source": [ + "# Advanced usage" + ] + }, + { + "cell_type": "markdown", + "id": "VfYIyXWcLKfg", + "metadata": { + "id": "VfYIyXWcLKfg" + }, + "source": [ + "## skip_checks\n", + "\n", + "By default rollout performs many checks on the dimensions of its arguments. This it allows it to infer dimensions such as `nbatch` and `nstep`, tile arguments that were not fully specified, and allocate the returned `state` and `sensordata` arrays.\n", + "\n", + "However, these check take time, particularly if `state` and `sensordata` are large or if there are many models and `nstep` is low. So advanced users may want to use the `skip_checks=True` argument in order to achieve additional performance.\n", + "\n", + "If used, certain arguments become non-optional, and all signals must be fully defined (no implicit tiling). In particular:\n", + "* `model` must be a list of length `nbatch`\n", + "* `data` must be a list of length `nthread`\n", + "* `nstep` must be specified\n", + "* `initial_state` must be an array of shape `nbatch x nstate`\n", + "* `control` is optional, but if passed must be an array of shape `nbatch x nstep x ncontrol`\n", + "* `state` is optional, but must be passed if state is to be returned and must be of shape `nbatch x nstep x nstate`\n", + "* `sensordata` is optional, but must be passed if sensor data is to be returned and must be of shape `nbatch x nstep x nsensordata`\n", + "\n", + "As an extreme example, we pass 10,000 humanoid models to `rollout` and simulate 1 step each with and without checks." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95", + "metadata": { + "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95" + }, + "outputs": [], + "source": [ + "nbatch = 1000\n", + "nstep = [1, 10, 100, 500]\n", + "ntiming = 5\n", + "\n", + "top_data = init_top(top_model)\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", + "\n", + "# Note: state, sensordata array automatically allocated and return\n", + "def rollout_with_checks(nstep):\n", + " state, sensordata = rollout.rollout([top_model]*nbatch, top_datas, initial_state, nstep=nstep)\n", + "\n", + "# Note: state, sensordata arrays have to be preallocated\n", + "state = None\n", + "sensordata = None\n", + "def rollout_skip_checks(nstep):\n", + " # Note initial state must be tiled\n", + " rollout.rollout([top_model]*nbatch, top_datas, initial_state_tiled, nstep=nstep,\n", + " state=state, sensordata=sensordata, skip_checks=True)\n", + "\n", + "t_with_checks = benchmark(lambda x: rollout_with_checks(x), nstep, ntiming=ntiming)\n", + "t_skip_checks = benchmark(lambda x: rollout_skip_checks(x), nstep, ntiming=ntiming)\n", + "\n", + "steps_per_second = (nbatch * np.array(nstep)) / np.array(t_with_checks)\n", + "steps_per_second_skip_checks = (nbatch * np.array(nstep)) / np.array(t_skip_checks)\n", + "\n", + "plt.loglog(nstep, steps_per_second, label='with checks')\n", + "plt.loglog(nstep, steps_per_second_skip_checks, label='skip checks')\n", + "plt.ylabel('steps per second')\n", + "plt.xlabel('nstep')\n", + "plt.legend()\n", + "plt.grid(True, which=\"both\", axis=\"both\")" + ] + }, + { + "cell_type": "markdown", + "id": "92627030-4726-4689-be8b-f1ba75905104", + "metadata": { + "id": "92627030-4726-4689-be8b-f1ba75905104" + }, + "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", + "Notice that the version with checks can use the non-tiled `initial_state`, however the skip checks version must used the tiled version, `initial_state_tiled`." + ] + }, + { + "cell_type": "markdown", + "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c", + "metadata": { + "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c" + }, + "source": [ + "## Reusing threadpools (`Rollout` class)\n", + "\n", + "The `rollout` module provided the class `Rollout` in addition to the method `rollout`. The class `Rollout` is designed allow safe reuse of the internally managed thread pool.\n", + "\n", + "Reuse can speed things up considerably when rollouts are short. Let's find out how the speedup changes for the tippe top model by rolling it out with increasing numbers of steps." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49", + "metadata": { + "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49" + }, + "outputs": [], + "source": [ + "nbatch = 100\n", + "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_datas = [copy.copy(top_data) for _ in range(nthread)]\n", + "\n", + "def rollout_method(nstep):\n", + " for i in range(20):\n", + " rollout.rollout(top_model, top_datas, initial_states, nstep=nstep)\n", + "\n", + "def rollout_class(nstep):\n", + " with rollout.Rollout(nthread=nthread) as rollout_:\n", + " for i in range(20):\n", + " rollout_.rollout(top_model, top_datas, initial_states, nstep=nstep)\n", + "\n", + "t_method = benchmark(lambda x: rollout_method(x), nsteps, ntiming)\n", + "t_class = benchmark(lambda x: rollout_class(x), nsteps, ntiming)\n", + "\n", + "plt.loglog(nsteps, nbatch * np.array(nsteps) / t_method, label='recreating threadpools')\n", + "plt.loglog(nsteps, nbatch * np.array(nsteps) / t_class, label='reusing threadpool')\n", + "plt.xlabel('nstep')\n", + "plt.ylabel('steps per second')\n", + "plt.legend()\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671", + "metadata": { + "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671" + }, + "source": [ + "## Reusing threadpools (`rollout` method)" + ] + }, + { + "cell_type": "markdown", + "id": "0aba4dd6", + "metadata": { + "id": "0aba4dd6" + }, + "source": [ + "`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:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "id": "b6aa6801", + "metadata": { + "id": "b6aa6801" + }, + "outputs": [], + "source": [ + "nbatch = 1000\n", + "nstep = 1\n", + "\n", + "top_data = init_top(top_model)\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", + "\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.shutdown_persistent_pool() # Shutdown the pool manually when finished" + ] + }, + { + "cell_type": "markdown", + "id": "144378d3", + "metadata": { + "id": "144378d3" + }, + "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, + "id": "7f46a6d8", + "metadata": { + "id": "7f46a6d8" + }, + "outputs": [], + "source": [ + "thread1 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n", + "thread2 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n", + "\n", + "thread1.start()\n", + "#thread2.start() # Do not do this! rollout will be using the same persistent threadpool from two threads and may crash the interpreter\n", + "thread1.join()\n", + "#thread2.join()\n", + "rollout.shutdown_persistent_pool()" + ] + }, + { + "cell_type": "markdown", + "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a", + "metadata": { + "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a" + }, + "source": [ + "## chunk_size" + ] + }, + { + "cell_type": "markdown", + "id": "2573fff5", + "metadata": { + "id": "2573fff5" + }, + "source": [ + "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." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "id": "a1be8f93", + "metadata": { + "id": "a1be8f93" + }, + "outputs": [], + "source": [ + "nbatch = 100\n", + "nstep = 1\n", + "ntiming = 20\n", + "\n", + "# Load model\n", + "hopper_model = mujoco.MjModel.from_xml_path(hopper_path)\n", + "hopper_data = mujoco.MjData(hopper_model)\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", + "\n", + "def rollout_chunk_size(chunk_size=None):\n", + " rollout.rollout(hopper_model, hopper_datas, initial_states, nstep=nstep, chunk_size=chunk_size)\n", + "\n", + "# Rollout with different chunk sizes\n", + "default_chunk_size = int(max(1.0, 0.1 * nbatch / nthread))\n", + "chunk_sizes = sorted([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, default_chunk_size])\n", + "t_chunk_size = benchmark(lambda x: rollout_chunk_size(x), chunk_sizes, ntiming=ntiming)\n", + "\n", + "# Get optimal chunk size\n", + "steps_per_second = nbatch * nstep / t_chunk_size\n", + "default_index = [i for i, c in enumerate(chunk_sizes) if c == default_chunk_size][0]\n", + "optimal_index = np.argmax(steps_per_second)\n", + "plt.loglog(chunk_sizes, steps_per_second, color='b')\n", + "plt.plot(chunk_sizes[default_index], steps_per_second[default_index], marker='o', color='r', label='default chunk size')\n", + "plt.plot(chunk_sizes[optimal_index], steps_per_second[optimal_index], marker='o', color='g', label='optimal chunk size')\n", + "plt.ylabel('steps per second')\n", + "plt.xlabel('chunk size')\n", + "plt.legend()\n", + "plt.grid()\n", + "\n", + "print(f'default chunk size: {default_chunk_size} \\t steps per second: {steps_per_second[default_index]:0.1f}')\n", + "print(f'optimal chunk size: {chunk_sizes[optimal_index]} \\t steps per second: {steps_per_second[optimal_index]:0.1f}')" + ] + }, + { + "cell_type": "markdown", + "id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7", + "metadata": { + "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": [ + "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", + "\n", + "The simulation is run three times. Once with a 6000 step rollout, once with 100 chunks of 60 steps with warmstarting, and once more in 100 chunks of 60 steps without warmstarting." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb", + "metadata": { + "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb" + }, + "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", + "chunks = 100\n", + "steps_per_chunk = 60\n", + "nstep = steps_per_chunk*chunks\n", + "\n", + "# Get initial states\n", + "data = init_top(model)\n", + "initial_state = get_state(model, data)\n", + "\n", + "start = time.time()\n", + "# Rollout with nstep steps\n", + "state_all, _ = rollout.rollout(model, data, initial_state, nstep=nstep)\n", + "\n", + "# Rollout in chunks with warmstarting\n", + "state_chunks = []\n", + "state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n", + "state_chunks.append(state_chunk)\n", + "for _ in range(chunks-1):\n", + " state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk, initial_warmstart=data.qacc_warmstart)\n", + " state_chunks.append(state_chunk)\n", + "state_all_chunked_warmstart = np.concatenate(state_chunks, axis=1)\n", + "\n", + "# Rollout in chunks without warmstarting\n", + "state_chunks = []\n", + "state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n", + "state_chunks.append(state_chunk)\n", + "first_warmstart = None\n", + "for i in range(chunks-1):\n", + " state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk)\n", + " state_chunks.append(state_chunk)\n", + "state_all_chunked = np.concatenate(state_chunks, axis=1)\n", + "end = time.time()\n", + "\n", + "# Render the rollouts\n", + "start_render = time.time()\n", + "framerate = 60\n", + "state_render = np.concatenate((state_all, state_all_chunked, state_all_chunked_warmstart), axis=0)\n", + "camera = 'distant'\n", + "frames1 = render_many(model, data, state_all, framerate, shape=(240, 320), transparent=False, camera=camera)\n", + "frames2 = render_many(model, data, state_all_chunked_warmstart, framerate, shape=(240, 320), transparent=False, camera=camera)\n", + "frames3 = render_many(model, data, state_all_chunked, framerate, shape=(240, 320), transparent=False, camera=camera)\n", + "media.show_video(np.concatenate((frames1, frames2, frames3), axis=2))\n", + "end_render = time.time()\n", + "\n", + "print(f'Rollout took {end-start:.1f} seconds')\n", + "print(f'Rendering took {end_render-start_render:.1f} seconds')" + ] + }, + { + "cell_type": "markdown", + "id": "7c2cf4fa", + "metadata": { + "id": "7c2cf4fa" + }, + "source": [ + "As expected, the middle animation (with warmstarting) matches the continuous rollout on the left. However, the model that did not use warmstarting diverged." + ] + }, { "cell_type": "markdown", "id": "7944637f", @@ -711,16 +1093,8 @@ "\n", "The `rollout.rollout` function in the `mujoco` Python library runs batches of simulations for a fixed number steps. It can run in single or multi-threaded modes. The speedup over pure Python is significant because `rollout` can be easily configured to use multithreading.\n", "\n", - "To show the speedup, we will run benchmarks with the \"tippe top\", \"humanoid\", and \"humanoid100\" models." - ] - }, - { - "cell_type": "markdown", - "id": "7ef2bd0b", - "metadata": { - "id": "7ef2bd0b" - }, - "source": [ + "To show the speedup, we will run benchmarks with the \"tippe top\", \"humanoid\", and \"humanoid100\" models.\n", + "\n", "## Python rollouts versus `rollout`\n", "\n", "The benchmark runs the three models with varying batch and step counts.\n", @@ -792,11 +1166,12 @@ "execution_count": 0, "id": "0301e3ee", "metadata": { + "cellView": "form", "id": "0301e3ee" }, "outputs": [], "source": [ - "#@title Benchmarking and plotting code\n", + "#@title Benchmarking utilities\n", "\n", "top_model = mujoco.MjModel.from_xml_string(tippe_top)\n", "def init_top(model):\n", @@ -913,8 +1288,11 @@ "nbatch = [1, 10, 100, 500, 1000] # Batch sizes to benchmark\n", "nstep = sorted([1, 10, 100, 1000, 2000, 4000]) # Step counts to benchmark\n", "\n", - "top_benchmark_results = benchmark_rollout(top_model, top_data, init_top, nbatch, nstep, nominal_nbatch, nominal_nstep)\n", - "plot_benchmark(top_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)" + "top_benchmark_results = benchmark_rollout(top_model, top_data, init_top,\n", + " nbatch, nstep,\n", + " nominal_nbatch, nominal_nstep)\n", + "plot_benchmark(top_benchmark_results, nbatch, nstep,\n", + " nominal_nbatch, nominal_nstep)" ] }, { @@ -941,8 +1319,11 @@ "nbatch = [1, 10, 100, 200, 400] # Batch sizes to benchmark\n", "nstep = sorted([1, 10, 100, 500, 1000]) # Step counts to benchmark\n", "\n", - "humanoid_benchmark_results = benchmark_rollout(humanoid_model, humanoid_data, init_humanoid, nbatch, nstep, nominal_nbatch, nominal_nstep)\n", - "plot_benchmark(humanoid_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)" + "humanoid_benchmark_results = benchmark_rollout(humanoid_model, humanoid_data,\n", + " init_humanoid, nbatch, nstep,\n", + " nominal_nbatch, nominal_nstep)\n", + "plot_benchmark(humanoid_benchmark_results, nbatch, nstep,\n", + " nominal_nbatch, nominal_nstep)" ] }, { @@ -964,386 +1345,22 @@ }, "outputs": [], "source": [ - "nominal_nbatch = 100 # Batch size to use when testing different nstep\n", - "nominal_nstep = 200 # Step count to use when testing different nbatch\n", - "nbatch = [1, 10, 50, 100, 200] # Batch sizes to benchmark\n", - "nstep = sorted([1, 10, 100, 200, 400]) # Step counts to benchmark\n", + "nominal_nbatch = 100 # Batch size to use when testing different nstep\n", + "nominal_nstep = 200 # Step count to use when testing different nbatch\n", + "nbatch = [1, 10, 50, 100, 200] # Batch sizes to benchmark\n", + "nstep = sorted([1, 10, 100, 200, 400]) # Step counts to benchmark\n", "\n", - "humanoid100_benchmark_results = benchmark_rollout(humanoid100_model, humanoid100_data, init_humanoid100, nbatch, nstep, nominal_nbatch, nominal_nstep)\n", - "plot_benchmark(humanoid100_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)" - ] - }, - { - "cell_type": "markdown", - "id": "0961c3ec-a691-4875-9a55-227a3d29c472", - "metadata": { - "id": "0961c3ec-a691-4875-9a55-227a3d29c472" - }, - "source": [ - "# Advanced usage\n", - "## skip_checks=True\n", - "\n", - "By default rollout performs many checks on the dimensions of its arguments. This it allows it to infer dimensions such as `nbatch` and `nstep`, tile arguments that were not fully specified, and allocate the returned `state` and `sensordata` arrays.\n", - "\n", - "However, these check take time, particularly if `state` and `sensordata` are large or if there are many models and `nstep` is low. So advanced users may want to use the `skip_checks=True` argument in order to achieve additional performance.\n", - "\n", - "If used, certain arguments become non-optional, and all signals must be fully defined (no implicit tiling). In particular:\n", - "* `model` must be a list of length `nbatch`\n", - "* `data` must be a list of length `nthread`\n", - "* `nstep` must be specified\n", - "* `initial_state` must be an array of shape `nbatch x nstate`\n", - "* `control` is optional, but if passed must be an array of shape `nbatch x nstep x ncontrol`\n", - "* `state` is optional, but must be passed if state is to be returned and must be of shape `nbatch x nstep x nstate`\n", - "* `sensordata` is optional, but must be passed if sensor data is to be returned and must be of shape `nbatch x nstep x nsensordata`\n", - "\n", - "As an extreme example, we pass 10,000 humanoid models to `rollout` and simulate 1 step each with and without checks." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95", - "metadata": { - "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95" - }, - "outputs": [], - "source": [ - "nbatch = 1000\n", - "nstep = [1, 10, 100, 500]\n", - "ntiming = 5\n", - "\n", - "top_data = init_top(top_model)\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", - "\n", - "# Note: state, sensordata array automatically allocated and return\n", - "def rollout_with_checks(nstep):\n", - " state, sensordata = rollout.rollout([top_model]*nbatch, top_datas, initial_state, nstep=nstep)\n", - "\n", - "# Note: state, sensordata arrays have to be preallocated\n", - "state = None\n", - "sensordata = None\n", - "def rollout_skip_checks(nstep):\n", - " # Note initial state must be tiled\n", - " rollout.rollout([top_model]*nbatch, top_datas, initial_state_tiled, nstep=nstep,\n", - " state=state, sensordata=sensordata, skip_checks=True)\n", - "\n", - "t_with_checks = benchmark(lambda x: rollout_with_checks(x), nstep, ntiming=ntiming)\n", - "t_skip_checks = benchmark(lambda x: rollout_skip_checks(x), nstep, ntiming=ntiming)\n", - "\n", - "steps_per_second = (nbatch * np.array(nstep)) / np.array(t_with_checks)\n", - "steps_per_second_skip_checks = (nbatch * np.array(nstep)) / np.array(t_skip_checks)\n", - "\n", - "plt.loglog(nstep, steps_per_second, label='with checks')\n", - "plt.loglog(nstep, steps_per_second_skip_checks, label='skip checks')\n", - "plt.ylabel('steps per second')\n", - "plt.xlabel('nstep')\n", - "plt.legend()\n", - "plt.grid(True, which=\"both\", axis=\"both\")" - ] - }, - { - "cell_type": "markdown", - "id": "92627030-4726-4689-be8b-f1ba75905104", - "metadata": { - "id": "92627030-4726-4689-be8b-f1ba75905104" - }, - "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", - "Notice that the version with checks can use the non-tiled `initial_state`, however the skip checks version must used the tiled version, `initial_state_tiled`." - ] - }, - { - "cell_type": "markdown", - "id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7", - "metadata": { - "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": [ - "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", - "\n", - "The simulation is run three times. Once with a 6000 step rollout, once with 100 chunks of 60 steps with warmstarting, and once more in 100 chunks of 60 steps without warmstarting." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb", - "metadata": { - "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb" - }, - "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", - "chunks = 100\n", - "steps_per_chunk = 60\n", - "nstep = steps_per_chunk*chunks\n", - "\n", - "# Get initial states\n", - "data = init_top(model)\n", - "initial_state = get_state(model, data)\n", - "\n", - "start = time.time()\n", - "# Rollout with nstep steps\n", - "state_all, _ = rollout.rollout(model, data, initial_state, nstep=nstep)\n", - "\n", - "# Rollout in chunks with warmstarting\n", - "state_chunks = []\n", - "state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n", - "state_chunks.append(state_chunk)\n", - "for _ in range(chunks-1):\n", - " state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk, initial_warmstart=data.qacc_warmstart)\n", - " state_chunks.append(state_chunk)\n", - "state_all_chunked_warmstart = np.concatenate(state_chunks, axis=1)\n", - "\n", - "# Rollout in chunks without warmstarting\n", - "state_chunks = []\n", - "state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n", - "state_chunks.append(state_chunk)\n", - "first_warmstart = None\n", - "for i in range(chunks-1):\n", - " state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk)\n", - " state_chunks.append(state_chunk)\n", - "state_all_chunked = np.concatenate(state_chunks, axis=1)\n", - "end = time.time()\n", - "\n", - "# Render the rollouts\n", - "start_render = time.time()\n", - "framerate = 60\n", - "state_render = np.concatenate((state_all, state_all_chunked, state_all_chunked_warmstart), axis=0)\n", - "camera = 'distant'\n", - "frames1 = render_many(model, data, state_all, framerate, shape=(240, 320), transparent=False, camera=camera)\n", - "frames2 = render_many(model, data, state_all_chunked_warmstart, framerate, shape=(240, 320), transparent=False, camera=camera)\n", - "frames3 = render_many(model, data, state_all_chunked, framerate, shape=(240, 320), transparent=False, camera=camera)\n", - "media.show_video(np.concatenate((frames1, frames2, frames3), axis=2))\n", - "end_render = time.time()\n", - "\n", - "print(f'Rollout took {end-start:.1f} seconds')\n", - "print(f'Rendering took {end_render-start_render:.1f} seconds')" - ] - }, - { - "cell_type": "markdown", - "id": "7c2cf4fa", - "metadata": { - "id": "7c2cf4fa" - }, - "source": [ - "As expected, the middle animation (with warmstarting) matches the continuous rollout on the left. However, the model that did not use warmstarting diverged." - ] - }, - { - "cell_type": "markdown", - "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a", - "metadata": { - "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a" - }, - "source": [ - "## chunk_size" - ] - }, - { - "cell_type": "markdown", - "id": "2573fff5", - "metadata": { - "id": "2573fff5" - }, - "source": [ - "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." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "id": "a1be8f93", - "metadata": { - "id": "a1be8f93" - }, - "outputs": [], - "source": [ - "nbatch = 100\n", - "nstep = 1\n", - "ntiming = 20\n", - "\n", - "# Load model\n", - "hopper_model = mujoco.MjModel.from_xml_path(hopper_path)\n", - "hopper_data = mujoco.MjData(hopper_model)\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", - "\n", - "def rollout_chunk_size(chunk_size=None):\n", - " rollout.rollout(hopper_model, hopper_datas, initial_states, nstep=nstep, chunk_size=chunk_size)\n", - "\n", - "# Rollout with different chunk sizes\n", - "default_chunk_size = int(max(1.0, 0.1 * nbatch / nthread))\n", - "chunk_sizes = sorted([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, default_chunk_size])\n", - "t_chunk_size = benchmark(lambda x: rollout_chunk_size(x), chunk_sizes, ntiming=ntiming)\n", - "\n", - "# Get optimal chunk size\n", - "steps_per_second = nbatch * nstep / t_chunk_size\n", - "default_index = [i for i, c in enumerate(chunk_sizes) if c == default_chunk_size][0]\n", - "optimal_index = np.argmax(steps_per_second)\n", - "plt.loglog(chunk_sizes, steps_per_second, color='b')\n", - "plt.plot(chunk_sizes[default_index], steps_per_second[default_index], marker='o', color='r', label='default chunk size')\n", - "plt.plot(chunk_sizes[optimal_index], steps_per_second[optimal_index], marker='o', color='g', label='optimal chunk size')\n", - "plt.ylabel('steps per second')\n", - "plt.xlabel('chunk size')\n", - "plt.legend()\n", - "plt.grid()\n", - "\n", - "print(f'default chunk size: {default_chunk_size} \\t steps per second: {steps_per_second[default_index]:0.1f}')\n", - "print(f'optimal chunk size: {chunk_sizes[optimal_index]} \\t steps per second: {steps_per_second[optimal_index]:0.1f}')" - ] - }, - { - "cell_type": "markdown", - "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c", - "metadata": { - "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c" - }, - "source": [ - "## Reusing threadpools with the class `Rollout`\n", - "\n", - "The `rollout` module provided the class `Rollout` in addition to the method `rollout`. The class `Rollout` is designed allow safe reuse of the internally managed thread pool.\n", - "\n", - "Reuse can speed things up considerably when rollouts are short. Let's find out how the speedup changes for the tippe top model by rolling it out with increasing numbers of steps." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49", - "metadata": { - "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49" - }, - "outputs": [], - "source": [ - "nbatch = 100\n", - "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_datas = [copy.copy(top_data) for _ in range(nthread)]\n", - "\n", - "def rollout_method(nstep):\n", - " for i in range(20):\n", - " rollout.rollout(top_model, top_datas, initial_states, nstep=nstep)\n", - "\n", - "def rollout_class(nstep):\n", - " with rollout.Rollout(nthread=nthread) as rollout_:\n", - " for i in range(20):\n", - " rollout_.rollout(top_model, top_datas, initial_states, nstep=nstep)\n", - "\n", - "t_method = benchmark(lambda x: rollout_method(x), nsteps, ntiming)\n", - "t_class = benchmark(lambda x: rollout_class(x), nsteps, ntiming)\n", - "\n", - "plt.loglog(nsteps, nbatch * np.array(nsteps) / t_method, label='recreating threadpools')\n", - "plt.loglog(nsteps, nbatch * np.array(nsteps) / t_class, label='reusing threadpool')\n", - "plt.xlabel('nstep')\n", - "plt.ylabel('steps per second')\n", - "plt.legend()\n", - "plt.grid()" - ] - }, - { - "cell_type": "markdown", - "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671", - "metadata": { - "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671" - }, - "source": [ - "## Reusing threadpools with the method `rollout`" - ] - }, - { - "cell_type": "markdown", - "id": "0aba4dd6", - "metadata": { - "id": "0aba4dd6" - }, - "source": [ - "`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:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "id": "b6aa6801", - "metadata": { - "id": "b6aa6801" - }, - "outputs": [], - "source": [ - "nbatch = 1000\n", - "nstep = 1\n", - "\n", - "top_data = init_top(top_model)\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", - "\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.shutdown_persistent_pool() # Shutdown the pool manually when finished" - ] - }, - { - "cell_type": "markdown", - "id": "144378d3", - "metadata": { - "id": "144378d3" - }, - "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, - "id": "7f46a6d8", - "metadata": { - "id": "7f46a6d8" - }, - "outputs": [], - "source": [ - "thread1 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n", - "thread2 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n", - "\n", - "thread1.start()\n", - "#thread2.start() # Do not do this! rollout will be using the same persistent threadpool from two threads and may crash the interpreter\n", - "thread1.join()\n", - "#thread2.join()\n", - "rollout.shutdown_persistent_pool()" + "humanoid100_benchmark_results = benchmark_rollout(\n", + " humanoid100_model,\n", + " humanoid100_data,\n", + " init_humanoid100,\n", + " nbatch,\n", + " nstep,\n", + " nominal_nbatch,\n", + " nominal_nstep,\n", + ")\n", + "plot_benchmark(humanoid100_benchmark_results, nbatch, nstep,\n", + " nominal_nbatch, nominal_nstep)" ] }, { @@ -1539,7 +1556,7 @@ "id": "fb2caa72" }, "source": [ - "### MJX Multiple Humanoids in One Model\n", + "### MJX Multiple Humanoids in one model\n", "\n", "The MJX [documentation](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits) contains a chart comparing the speed of native MuJoCo vs MJX on a variety of devices.\n", "\n", @@ -1566,7 +1583,9 @@ "t_mjx = []\n", "for i in range(1, max_humanoids+1):\n", " print(f'Running benchmark on {i} humanoids')\n", - " model = mujoco.MjModel.from_xml_path(f'mujoco/mjx/mujoco/mjx/test_data/humanoid/{i:02d}_humanoids.xml')\n", + " model = mujoco.MjModel.from_xml_path(\n", + " f'mujoco/mjx/mujoco/mjx/test_data/humanoid/{i:02d}_humanoids.xml'\n", + " )\n", " data = mujoco.MjData(model)\n", "\n", " mjx_model = mjx.put_model(model)\n", @@ -1574,7 +1593,10 @@ " batch = jax.vmap(lambda x: mjx_data)(jp.array(list(range(nbatch))))\n", "\n", " start = time.perf_counter()\n", - " rollout.rollout(model, [copy.copy(data) for _ in range(nthread)], initial_state=get_state(model, data, nbatch), nstep=humanoid_nstep)\n", + " datas = [copy.copy(data) for _ in range(nthread)]\n", + " initial_state = get_state(model, data, nbatch)\n", + " rollout.rollout(model, datas, initial_state=initial_state,\n", + " nstep=humanoid_nstep)\n", " end = time.perf_counter()\n", " t_rollout.append(end-start)\n", "\n", @@ -1633,13 +1655,15 @@ "accelerator": "GPU", "colab": { "collapsed_sections": [ - "7944637f", - "0961c3ec-a691-4875-9a55-227a3d29c472", - "d1133084", + "VfYIyXWcLKfg", + "d32a77b5-24bd-4d17-80ac-15cc4d03731c", + "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671", + "78c1f864-5238-4e27-a7ec-d03c45484d9a", + "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7", "a2dafd2e", "205da5da" ], - "gpuType": "T4", + "gpuType": "A100", "private_outputs": true, "toc_visible": true },