From cfc59164ebb3547e1246affd82a2f683721e1dd9 Mon Sep 17 00:00:00 2001
From: Levi Burner <leviburner@gmail.com>
Date: Tue, 25 Feb 2025 03:25:51 -0500
Subject: [PATCH] rollout notebook: fixups to benchmarks, plots, text

1. MJX benchmarks use jax.lax.scan and block_until_ready
2. Benchmark parameters updated to match RL use case
3. humanoid benchmarks are warmstarted by 2 seconds, humanoid100 by 4 seconds
4. benchmarks use skip_checks and preallocated threadpools
5. Improved advanced usage plots
6. Benchmark no longer includes model initialization in timed section
7. Update notebook text as needed
---
 python/rollout.ipynb | 382 +++++++++++++++++++++++++------------------
 1 file changed, 223 insertions(+), 159 deletions(-)

diff --git a/python/rollout.ipynb b/python/rollout.ipynb
index e87a5067..fb559c3d 100644
--- a/python/rollout.ipynb
+++ b/python/rollout.ipynb
@@ -17,9 +17,9 @@
     "\n",
     "This notebook describes the `rollout` module included in the MuJoCo Python library. It performs simulation \"rollouts\" with an underlying C++ function. The rollouts can be multithreaded.\n",
     "\n",
-    "Below, the usage of each argument is explained with examples. An example of using `rollout` with minimize is also given. Then `rollout` is benchmarked against pure python and MJX. Finally, some examples for advanced use cases are provided.\n",
+    "Below, the usage of each argument is explained with examples. Then some examples for advanced use cases are provided. Finally, `rollout` is benchmarked against pure python and MJX. \n",
     "\n",
-    "Note the benchmarks were designed to run on a AMD 5800X3D and an RTX 4090. They do not run in a reasonable amount of time on a typical free colab runtime.\n",
+    "Note the benchmarks were designed to run on >16 thread CPU and an RTX 4090 or A100. They do not run in a reasonable amount of time on a typical free colab runtime.\n",
     "\n",
     "<!-- Copyright 2025 DeepMind Technologies Limited\n",
     "\n",
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
    "metadata": {
     "editable": true,
@@ -100,7 +100,6 @@
     "try:\n",
     "  print('Checking that the installation succeeded:')\n",
     "  import mujoco\n",
-    "  from mujoco import minimize\n",
     "  from mujoco import rollout\n",
     "  from mujoco import mjx\n",
     "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
@@ -123,7 +122,6 @@
     "import time\n",
     "from multiprocessing import cpu_count\n",
     "import threading\n",
-    "import itertools\n",
     "import numpy as np\n",
     "import jax\n",
     "import jax.numpy as jp\n",
@@ -168,7 +166,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "082482c7",
    "metadata": {
     "editable": true,
@@ -195,14 +193,23 @@
     "  x, y = np.meshgrid(rows, cols)\n",
     "  return np.stack((x.flatten(), y.flatten())).T\n",
     "\n",
-    "def benchmark(f, x_list=[None], ntiming=1):\n",
+    "def benchmark(f, x_list=[None], ntiming=1, f_init=None):\n",
     "  x_times_list = []\n",
     "  for x in x_list:\n",
-    "    times = [time.perf_counter()]\n",
+    "    times = []\n",
     "    for i in range(ntiming):\n",
-    "      f(x)\n",
-    "      times.append(time.perf_counter())\n",
-    "    x_times_list.append(np.mean(np.diff(times)))\n",
+    "      if f_init is not None:\n",
+    "        x_init = f_init(x)\n",
+    "\n",
+    "      start = time.perf_counter()\n",
+    "      if f_init is not None:\n",
+    "        f(x, x_init)\n",
+    "      else:\n",
+    "        f(x)\n",
+    "      end = time.perf_counter()\n",
+    "      times.append(end - start)\n",
+    "\n",
+    "    x_times_list.append(np.mean(times))\n",
     "  return np.array(x_times_list)\n",
     "\n",
     "def render_many(model, data, state, framerate, camera=-1, shape=(480, 640),\n",
@@ -281,7 +288,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "849b93e5",
    "metadata": {
     "id": "849b93e5"
@@ -385,7 +392,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
    "metadata": {
     "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5"
@@ -404,12 +411,14 @@
    },
    "source": [
     "### Example: different initial states\n",
-    "`rollout` is designed to run `nbatch` rollouts in parallel for `nstep` steps. Let's simulate 100 tippe tops with different initial rotation speeds."
+    "`rollout` is designed to run `nbatch` rollouts in parallel for `nstep` steps. Let's simulate 100 tippe tops with different initial rotation speeds.\n",
+    "\n",
+    "**Note:** Using multithreading with rollout is enabled by passing one MjData per thread, as is done below."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
    "metadata": {
     "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe"
@@ -426,7 +435,7 @@
     "\n",
     "# Run the rollout\n",
     "start = time.time()\n",
-    "top_datas = [copy.copy(top_data) for _ in range(nthread)] # 1 MjData per rollout\n",
+    "top_datas = [copy.copy(top_data) for _ in range(nthread)] # 1 MjData per thread\n",
     "state, sensordata = rollout.rollout(top_model, top_datas, initial_states,\n",
     "                                    nstep=int(top_nstep*1.5))\n",
     "end = time.time()\n",
@@ -454,7 +463,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "957b8566-da31-410b-b385-e78241c5247a",
    "metadata": {
     "id": "957b8566-da31-410b-b385-e78241c5247a"
@@ -488,7 +497,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
    "metadata": {
     "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2"
@@ -569,7 +578,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
    "metadata": {
     "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03"
@@ -601,7 +610,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "2a184873-8d24-45da-b444-8d21f5dcd733",
    "metadata": {
     "id": "2a184873-8d24-45da-b444-8d21f5dcd733"
@@ -609,7 +618,7 @@
    "outputs": [],
    "source": [
     "# Episode parameters.\n",
-    "duration = 3       # (seconds)\n",
+    "duration = 3        # (seconds)\n",
     "framerate = 120     # (Hz)\n",
     "\n",
     "# Generate 100 different control sequences\n",
@@ -667,7 +676,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
    "metadata": {
     "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5"
@@ -738,7 +747,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
    "metadata": {
     "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95"
@@ -777,6 +786,8 @@
     "plt.loglog(nstep, steps_per_second_skip_checks, label='skip checks')\n",
     "plt.ylabel('steps per second')\n",
     "plt.xlabel('nstep')\n",
+    "ticker = matplotlib.ticker.FuncFormatter(lambda x, p: format(int(x), ','))\n",
+    "plt.gca().yaxis.set_minor_formatter(ticker)\n",
     "plt.legend()\n",
     "plt.grid(True, which=\"both\", axis=\"both\")"
    ]
@@ -809,7 +820,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
    "metadata": {
     "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49"
@@ -842,8 +853,10 @@
     "plt.loglog(nsteps, nbatch * np.array(nsteps) / t_class, label='reusing threadpool')\n",
     "plt.xlabel('nstep')\n",
     "plt.ylabel('steps per second')\n",
+    "ticker = matplotlib.ticker.FuncFormatter(lambda x, p: format(int(x), ','))\n",
+    "plt.gca().yaxis.set_minor_formatter(ticker)\n",
     "plt.legend()\n",
-    "plt.grid()"
+    "plt.grid(True, which=\"both\", axis=\"both\")"
    ]
   },
   {
@@ -862,7 +875,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "b6aa6801",
    "metadata": {
     "id": "b6aa6801"
@@ -895,7 +908,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "7f46a6d8",
    "metadata": {
     "id": "7f46a6d8"
@@ -928,7 +941,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "a1be8f93",
    "metadata": {
     "id": "a1be8f93"
@@ -964,8 +977,10 @@
     "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",
+    "ticker = matplotlib.ticker.FuncFormatter(lambda x, p: format(int(x), ','))\n",
+    "plt.gca().yaxis.set_minor_formatter(ticker)\n",
     "plt.legend()\n",
-    "plt.grid()\n",
+    "plt.grid(True, which=\"both\", axis=\"both\")\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}')"
@@ -989,48 +1004,49 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
    "metadata": {
     "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb"
    },
    "outputs": [],
    "source": [
-    "model = copy.copy(top_model)\n",
+    "top_model_cg = copy.copy(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",
+    "top_model_cg.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 = mujoco.MjData(top_model)\n",
-    "mujoco.mj_resetDataKeyframe(top_model, data, 0)\n",
-    "initial_state = get_state(model, data)\n",
+    "top_data_cg = mujoco.MjData(top_model_cg)\n",
+    "mujoco.mj_resetDataKeyframe(top_model_cg, top_data_cg, 0)\n",
+    "initial_state = get_state(top_model_cg, top_data_cg)\n",
     "\n",
     "start = time.time()\n",
     "# Rollout with nstep steps\n",
-    "state_all, _  = rollout.rollout(model, data, initial_state, nstep=nstep)\n",
+    "state_all, _  = rollout.rollout(top_model_cg, top_data_cg, 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_chunk, _ = rollout.rollout(top_model_cg, top_data_cg, 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_chunk, _ = rollout.rollout(top_model_cg, top_data_cg, state_chunks[-1][0, -1, :],\n",
+    "                                   nstep=steps_per_chunk, initial_warmstart=top_data_cg.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_chunk, _ = rollout.rollout(top_model_cg, top_data_cg, 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_chunk, _ = rollout.rollout(top_model_cg, top_data_cg, 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",
@@ -1040,9 +1056,9 @@
     "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",
+    "frames1 = render_many(top_model_cg, top_data_cg, state_all, framerate, shape=(240, 320), transparent=False, camera=camera)\n",
+    "frames2 = render_many(top_model_cg, top_data_cg, state_all_chunked_warmstart, framerate, shape=(240, 320), transparent=False, camera=camera)\n",
+    "frames3 = render_many(top_model_cg, top_data_cg, 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",
@@ -1082,16 +1098,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "cb6355dd",
    "metadata": {
     "id": "cb6355dd"
    },
    "outputs": [],
    "source": [
-    "def python_rollout(model, init_model, nbatch, nstep):\n",
+    "def python_rollout(model, data, nbatch, nstep):\n",
     "  for i in range(nbatch):\n",
-    "    data = init_model(model)\n",
     "    for i in range(nstep):\n",
     "      mujoco.mj_step(model, data)"
    ]
@@ -1113,20 +1128,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "74f143e2",
    "metadata": {
     "id": "74f143e2"
    },
    "outputs": [],
    "source": [
-    "def nthread_rollout(model, init_model, nbatch, nstep, nthread):\n",
-    "  # Initialize the MjData for the given model using the provided initializer\n",
-    "  data = init_model(model)\n",
-    "  rollout.rollout(model,\n",
-    "                  [copy.copy(data) for _ in range(nthread)], # Create one MjData per thread\n",
-    "                  np.tile(get_state(model, data), (nbatch, 1)), # Tile the initial condition nbatch times\n",
-    "                  nstep=nstep)"
+    "def nthread_rollout(model, data, nbatch, nstep, nthread, rollout_):\n",
+    "  rollout_.rollout([model]*nbatch,\n",
+    "                   [copy.copy(data) for _ in range(nthread)], # Create one MjData per thread\n",
+    "                   np.tile(get_state(model, data), (nbatch, 1)), # Tile the initial condition nbatch times\n",
+    "                   nstep=nstep,\n",
+    "                   skip_checks=True)"
    ]
   },
   {
@@ -1141,7 +1155,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "0301e3ee",
    "metadata": {
     "cellView": "form",
@@ -1157,49 +1171,69 @@
     "  # Set to the state to a spinning top (keyframe 0)\n",
     "  mujoco.mj_resetDataKeyframe(model, data, 0)\n",
     "  return data\n",
-    "top_data = init_top(top_model)\n",
     "\n",
     "# Create and initialize humanoid model\n",
+    "# Step for 2 seconds to get a stable set of contacts to benchmark\n",
     "humanoid_model = mujoco.MjModel.from_xml_path(humanoid_path)\n",
+    "humanoid_data = mujoco.MjData(humanoid_model)\n",
+    "humanoid_data.qvel[2] = 4 # Make the humanoid jump\n",
+    "while humanoid_data.time < 2.0:\n",
+    "  mujoco.mj_step(humanoid_model, humanoid_data)\n",
+    "humanoid_initial_state = get_state(humanoid_model, humanoid_data)\n",
     "def init_humanoid(model):\n",
     "  data = mujoco.MjData(model)\n",
-    "  data.qvel[2] = 4 # Make the humanoid jump\n",
+    "  mujoco.mj_setState(model, data, humanoid_initial_state.flatten(),\n",
+    "                     mujoco.mjtState.mjSTATE_FULLPHYSICS)\n",
     "  return data\n",
-    "humanoid_data = init_humanoid(humanoid_model)\n",
     "\n",
     "# Create and initialize humanoid100 model\n",
+    "# Step for 4 seconds to get a stable set of contacts to benchmark\n",
     "humanoid100_model = mujoco.MjModel.from_xml_path(humanoid100_path)\n",
+    "humanoid100_data = mujoco.MjData(humanoid100_model)\n",
+    "while humanoid100_data.time < 4.0:\n",
+    "  mujoco.mj_step(humanoid100_model, humanoid100_data)\n",
+    "humanoid100_initial_state = get_state(humanoid100_model, humanoid100_data)\n",
     "def init_humanoid100(model):\n",
     "  data = mujoco.MjData(model)\n",
+    "  mujoco.mj_setState(model, data, humanoid100_initial_state.flatten(),\n",
+    "                     mujoco.mjtState.mjSTATE_FULLPHYSICS)\n",
     "  return data\n",
-    "humanoid100_data = init_humanoid100(humanoid100_model)\n",
     "\n",
-    "def benchmark_rollout(model, data, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1):\n",
+    "def benchmark_rollout(model, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1):\n",
     "  print('Benchmarking pure python', end='\\r')\n",
     "  start = time.time()\n",
-    "  t_python_nbatch = benchmark(lambda x: python_rollout(model, init_model, x,  nominal_nstep), nbatch, ntiming)\n",
-    "  t_python_nstep  = benchmark(lambda x: python_rollout(model, init_model, nominal_nbatch, x), nstep,  ntiming)\n",
+    "  t_python_nbatch = benchmark(lambda x, data: python_rollout(model, data, x, nominal_nstep), nbatch, ntiming,\n",
+    "                              f_init=lambda x: init_model(model))\n",
+    "  t_python_nstep  = benchmark(lambda x, data: python_rollout(model, data, nominal_nbatch, x), nstep,  ntiming,\n",
+    "                              f_init=lambda x: init_model(model))\n",
     "  end = time.time()\n",
     "  print(f'Benchmarking pure python took {end-start:0.1f} seconds')\n",
     "\n",
     "  print('Benchmarking single threaded rollout', end='\\r')\n",
-    "  start = time.time()\n",
-    "  t_rollout_single_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep,  nthread=1), nbatch, ntiming)\n",
-    "  t_rollout_single_nstep  = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread=1), nstep,  ntiming)\n",
-    "  end = time.time()\n",
+    "  with rollout.Rollout(nthread=0) as rollout_:\n",
+    "    start = time.time()\n",
+    "    t_rollout_single_nbatch = benchmark(lambda x, data: nthread_rollout(model, data, x, nominal_nstep,  nthread=1, rollout_=rollout_),\n",
+    "                                        nbatch, ntiming,\n",
+    "                                        f_init=lambda x: init_model(model))\n",
+    "    t_rollout_single_nstep  = benchmark(lambda x, data: nthread_rollout(model, data, nominal_nbatch, x, nthread=1, rollout_=rollout_),\n",
+    "                                        nstep,  ntiming, f_init=lambda x: init_model(model))\n",
+    "    end = time.time()\n",
     "  print(f'Benchmarking single threaded rollout took {end-start:0.1f} seconds')\n",
     "\n",
     "  print(f'Benchmarking multithreaded rollout using {nthread} threads', end='\\r')\n",
-    "  start = time.time()\n",
-    "  t_rollout_multi_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep,  nthread), nbatch, ntiming)\n",
-    "  t_rollout_multi_nstep  = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread), nstep,  ntiming)\n",
-    "  end = time.time()\n",
+    "  with rollout.Rollout(nthread=nthread) as rollout_:\n",
+    "    start = time.time()\n",
+    "    t_rollout_multi_nbatch = benchmark(lambda x, data: nthread_rollout(model, data, x, nominal_nstep,  nthread, rollout_=rollout_),\n",
+    "                                       nbatch, ntiming, f_init=lambda x: init_model(model))\n",
+    "    t_rollout_multi_nstep  = benchmark(lambda x, data: nthread_rollout(model, data, nominal_nbatch, x, nthread, rollout_=rollout_),\n",
+    "                                       nstep,  ntiming, f_init=lambda x: init_model(model))\n",
+    "    end = time.time()\n",
     "  print(f'Benchmarking multithreaded rollout using {nthread} threads took {end-start:0.1f} seconds')\n",
     "\n",
     "  return (t_python_nbatch, t_rollout_single_nbatch, t_rollout_multi_nbatch,\n",
     "          t_python_nstep, t_rollout_single_nstep, t_rollout_multi_nstep)\n",
     "\n",
-    "def plot_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep):\n",
+    "def plot_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep, title):\n",
     "  (t_python_nbatch, t_rollout_single_nbatch, t_rollout_multi_nbatch,\n",
     "   t_python_nstep, t_rollout_single_nstep, t_rollout_multi_nstep) = results\n",
     "\n",
@@ -1239,6 +1273,7 @@
     "\n",
     "  ax2.legend(loc=(1.04, 0.0))\n",
     "  fig.set_size_inches(10, 4)\n",
+    "  plt.suptitle(title)\n",
     "  plt.tight_layout()"
    ]
   },
@@ -1254,23 +1289,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "f7e54830",
    "metadata": {
     "id": "f7e54830"
    },
    "outputs": [],
    "source": [
-    "nominal_nbatch = 100 # Batch size to use when testing different nstep\n",
-    "nominal_nstep = 1000 # Step count to use when testing different nbatch\n",
-    "nbatch = [1, 10, 100, 500, 1000] # Batch sizes to benchmark\n",
-    "nstep = sorted([1, 10, 100, 1000, 2000, 4000]) # Step counts to benchmark\n",
+    "nominal_nbatch = 256 # Batch size to use when testing different nstep\n",
+    "nominal_nstep = 5 # Step count to use when testing different nbatch\n",
+    "nbatch = [1, 256, 2048, 8192]\n",
+    "nstep = [1, 10, 100, 1000]\n",
     "\n",
-    "top_benchmark_results = benchmark_rollout(top_model, top_data, init_top,\n",
+    "top_benchmark_results = benchmark_rollout(top_model, init_top,\n",
     "                                          nbatch, nstep,\n",
     "                                          nominal_nbatch, nominal_nstep)\n",
     "plot_benchmark(top_benchmark_results, nbatch, nstep,\n",
-    "               nominal_nbatch, nominal_nstep)"
+    "               nominal_nbatch, nominal_nstep,\n",
+    "               title='Tippe Top')"
    ]
   },
   {
@@ -1285,23 +1321,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "c9e58c6c",
    "metadata": {
     "id": "c9e58c6c"
    },
    "outputs": [],
    "source": [
-    "nominal_nbatch = 200 # Batch size to use when testing different nstep\n",
-    "nominal_nstep = 500 # Step count to use when testing different nbatch\n",
-    "nbatch = [1, 10, 100, 200, 400] # Batch sizes to benchmark\n",
-    "nstep = sorted([1, 10, 100, 500, 1000]) # Step counts to benchmark\n",
+    "nominal_nbatch = 256 # Batch size to use when testing different nstep\n",
+    "nominal_nstep = 5 # Step count to use when testing different nbatch\n",
+    "nbatch = [1, 256, 2048, 8192] # Batch sizes to benchmark\n",
+    "nstep = [1, 10, 100, 1000] # Step counts to benchmark\n",
     "\n",
-    "humanoid_benchmark_results = benchmark_rollout(humanoid_model, humanoid_data,\n",
-    "                                               init_humanoid, nbatch, nstep,\n",
+    "humanoid_benchmark_results = benchmark_rollout(humanoid_model, init_humanoid,\n",
+    "                                               nbatch, nstep,\n",
     "                                               nominal_nbatch, nominal_nstep)\n",
     "plot_benchmark(humanoid_benchmark_results, nbatch, nstep,\n",
-    "               nominal_nbatch, nominal_nstep)"
+    "               nominal_nbatch, nominal_nstep,\n",
+    "               title='Humanoid')"
    ]
   },
   {
@@ -1316,21 +1353,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "83d775d4",
    "metadata": {
     "id": "83d775d4"
    },
    "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 = 128 # Batch size to use when testing different nstep\n",
+    "nominal_nstep = 5 # Step count to use when testing different nbatch\n",
+    "nbatch = [1, 64, 128, 256] # Batch sizes to benchmark\n",
+    "nstep = [1, 10, 100, 1000] # Step counts to benchmark\n",
     "\n",
     "humanoid100_benchmark_results = benchmark_rollout(\n",
     "    humanoid100_model,\n",
-    "    humanoid100_data,\n",
     "    init_humanoid100,\n",
     "    nbatch,\n",
     "    nstep,\n",
@@ -1338,7 +1374,8 @@
     "    nominal_nstep,\n",
     ")\n",
     "plot_benchmark(humanoid100_benchmark_results, nbatch, nstep,\n",
-    "               nominal_nbatch, nominal_nstep)"
+    "               nominal_nbatch, nominal_nstep,\n",
+    "               title='Humanoid100')"
    ]
   },
   {
@@ -1360,14 +1397,14 @@
    "source": [
     "Next we will benchmark `rollout` and MJX using the tippe top and humanoid models (humanoid100 is not supported by MJX).\n",
     "\n",
-    "The benchmark below takes about 5.5 minutes on an AMD 5800X3D and an NVIDIA 4090. Almost half the time is spent compiling the JIT functions. The JIT functions are cached so that subsequent runs of the benchmark run much faster.\n",
+    "The next two benchmarks take about 16.5 minutes total on an AMD 5800X3D and an NVIDIA 4090. Most of the time is spent JIT compiling the MJX functions. The JIT functions are cached so that subsequent runs of the benchmark run much faster.\n",
     "\n",
-    "**Note:** MJX is most useful when coupled with something else that runs best on a GPU, like a neural network. Without any such additional workload, CPU based simulation will sometimes be faster, especially when using less than state-of-the-art GPUs. In the results below, the tippe top model runs faster on the 4090 with batch sizes in the 1000's, however the humanoid model always runs slower than the CPU."
+    "**Note:** MJX is most useful when coupled with something else that runs best on a GPU, like a neural network. Without any such additional workload, CPU based simulation will sometimes be faster, especially when using less than state-of-the-art GPUs."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "7c86d157",
    "metadata": {
     "cellView": "form",
@@ -1376,59 +1413,77 @@
    "outputs": [],
    "source": [
     "#@title MJX helper functions\n",
-    "def init_mjx_batch(model, init_model, nbatch, skip_jit=False):\n",
+    "def init_mjx_batch(model, init_model, nbatch, nstep, skip_jit=False):\n",
     "  data = init_model(model)\n",
     "\n",
     "  # Make MJX versions of model and data\n",
     "  mjx_model = mjx.put_model(model)\n",
     "  mjx_data = mjx.put_data(model, data)\n",
     "\n",
-    "  jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
     "  batch = jax.vmap(lambda x: mjx_data)(jp.array(list(range(nbatch))))\n",
+    "  jax.block_until_ready(batch)\n",
     "\n",
-    "  # Trigger JIT for model/batch so as not to include JIT time in benchmarking information\n",
     "  if not skip_jit:\n",
-    "    batch = jit_step(mjx_model, batch)\n",
-    "\n",
-    "  return mjx_model, mjx_data, jit_step, batch\n",
-    "\n",
-    "def mjx_rollout(model, init_model, nbatch, nstep, jit_step=None):\n",
-    "  # Initialize model, skip JIT of stepping function if possible\n",
-    "  if jit_step is None:\n",
-    "    mjx_model, _, jit_step, batch = init_mjx_batch(model, init_model, nbatch)\n",
+    "    start = time.time()\n",
+    "    jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
+    "    def unroll(d, _):\n",
+    "      d = jit_step(mjx_model, d)\n",
+    "      return d, None\n",
+    "    jit_unroll = jax.jit(lambda d: jax.lax.scan(unroll, d, None, length=nstep, unroll=4)[0])\n",
+    "    jit_unroll = jit_unroll.lower(batch).compile()\n",
+    "    end = time.time()\n",
+    "    jit_time = end - start\n",
     "  else:\n",
-    "    mjx_model, _, _, batch = init_mjx_batch(model, init_model, nbatch, skip_jit=True)\n",
+    "    jit_unroll = None\n",
+    "    jit_time = 0.0\n",
     "\n",
-    "  for _ in range(nstep):\n",
-    "    batch = jit_step(mjx_model, batch)\n",
+    "  return mjx_model, mjx_data, jit_unroll, batch, jit_time\n",
     "\n",
-    "def benchmark_mjx(model, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1, jit_steps=None):\n",
+    "def mjx_rollout(batch, jit_unroll):\n",
+    "  batch = jit_unroll(batch)\n",
+    "  jax.block_until_ready(batch)\n",
+    "\n",
+    "def benchmark_mjx(model, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1, jit_unroll_cache=None):\n",
     "  print(f'Benchmarking multithreaded rollout using {nthread} threads', end=\"\\r\")\n",
-    "  start = time.time()\n",
-    "  t_rollout_multi_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep,  nthread), nbatch, ntiming)\n",
-    "  t_rollout_multi_nstep  = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread), nstep,  ntiming)\n",
-    "  end = time.time()\n",
+    "  with rollout.Rollout(nthread=nthread) as rollout_:\n",
+    "    start = time.time()\n",
+    "    t_rollout_multi_nbatch = benchmark(lambda x, data: nthread_rollout(model, data, x, nominal_nstep,  nthread, rollout_),\n",
+    "                                       nbatch, ntiming, f_init=lambda x: init_model(model))\n",
+    "    t_rollout_multi_nstep  = benchmark(lambda x, data: nthread_rollout(model, data, nominal_nbatch, x, nthread, rollout_),\n",
+    "                                       nstep,  ntiming, f_init=lambda x: init_model(model))\n",
+    "    end = time.time()\n",
     "  print(f'Benchmarking multithreaded rollout using {nthread} threads took {end-start:0.1f} seconds')\n",
     "\n",
     "  print('Running JIT for MJX', end='\\r')\n",
-    "  start = time.time()\n",
-    "  if jit_steps is None: jit_steps = {}\n",
-    "  for n in nbatch + [nominal_nbatch,]:\n",
-    "    if n not in jit_steps:\n",
-    "      _, _, jit_steps[n], _ = init_mjx_batch(model, init_model, n)\n",
-    "  end = time.time()\n",
-    "  print(f'Running JIT for MJX took {end-start:0.1f} seconds')\n",
+    "  total_jit = 0.0\n",
+    "  if jit_unroll_cache is None:\n",
+    "    jit_unroll_cache = {}\n",
+    "  if f'nbatch_{nominal_nstep}' not in jit_unroll_cache:\n",
+    "    jit_unroll_cache[f'nbatch_{nominal_nstep}'] = {}\n",
+    "  if f'nstep_{nominal_nbatch}' not in jit_unroll_cache:\n",
+    "    jit_unroll_cache[f'nstep_{nominal_nbatch}'] = {}\n",
+    "  for n in nbatch:\n",
+    "    if n not in jit_unroll_cache[f'nbatch_{nominal_nstep}']:\n",
+    "      _, _, jit_unroll_cache[f'nbatch_{nominal_nstep}'][n], _, jit_time = init_mjx_batch(model, init_model, n, nominal_nstep)\n",
+    "      total_jit += jit_time\n",
+    "  for n in nstep:\n",
+    "    if n not in jit_unroll_cache[f'nstep_{nominal_nbatch}']:\n",
+    "      _, _, jit_unroll_cache[f'nstep_{nominal_nbatch}'][n], _, jit_time = init_mjx_batch(model, init_model, nominal_nbatch, n)\n",
+    "      total_jit += jit_time\n",
+    "  print(f'Running JIT for MJX took {total_jit:0.1f} seconds')\n",
     "\n",
     "  print('Benchmarking MJX', end='\\r')\n",
     "  start = time.time()\n",
-    "  t_mjx_nbatch = benchmark(lambda x: mjx_rollout(model, init_model, x, nominal_nstep, jit_steps[x]), nbatch, ntiming)\n",
-    "  t_mjx_nstep  = benchmark(lambda x: mjx_rollout(model, init_model, nominal_nbatch, x, jit_steps[nominal_nbatch]), nstep, ntiming)\n",
+    "  t_mjx_nbatch = benchmark(lambda x, x_init: mjx_rollout(x_init[3], jit_unroll_cache[f'nbatch_{nominal_nstep}'][x]),\n",
+    "                           nbatch, ntiming, f_init=lambda x: init_mjx_batch(model, init_model, x, nominal_nstep, skip_jit=True))\n",
+    "  t_mjx_nstep  = benchmark(lambda x, x_init: mjx_rollout(x_init[3], jit_unroll_cache[f'nstep_{nominal_nbatch}'][x]),\n",
+    "                           nstep, ntiming, f_init=lambda x: init_mjx_batch(model, init_model, nominal_nbatch, x, skip_jit=True))\n",
     "  end = time.time()\n",
     "  print(f'Benchmarking MJX took {end-start:0.1f} seconds')\n",
     "\n",
     "  return t_rollout_multi_nbatch, t_rollout_multi_nstep, t_mjx_nbatch, t_mjx_nstep\n",
     "\n",
-    "def plot_mjx_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep):\n",
+    "def plot_mjx_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep, title):\n",
     "  t_rollout_multi_nbatch, t_rollout_multi_nstep, t_mjx_nbatch, t_mjx_nstep = results\n",
     "\n",
     "  width = 0.333\n",
@@ -1463,12 +1518,12 @@
     "\n",
     "  ax2.legend(loc=(1.04, 0.0))\n",
     "  fig.set_size_inches(10, 4)\n",
+    "  plt.suptitle(title)\n",
     "  plt.tight_layout()\n",
     "\n",
     "# Caches for jit_step functions, they take a long time to compile\n",
-    "top_jit_steps = {}\n",
-    "humanoid_jit_steps = {}\n",
-    "humanoid100_jit_steps = {}"
+    "top_jit_unroll_cache = {}\n",
+    "humanoid_jit_unroll_cache = {}"
    ]
   },
   {
@@ -1483,20 +1538,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "98c580b0",
    "metadata": {
     "id": "98c580b0"
    },
    "outputs": [],
    "source": [
-    "nominal_nbatch = 50000 # Batch size to use when testing different nstep\n",
-    "nominal_nstep = 200 # Step count to use when testing different nbatch\n",
-    "nbatch = [100, 1000, 10000, 50000, 100000] # Batch sizes to benchmark\n",
+    "nominal_nbatch = 16384 # Batch size to use when testing different nstep\n",
+    "nominal_nstep = 5 # Step count to use when testing different nbatch\n",
+    "nbatch = [4096, 16384, 65536, 131072] # Batch sizes to benchmark\n",
     "nstep = [1, 10, 100, 200] # Step counts to benchmark\n",
     "\n",
-    "mjx_top_results = benchmark_mjx(top_model, init_top, nbatch, nstep, nominal_nbatch, nominal_nstep, jit_steps=top_jit_steps)\n",
-    "plot_mjx_benchmark(mjx_top_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
+    "mjx_top_results = benchmark_mjx(top_model, init_top, nbatch, nstep, nominal_nbatch, nominal_nstep,\n",
+    "                                jit_unroll_cache=top_jit_unroll_cache)\n",
+    "plot_mjx_benchmark(mjx_top_results, nbatch, nstep, nominal_nbatch, nominal_nstep, title='MJX Tippe Top')"
    ]
   },
   {
@@ -1511,20 +1567,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "53166ae1",
    "metadata": {
     "id": "53166ae1"
    },
    "outputs": [],
    "source": [
-    "nominal_nbatch = 10000 # Batch size to use when testing different nstep\n",
-    "nominal_nstep = 200 # Step count to use when testing different nbatch\n",
-    "nbatch = [100, 1000, 10000, 30000] # Batch sizes to benchmark\n",
-    "nstep = [1, 10, 100, 200, 400] # Step counts to benchmark\n",
+    "nominal_nbatch = 4096 # Batch size to use when testing different nstep\n",
+    "nominal_nstep = 5 # Step count to use when testing different nbatch\n",
+    "nbatch = [1024, 4096, 16384, 32768] # Batch sizes to benchmark\n",
+    "nstep = [1, 10, 100, 200] # Step counts to benchmark\n",
     "\n",
-    "mjx_humanoid_results = benchmark_mjx(humanoid_model, init_humanoid, nbatch, nstep, nominal_nbatch, nominal_nstep, jit_steps=humanoid_jit_steps)\n",
-    "plot_mjx_benchmark(mjx_humanoid_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
+    "mjx_humanoid_results = benchmark_mjx(humanoid_model, init_humanoid, nbatch, nstep, nominal_nbatch, nominal_nstep,\n",
+    "                                     jit_unroll_cache=humanoid_jit_unroll_cache)\n",
+    "plot_mjx_benchmark(mjx_humanoid_results, nbatch, nstep, nominal_nbatch, nominal_nstep, title='MJX Humanoid')"
    ]
   },
   {
@@ -1538,14 +1595,14 @@
     "\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",
-    "Here we will produce a similar plot to compare MJX and with `rollout`. On a 5800X3D and 4090 devices the benchmark takes about 6.5 minutes to run.\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 was run on different devices and in particular an A100. Additionally, to run on a 4090 the batch size was redued from 8192 to 4096."
+    "**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."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "3d6be608",
    "metadata": {
     "id": "3d6be608"
@@ -1556,41 +1613,48 @@
     "nbatch = 8192 // 2 # The original benchmark ran with a batch size of 8192, but on a 4090 we can only fit about 4096 humanoids\n",
     "nstep = 200\n",
     "\n",
-    "jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
+    "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
     "t_rollout = []\n",
     "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(\n",
+    "  nhumanoid_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",
+    "  nhumanoid_data = mujoco.MjData(nhumanoid_model)\n",
     "\n",
-    "  mjx_model = mjx.put_model(model)\n",
-    "  mjx_data = mjx.put_data(model, data)\n",
+    "  mjx_model = mjx.put_model(nhumanoid_model)\n",
+    "  mjx_data = mjx.put_data(nhumanoid_model, nhumanoid_data)\n",
     "  batch = jax.vmap(lambda x: mjx_data)(jp.array(list(range(nbatch))))\n",
+    "  jax.block_until_ready(batch)\n",
     "\n",
-    "  start = time.perf_counter()\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",
+    "  with rollout.Rollout(nthread=nthread) as rollout_:\n",
+    "    initial_state = get_state(nhumanoid_model, nhumanoid_data, nbatch)\n",
+    "    start = time.perf_counter()\n",
+    "    rollout_.rollout([nhumanoid_model]*nbatch,\n",
+    "                     [copy.copy(nhumanoid_data) for _ in range(nthread)],\n",
+    "                     initial_state=initial_state,\n",
+    "                     nstep=nstep, skip_checks=True)\n",
+    "    end = time.perf_counter()\n",
     "  t_rollout.append(end-start)\n",
     "\n",
     "  # Trigger JIT for model/batch so as not to include JIT time in benchmarking information\n",
-    "  batch = jit_step(mjx_model, batch)\n",
+    "  def unroll(d, _):\n",
+    "    d = jit_step(mjx_model, d)\n",
+    "    return d, None\n",
+    "  jit_unroll = jax.jit(lambda d: jax.lax.scan(unroll, d, None, length=nstep, unroll=4)[0])\n",
+    "  jit_unroll = jit_unroll.lower(batch).compile()\n",
     "\n",
     "  start = time.perf_counter()\n",
-    "  for _ in range(nstep):\n",
-    "    batch = jit_step(mjx_model, batch)\n",
+    "  jit_unroll(batch)\n",
+    "  jax.block_until_ready(batch)\n",
     "  end = time.perf_counter()\n",
     "  t_mjx.append(end-start)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 0,
+   "execution_count": null,
    "id": "b6c5fc2e",
    "metadata": {
     "id": "b6c5fc2e"