diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index 234673e8..83e1bd30 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -187,11 +187,9 @@
       },
       "source": [
         "# Training a Policy with MJX\n",
-        "MJX is an implementation of MuJoCo written in [jax](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate training RL policies with MJX.\n",
+        "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate how to train RL policies with MJX.\n",
         "\n",
-        "First, we implement an environment `State` so that we can plug into the [brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and etc. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n",
-        "\n",
-        "`MjxEnv` is an implementation of the `brax.envs.base.Env` class that initializes `mjx.Model` and `mjx.Data` objects. Inheriting from `brax.envs.base.Env` allows us to use the training agents implemented in brax. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`."
+        "First, we implement an environment `State` so that we can plug into the [Brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n"
       ]
     },
     {
@@ -202,7 +200,7 @@
       },
       "outputs": [],
       "source": [
-        "#@title State and MjxEnv\n",
+        "#@title State\n",
         "\n",
         "@struct.dataclass\n",
         "class State(Base):\n",
@@ -223,8 +221,30 @@
         "  reward: jax.Array\n",
         "  done: jax.Array\n",
         "  metrics: Dict[str, jax.Array] = struct.field(default_factory=dict)\n",
-        "  info: Dict[str, Any] = struct.field(default_factory=dict)\n",
+        "  info: Dict[str, Any] = struct.field(default_factory=dict)\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "acpXtDLNXLV9"
+      },
+      "source": [
         "\n",
+        "Next, we implement `MjxEnv`, an environment class we'll use through the notebook. `MjxEnv` initializes a `mjx.Model` and `mjx.Data` object. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`.\n",
+        "\n",
+        "`MjxEnv` also inherits from `brax.envs.base.Env` which allows us to use the training agents implemented in brax."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ccujYeJ5XOhx"
+      },
+      "outputs": [],
+      "source": [
+        "#@title MjxEnv\n",
         "\n",
         "class MjxEnv(Env):\n",
         "  \"\"\"API for driving an MJX system for training and inference in brax.\"\"\"\n",
@@ -305,7 +325,7 @@
         "id": "iPlFu4CiIgBN"
       },
       "source": [
-        "Now we can define environment implementations for Humanoid and the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0). The environments define the reward and stepping logic. `reset` initializes a `State`, and `step` steps through the physics step and reward logic."
+        "Finally we can implement a real environment. We choose to first implement the Humanoid environment. Notice that `reset` initializes a `State`, and `step` steps through the physics step and reward logic. The reward and stepping logic train the Humanoid to run forwards."
       ]
     },
     {
@@ -444,6 +464,372 @@
         "envs.register_environment('humanoid', Humanoid)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "P1K6IznI2y83"
+      },
+      "source": [
+        "## Visualize a Rollout\n",
+        "\n",
+        "Let's instantiate the environment and visualize a short rollout.\n",
+        "\n",
+        "NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "EhKLFK54C1CH"
+      },
+      "outputs": [],
+      "source": [
+        "# instantiate the environment\n",
+        "env_name = 'humanoid'\n",
+        "env = envs.get_environment(env_name)\n",
+        "\n",
+        "# define the jit reset/step functions\n",
+        "jit_reset = jax.jit(env.reset)\n",
+        "jit_step = jax.jit(env.step)\n",
+        "\n",
+        "# instantiate the renderer\n",
+        "renderer = mujoco.Renderer(env.model)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9f2ME2WbA5Ip"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Define a render utility function\n",
+        "\n",
+        "def get_image(state: State, camera: str) -\u003e np.ndarray:\n",
+        "  \"\"\"Renders the environment state.\"\"\"\n",
+        "  d = mujoco.MjData(env.model)\n",
+        "  # write the mjx.Data into an mjData object\n",
+        "  mjx.device_get_into(d, state.pipeline_state)\n",
+        "  mujoco.mj_forward(env.model, d)\n",
+        "  # use the mjData object to update the renderer\n",
+        "  renderer.update_scene(d, camera=camera)\n",
+        "  return renderer.render()\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ph8u-v2Q2xLS"
+      },
+      "outputs": [],
+      "source": [
+        "# initialize the state\n",
+        "state = jit_reset(jax.random.PRNGKey(0))\n",
+        "rollout = [state]\n",
+        "images = [get_image(state, camera='side')]\n",
+        "\n",
+        "# grab a trajectory\n",
+        "for i in range(10):\n",
+        "  ctrl = -0.1 * jp.ones(env.sys.nu)\n",
+        "  state = jit_step(state, ctrl)\n",
+        "  rollout.append(state)\n",
+        "  images.append(get_image(state, camera='side'))\n",
+        "\n",
+        "media.show_video(images, fps=1.0 / env.dt)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BQDG6NQ1CbZD"
+      },
+      "source": [
+        "## Train Humanoid Policy\n",
+        "\n",
+        "Let's finally train a policy with PPO to make the Humanoid run forwards. Training takes about 13-14 minutes on a Tesla V100 GPU."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xLiddQYPApBw"
+      },
+      "outputs": [],
+      "source": [
+        "train_fn = functools.partial(\n",
+        "    ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
+        "    episode_length=1000, normalize_observations=True, action_repeat=1,\n",
+        "    unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
+        "    discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
+        "    batch_size=1024, seed=0)\n",
+        "\n",
+        "\n",
+        "x_data = []\n",
+        "y_data = []\n",
+        "ydataerr = []\n",
+        "times = [datetime.now()]\n",
+        "\n",
+        "max_y, min_y = 13000, 0\n",
+        "def progress(num_steps, metrics):\n",
+        "  times.append(datetime.now())\n",
+        "  x_data.append(num_steps)\n",
+        "  y_data.append(metrics['eval/episode_reward'])\n",
+        "  ydataerr.append(metrics['eval/episode_reward_std'])\n",
+        "\n",
+        "  plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
+        "  plt.ylim([min_y, max_y])\n",
+        "\n",
+        "  plt.xlabel('# environment steps')\n",
+        "  plt.ylabel('reward per episode')\n",
+        "  plt.title(f'y={y_data[-1]:.3f}')\n",
+        "\n",
+        "  plt.errorbar(\n",
+        "      x_data, y_data, yerr=ydataerr)\n",
+        "  plt.show()\n",
+        "\n",
+        "make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\n",
+        "\n",
+        "print(f'time to jit: {times[1] - times[0]}')\n",
+        "print(f'time to train: {times[-1] - times[1]}')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YYIch0HEApBx"
+      },
+      "source": [
+        "## Save and Load Policy\n",
+        "\n",
+        "We can save and load the policy using the brax model API."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Z8gI6qH6ApBx"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Save Model\n",
+        "model_path = '/tmp/mjx_brax_policy'\n",
+        "model.save_params(model_path, params)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "h4reaWgxApBx"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Load Model and Define Inference Function\n",
+        "params = model.load_params(model_path)\n",
+        "\n",
+        "inference_fn = make_inference_fn(params)\n",
+        "jit_inference_fn = jax.jit(inference_fn)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0G357XIfApBy"
+      },
+      "source": [
+        "## Visualize Policy\n",
+        "\n",
+        "Finally we can visualize the policy."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "osYasMw4ApBy"
+      },
+      "outputs": [],
+      "source": [
+        "eval_env = envs.get_environment(env_name)\n",
+        "\n",
+        "jit_reset = jax.jit(eval_env.reset)\n",
+        "jit_step = jax.jit(eval_env.step)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d-UhypudApBy"
+      },
+      "outputs": [],
+      "source": [
+        "# initialize the state\n",
+        "rng = jax.random.PRNGKey(0)\n",
+        "state = jit_reset(rng)\n",
+        "rollout = [state]\n",
+        "images = [get_image(state, camera='side')]\n",
+        "\n",
+        "# grab a trajectory\n",
+        "n_steps = 500\n",
+        "render_every = 2\n",
+        "\n",
+        "for i in range(n_steps):\n",
+        "  act_rng, rng = jax.random.split(rng)\n",
+        "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+        "  state = jit_step(state, ctrl)\n",
+        "  rollout.append(state)\n",
+        "  if i % render_every == 0:\n",
+        "    images.append(get_image(state, camera='side'))\n",
+        "\n",
+        "  if state.done:\n",
+        "    break\n",
+        "\n",
+        "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zR-heox6LARK"
+      },
+      "source": [
+        "# MJX Policy in MuJoCo\n",
+        "\n",
+        "Note that we can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "w6ixFi4dApBy"
+      },
+      "outputs": [],
+      "source": [
+        "mj_model = eval_env.model\n",
+        "mj_data = mujoco.MjData(mj_model)\n",
+        "\n",
+        "renderer = mujoco.Renderer(mj_model)\n",
+        "ctrl = jp.zeros(mj_model.nu)\n",
+        "\n",
+        "images = []\n",
+        "for i in range(n_steps):\n",
+        "  act_rng, rng = jax.random.split(rng)\n",
+        "\n",
+        "  obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n",
+        "  ctrl, _ = jit_inference_fn(obs, act_rng)\n",
+        "\n",
+        "  mj_data.ctrl = ctrl\n",
+        "  for _ in range(eval_env._physics_steps_per_control_step):\n",
+        "    mujoco.mj_step(mj_model, mj_data)  # Physics step using MuJoCo mj_step.\n",
+        "\n",
+        "  if i % render_every == 0:\n",
+        "    renderer.update_scene(mj_data, camera='side')\n",
+        "    images.append(renderer.render())\n",
+        "\n",
+        "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "65mIPj6DQNNa"
+      },
+      "source": [
+        "# Domain Randomization\n",
+        "\n",
+        "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "h8mhzKjHQuoL"
+      },
+      "outputs": [],
+      "source": [
+        "def domain_randomize(sys, rng):\n",
+        "  \"\"\"Randomizes the mjx.Model.\"\"\"\n",
+        "  @jax.vmap\n",
+        "  def rand(rng):\n",
+        "    _, key = jax.random.split(rng, 2)\n",
+        "    # friction\n",
+        "    friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n",
+        "    friction = sys.geom_friction.at[:, 0].set(friction)\n",
+        "    # actuator\n",
+        "    _, key = jax.random.split(key, 2)\n",
+        "    gain_range = (-10, -5)\n",
+        "    param = jax.random.uniform(\n",
+        "        key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
+        "    ) + sys.actuator_gainprm[:, 0]\n",
+        "    gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
+        "    bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
+        "    return friction, gain, bias\n",
+        "\n",
+        "  friction, gain, bias = rand(rng)\n",
+        "\n",
+        "  in_axes = jax.tree_map(lambda x: None, sys)\n",
+        "  in_axes = in_axes.tree_replace({\n",
+        "      'geom_friction': 0,\n",
+        "      'actuator_gainprm': 0,\n",
+        "      'actuator_biasprm': 0,\n",
+        "  })\n",
+        "\n",
+        "  sys = sys.tree_replace({\n",
+        "      'geom_friction': friction,\n",
+        "      'actuator_gainprm': gain,\n",
+        "      'actuator_biasprm': bias,\n",
+        "  })\n",
+        "\n",
+        "  return sys, in_axes"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gnsZo-GWSYYj"
+      },
+      "source": [
+        "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjModel` along with a dictionary specifying the axes that are batched."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1K45Kp2ASV9s"
+      },
+      "outputs": [],
+      "source": [
+        "rng = jax.random.PRNGKey(0)\n",
+        "rng = jax.random.split(rng, 10)\n",
+        "batched_sys, _ = domain_randomize(env.sys, rng)\n",
+        "\n",
+        "print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
+        "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
+        "\n",
+        "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
+        "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "efnxNOnpQFuC"
+      },
+      "source": [
+        "## Quadruped Env\n",
+        "\n",
+        "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0) and an environment that trains a joystick policy."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -483,13 +869,14 @@
         "                    # Penalize non-zero roll and pitch angles. L2 penalty.\n",
         "                    orientation=-5.0,\n",
         "                    # L2 regularization of joint torques, |tau|^2.\n",
-        "                    torques=-0.0002,\n",
+        "                    # torques=-0.0002,\n",
+        "                    torques=-0.002,\n",
         "                    # Penalize the change in the action and encourage smooth\n",
         "                    # actions. L2 regularization |action - last_action|^2\n",
-        "                    action_rate=-0.3,\n",
+        "                    action_rate=-0.1,\n",
         "                    # Encourage long swing steps.  However, it does not\n",
         "                    # encourage high clearances.\n",
-        "                    feet_air_time=0.1,\n",
+        "                    feet_air_time=0.2,\n",
         "                    # Encourage no motion at zero command, L2 regularization\n",
         "                    # |q - q_default|^2.\n",
         "                    stand_still=-0.5,\n",
@@ -550,7 +937,7 @@
         "    self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n",
         "\n",
         "  def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n",
-        "    lin_vel_x = [-0.6, 1.5]  # min max [m/s]\n",
+        "    lin_vel_x = [-0.6, 1.0]  # min max [m/s]\n",
         "    lin_vel_y = [-0.8, 0.8]  # min max [m/s]\n",
         "    ang_vel_yaw = [-0.7, 0.7]  # min max [rad/s]\n",
         "\n",
@@ -857,171 +1244,19 @@
         "envs.register_environment('barkour', BarkourEnv)"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "P1K6IznI2y83"
-      },
-      "source": [
-        "## Visualize a Rollout\n",
-        "\n",
-        "Let's visualize an environment rollout by running a few env steps with sinusoidal actuation."
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "cellView": "form",
-        "id": "EhKLFK54C1CH"
+        "id": "pi_yrcz-Qp3W"
       },
       "outputs": [],
       "source": [
-        "# instantiate the environment\n",
-        "\n",
-        "env_name = 'humanoid'  # @param ['barkour', 'humanoid']\n",
+        "env_name = 'barkour'\n",
         "env = envs.get_environment(env_name)\n",
         "\n",
-        "# define the jit reset/step functions\n",
-        "jit_reset = jax.jit(env.reset)\n",
-        "jit_step = jax.jit(env.step)\n",
-        "\n",
-        "# instantiate the renderer\n",
-        "renderer = mujoco.Renderer(env.model)\n",
-        "\n",
-        "def get_image(state: State):\n",
-        "  d = mujoco.MjData(env.model)\n",
-        "  # write the mjx.Data into an mjData object\n",
-        "  mjx.device_get_into(d, state.pipeline_state)\n",
-        "  mujoco.mj_forward(env.model, d)\n",
-        "  if env_name == 'barkour':\n",
-        "    camera='track'\n",
-        "  elif env_name == 'humanoid':\n",
-        "    camera = 'side'\n",
-        "  else:\n",
-        "    raise NotImplementedError(env_name)\n",
-        "  # use the mjData object to update the renderer\n",
-        "  renderer.update_scene(d, camera=camera)\n",
-        "  return renderer.render()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Ph8u-v2Q2xLS"
-      },
-      "outputs": [],
-      "source": [
-        "# initialize the state\n",
-        "state = jit_reset(jax.random.PRNGKey(0))\n",
-        "rollout = [state]\n",
-        "images = [get_image(state)]\n",
-        "\n",
-        "# grab a trajectory\n",
-        "n_steps = 500\n",
-        "render_every = 5\n",
-        "\n",
-        "for i in range(n_steps):\n",
-        "  ctrl = jp.ones(env.sys.nu) * jp.sin(2.0 * jp.pi * i / 500)\n",
-        "  state = jit_step(state, ctrl)\n",
-        "  rollout.append(state)\n",
-        "  if i % render_every == 0:\n",
-        "    images.append(get_image(state))\n",
-        "\n",
-        "media.show_video(images, fps=1.0 / env.dt / render_every)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "65mIPj6DQNNa"
-      },
-      "source": [
-        "# Domain Randomization\n",
-        "\n",
-        "We have all the pieces to train a policy, but we might also want to include randomization over certain `mjModel` parameters. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "h8mhzKjHQuoL"
-      },
-      "outputs": [],
-      "source": [
-        "def domain_randomize(sys, rng):\n",
-        "  \"\"\"Randomizes the mjx.Model.\"\"\"\n",
-        "  @jax.vmap\n",
-        "  def rand(rng):\n",
-        "    _, key = jax.random.split(rng, 2)\n",
-        "    # friction\n",
-        "    friction = jax.random.uniform(key, (1,), minval=0.8, maxval=1.2)\n",
-        "    friction = sys.geom_friction.at[:, 0].set(friction)\n",
-        "    # actuator\n",
-        "    _, key = jax.random.split(key, 2)\n",
-        "    gain_range = (-5, 5)\n",
-        "    param = jax.random.uniform(\n",
-        "        key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
-        "    ) + sys.actuator_gainprm[:, 0]\n",
-        "    gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
-        "    bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
-        "    return friction, gain, bias\n",
-        "\n",
-        "  friction, gain, bias = rand(rng)\n",
-        "\n",
-        "  in_axes = jax.tree_map(lambda x: None, sys)\n",
-        "  in_axes = in_axes.tree_replace({\n",
-        "      'geom_friction': 0,\n",
-        "      'actuator_gainprm': 0,\n",
-        "      'actuator_biasprm': 0,\n",
-        "  })\n",
-        "\n",
-        "  sys = sys.tree_replace({\n",
-        "      'geom_friction': friction,\n",
-        "      'actuator_gainprm': gain,\n",
-        "      'actuator_biasprm': bias,\n",
-        "  })\n",
-        "\n",
-        "  return sys, in_axes"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "gnsZo-GWSYYj"
-      },
-      "source": [
-        "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1K45Kp2ASV9s"
-      },
-      "outputs": [],
-      "source": [
-        "rng = jax.random.PRNGKey(0)\n",
-        "rng = jax.random.split(rng, 10)\n",
-        "batched_sys, _ = domain_randomize(env.sys, rng)\n",
-        "\n",
-        "print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
-        "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
-        "\n",
-        "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
-        "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "LdqPYBlPVWwc"
-      },
-      "source": [
-        "Note that we'll pass the `domain_randomize` function into the brax trainer; the brax training algorithms will call the randomizer over the environment batch dimension."
+        "# re-instantiate the renderer\n",
+        "renderer = mujoco.Renderer(env.model)"
       ]
     },
     {
@@ -1030,9 +1265,9 @@
         "id": "nxaNFP9mA23H"
       },
       "source": [
-        "# Train Policy\n",
+        "## Train Policy\n",
         "\n",
-        "Let's finally train a policy with PPO. Training takes about 12-13 minutes on a Tesla V100 GPU for both environments."
+        "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes about 14 minutes on a Tesla V100 GPU."
       ]
     },
     {
@@ -1043,59 +1278,28 @@
       },
       "outputs": [],
       "source": [
+        "make_networks_factory = functools.partial(\n",
+        "    ppo_networks.make_ppo_networks,\n",
+        "        policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
+        "train_fn = functools.partial(\n",
+        "      ppo.train,\n",
+        "      num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n",
+        "      episode_length=1000, normalize_observations=True,\n",
+        "      action_repeat=1, unroll_length=20, num_minibatches=8, gae_lambda=0.95,\n",
+        "      num_updates_per_batch=4, discounting=0.99, learning_rate=3e-4,\n",
+        "      entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n",
+        "      network_factory=make_networks_factory,\n",
+        "      num_resets_per_eval=10,\n",
+        "      randomization_fn=domain_randomize, seed=0)\n",
         "\n",
-        "if env_name == 'barkour':\n",
-        "  make_networks_factory = functools.partial(\n",
-        "      ppo_networks.make_ppo_networks,\n",
-        "          policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
-        "  train_fn = functools.partial(\n",
-        "        ppo.train,\n",
-        "        num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n",
-        "        episode_length=1000, normalize_observations=True,\n",
-        "        action_repeat=1, unroll_length=20, num_minibatches=32, gae_lambda=0.95,\n",
-        "        num_updates_per_batch=4, discounting=0.99, learning_rate=3.0e-4,\n",
-        "        entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n",
-        "        network_factory=make_networks_factory,\n",
-        "        num_resets_per_eval=10,\n",
-        "        randomization_fn=domain_randomize, seed=0)\n",
-        "elif env_name == 'humanoid':\n",
-        "  train_fn = functools.partial(\n",
-        "      ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
-        "      episode_length=1000, normalize_observations=True, action_repeat=1,\n",
-        "      unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
-        "      discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
-        "      batch_size=1024, seed=0)\n",
-        "else:\n",
-        "  raise NotImplementedError(\n",
-        "      f'env_name: {env_name} is not implemented in this notebook.')\n",
-        "\n",
-        "\n",
-        "max_y = {'barkour':30, 'humanoid': 13000}[env_name]\n",
-        "min_y = {'barkour': -15}.get(env_name, 0)\n",
         "\n",
         "x_data = []\n",
         "y_data = []\n",
         "ydataerr = []\n",
         "times = [datetime.now()]\n",
+        "max_y, min_y = 30, 0\n",
         "\n",
-        "def progress(num_steps, metrics):\n",
-        "  times.append(datetime.now())\n",
-        "  x_data.append(num_steps)\n",
-        "  y_data.append(metrics['eval/episode_reward'])\n",
-        "  ydataerr.append(metrics['eval/episode_reward_std'])\n",
-        "\n",
-        "  plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
-        "  plt.ylim([min_y, max_y])\n",
-        "\n",
-        "  plt.xlabel('# environment steps')\n",
-        "  plt.ylabel('reward per episode')\n",
-        "  plt.title(f'y={y_data[-1]:.3f}')\n",
-        "\n",
-        "  plt.errorbar(\n",
-        "      x_data, y_data, yerr=ydataerr)\n",
-        "  plt.show()\n",
-        "\n",
-        "# Reset environments since internals may be overwritten by tracers due to\n",
+        "# Reset environments since internals may be overwritten by tracers from the\n",
         "# domain randomization function.\n",
         "env = envs.get_environment(env_name)\n",
         "eval_env = envs.get_environment(env_name)\n",
@@ -1107,30 +1311,6 @@
         "print(f'time to train: {times[-1] - times[1]}')"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-dRmWHsHB_1K"
-      },
-      "source": [
-        "# Save and Load Policy\n",
-        "\n",
-        "We can save and load the policy using the brax model API."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "u9sdN4Xa5JoH"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Save Model\n",
-        "model_path = '/tmp/mjx_brax_policy'\n",
-        "model.save_params(model_path, params)"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1139,7 +1319,9 @@
       },
       "outputs": [],
       "source": [
-        "#@title Load Model and Define Inference Function\n",
+        "# Save and reload params.\n",
+        "model_path = '/tmp/mjx_brax_quadruped_policy'\n",
+        "model.save_params(model_path, params)\n",
         "params = model.load_params(model_path)\n",
         "\n",
         "inference_fn = make_inference_fn(params)\n",
@@ -1152,9 +1334,9 @@
         "id": "L01IrN4oCIkC"
       },
       "source": [
-        "# Visualize Policy\n",
+        "## Visualize Policy\n",
         "\n",
-        "Finally we can visualize the policy. For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
+        "For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
       ]
     },
     {
@@ -1183,7 +1365,7 @@
         "# @markdown Commands **only used for Barkour Env**:\n",
         "x_vel = 1.0  #@param {type: \"number\"}\n",
         "y_vel = 0.0  #@param {type: \"number\"}\n",
-        "ang_vel = -0.0  #@param {type: \"number\"}\n",
+        "ang_vel = -0.5  #@param {type: \"number\"}\n",
         "\n",
         "the_command = jp.array([x_vel, y_vel, ang_vel])\n",
         "\n",
@@ -1192,7 +1374,7 @@
         "state = jit_reset(rng)\n",
         "state.info['command'] = the_command\n",
         "rollout = [state]\n",
-        "images = [get_image(state)]\n",
+        "images = [get_image(state, camera='track')]\n",
         "\n",
         "# grab a trajectory\n",
         "n_steps = 500\n",
@@ -1204,7 +1386,7 @@
         "  state = jit_step(state, ctrl)\n",
         "  rollout.append(state)\n",
         "  if i % render_every == 0:\n",
-        "    images.append(get_image(state))\n",
+        "    images.append(get_image(state, camera='track'))\n",
         "\n",
         "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
       ]
@@ -1216,7 +1398,12 @@
       "gpuClass": "premium",
       "gpuType": "V100",
       "private_outputs": true,
-      "provenance": [],
+      "provenance": [
+        {
+          "file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2",
+          "timestamp": 1697763780236
+        }
+      ],
       "toc_visible": true
     },
     "kernelspec": {