diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index f27461c6..8644de41 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -157,20 +157,25 @@
       "source": [
         "#@title Import MuJoCo, MJX, and Brax\n",
         "\n",
+        "\n",
         "from datetime import datetime\n",
         "import functools\n",
+        "from IPython.display import HTML\n",
         "import jax\n",
         "from jax import numpy as jp\n",
         "import numpy as np\n",
-        "from typing import Any, Dict, Tuple, Union\n",
+        "from typing import Any, Dict, Sequence, Tuple, Union\n",
         "\n",
+        "from brax import base\n",
         "from brax import envs\n",
         "from brax import math\n",
         "from brax.base import Base, Motion, Transform\n",
-        "from brax.envs.base import Env, State\n",
+        "from brax.envs.base import Env, MjxEnv, State\n",
+        "from brax.mjx.base import State as MjxState\n",
         "from brax.training.agents.ppo import train as ppo\n",
         "from brax.training.agents.ppo import networks as ppo_networks\n",
-        "from brax.io import model\n",
+        "from brax.io import html, mjcf, model\n",
+        "\n",
         "from etils import epath\n",
         "from flax import struct\n",
         "from matplotlib import pyplot as plt\n",
@@ -180,6 +185,211 @@
         "from mujoco import mjx\n"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Nj4-Xmx4DFaq"
+      },
+      "source": [
+        "# Introduction to MJX\n",
+        "\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 will demonstrate how to train RL policies with MJX.\n",
+        "\n",
+        "Before we get into hefty RL workloads, let's get started with a simpler example! The entrypoint into MJX is through MuJoCo, so first we load a MuJoCo model:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bNus3mbbDz6a"
+      },
+      "outputs": [],
+      "source": [
+        "xml = \"\"\"\n",
+        "\u003cmujoco\u003e\n",
+        "  \u003cworldbody\u003e\n",
+        "    \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n",
+        "    \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n",
+        "      \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n",
+        "      \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
+        "      \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
+        "    \u003c/body\u003e\n",
+        "  \u003c/worldbody\u003e\n",
+        "\u003c/mujoco\u003e\n",
+        "\"\"\"\n",
+        "\n",
+        "# Make model, data, and renderer\n",
+        "mj_model = mujoco.MjModel.from_xml_string(xml)\n",
+        "mj_data = mujoco.MjData(mj_model)\n",
+        "renderer = mujoco.Renderer(mj_model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Po5oykJbFQbj"
+      },
+      "source": [
+        "Next we take the MuJoCo model and data, and place them on the GPU device using MJX."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "TSpoOWqeEC3P"
+      },
+      "outputs": [],
+      "source": [
+        "mjx_model = mjx.put_model(mj_model)\n",
+        "mjx_data = mjx.put_data(mj_model, mj_data)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6rxMMSs4OJJf"
+      },
+      "source": [
+        "Below, we print the `qpos` from MuJoCo and MJX. Notice that the `qpos` for the mjData is a numpy array living on the CPU, while the `qpos` for `mjx.Data` is a JAX Array living on the GPU device."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ZOD582pfOLP-"
+      },
+      "outputs": [],
+      "source": [
+        "print(mj_data.qpos, type(mj_data.qpos))\n",
+        "print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZShF9-o_JLm3"
+      },
+      "source": [
+        "Let's run the simulation in MuJoCo and render the trajectory. This example is taken from the [MuJoCo tutorial](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HDlPlX05I3m-"
+      },
+      "outputs": [],
+      "source": [
+        "# enable joint visualization option:\n",
+        "scene_option = mujoco.MjvOption()\n",
+        "scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True\n",
+        "\n",
+        "duration = 3.8  # (seconds)\n",
+        "framerate = 60  # (Hz)\n",
+        "\n",
+        "frames = []\n",
+        "mujoco.mj_resetData(mj_model, mj_data)\n",
+        "while mj_data.time \u003c duration:\n",
+        "  mujoco.mj_step(mj_model, mj_data)\n",
+        "  if len(frames) \u003c mj_data.time * framerate:\n",
+        "    renderer.update_scene(mj_data, scene_option=scene_option)\n",
+        "    pixels = renderer.render()\n",
+        "    frames.append(pixels)\n",
+        "\n",
+        "# Simulate and display video.\n",
+        "media.show_video(frames, fps=framerate)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "m70b_RxBJOyd"
+      },
+      "source": [
+        "Now let's run the same exact simulation on the GPU device using MJX!\n",
+        "\n",
+        "In the example below, we use `mjx.step` instead of `mujoco.mj_step`, and we also [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) the `mjx.step` so that it runs efficiently on the GPU. After each step, we convert the `mjx.Data` back to `mjData` so that we can use the MuJoCo renderer.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Pr29xq0-JRQv"
+      },
+      "outputs": [],
+      "source": [
+        "\n",
+        "jit_step = jax.jit(mjx.step)\n",
+        "\n",
+        "frames = []\n",
+        "mujoco.mj_resetData(mj_model, mj_data)\n",
+        "mjx_data = mjx.put_data(mj_model, mj_data)\n",
+        "while mjx_data.time \u003c duration:\n",
+        "  mjx_data = jit_step(mjx_model, mjx_data)\n",
+        "  if len(frames) \u003c mjx_data.time * framerate:\n",
+        "    mj_data = mjx.get_data(mj_model, mjx_data)\n",
+        "    renderer.update_scene(mj_data, scene_option=scene_option)\n",
+        "    pixels = renderer.render()\n",
+        "    frames.append(pixels)\n",
+        "\n",
+        "media.show_video(frames, fps=framerate)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wXsQ4qO2KO3Q"
+      },
+      "source": [
+        "Running single threaded physics simulation on the GPU is not very [efficient](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits). The advantage with MJX is that we can run environments in parallel on a hardware accelerated device. Let's try it out!\n",
+        "\n",
+        "In the example below, we create 4096 copies of the `mjx.Data` and we run the `mjx.step` over the batched data. Since MJX is implemented in JAX, we take advantage of [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html) to run the `mjx.step` in parallel over all `mjx.Data`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rrdrcKRVK6w9"
+      },
+      "outputs": [],
+      "source": [
+        "rng = jax.random.PRNGKey(0)\n",
+        "rng = jax.random.split(rng, 4096)\n",
+        "batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)\n",
+        "\n",
+        "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
+        "batch = jit_step(mjx_model, batch)\n",
+        "\n",
+        "print(batch.qpos)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "x4lL220cOj0q"
+      },
+      "source": [
+        "We can copy the batched `mjx.Data` back to MuJoCo like we did before:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jtz7j1PDOnw5"
+      },
+      "outputs": [],
+      "source": [
+        "batched_mj_data = mjx.get_data(mj_model, batch)\n",
+        "print([d.qpos for d in batched_mj_data])"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -187,145 +397,10 @@
       },
       "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 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 environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "7DQ_rW4CkIB_"
-      },
-      "outputs": [],
-      "source": [
-        "#@title State\n",
+        "Running large batch physics simulation is useful for training RL policies. Here we demonstrate training RL policies with MJX using the RL library from [Brax](https://github.com/google/brax).\n",
         "\n",
-        "@struct.dataclass\n",
-        "class State(Base):\n",
-        "  \"\"\"Environment state for training and inference with brax.\n",
-        "\n",
-        "  Args:\n",
-        "    pipeline_state: the physics state, mjx.Data\n",
-        "    obs: environment observations\n",
-        "    reward: environment reward\n",
-        "    done: boolean, True if the current episode has terminated\n",
-        "    metrics: metrics that get tracked per environment step\n",
-        "    info: environment variables defined and updated by the environment reset\n",
-        "      and step functions\n",
-        "  \"\"\"\n",
-        "\n",
-        "  pipeline_state: mjx.Data\n",
-        "  obs: jax.Array\n",
-        "  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"
-      ]
-    },
-    {
-      "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",
-        "\n",
-        "  def __init__(\n",
-        "      self,\n",
-        "      mj_model: mujoco.MjModel,\n",
-        "      physics_steps_per_control_step: int = 1,\n",
-        "  ):\n",
-        "    \"\"\"Initializes MjxEnv.\n",
-        "\n",
-        "    Args:\n",
-        "      mj_model: mujoco.MjModel\n",
-        "      physics_steps_per_control_step: the number of times to step the physics\n",
-        "        pipeline for each environment step\n",
-        "    \"\"\"\n",
-        "    self.model = mj_model\n",
-        "    self.data = mujoco.MjData(mj_model)\n",
-        "    self.sys = mjx.device_put(mj_model)\n",
-        "    self._physics_steps_per_control_step = physics_steps_per_control_step\n",
-        "\n",
-        "  def pipeline_init(\n",
-        "      self, qpos: jax.Array, qvel: jax.Array\n",
-        "  ) -\u003e mjx.Data:\n",
-        "    \"\"\"Initializes the physics state.\"\"\"\n",
-        "    data = mjx.device_put(self.data)\n",
-        "    data = data.replace(qpos=qpos, qvel=qvel, ctrl=jp.zeros(self.sys.nu))\n",
-        "    data = mjx.forward(self.sys, data)\n",
-        "    return data\n",
-        "\n",
-        "  def pipeline_step(\n",
-        "      self, data: mjx.Data, ctrl: jax.Array\n",
-        "  ) -\u003e mjx.Data:\n",
-        "    \"\"\"Takes a physics step using the physics pipeline.\"\"\"\n",
-        "    def f(data, _):\n",
-        "      data = data.replace(ctrl=ctrl)\n",
-        "      return (\n",
-        "          mjx.step(self.sys, data),\n",
-        "          None,\n",
-        "      )\n",
-        "    data, _ = jax.lax.scan(f, data, (), self._physics_steps_per_control_step)\n",
-        "    return data\n",
-        "\n",
-        "  @property\n",
-        "  def dt(self) -\u003e jax.Array:\n",
-        "    \"\"\"The timestep used for each env step.\"\"\"\n",
-        "    return self.sys.opt.timestep * self._physics_steps_per_control_step\n",
-        "\n",
-        "  @property\n",
-        "  def observation_size(self) -\u003e int:\n",
-        "    rng = jax.random.PRNGKey(0)\n",
-        "    reset_state = self.unwrapped.reset(rng)\n",
-        "    return reset_state.obs.shape[-1]\n",
-        "\n",
-        "  @property\n",
-        "  def action_size(self) -\u003e int:\n",
-        "    return self.sys.nu\n",
-        "\n",
-        "  @property\n",
-        "  def backend(self) -\u003e str:\n",
-        "    return 'mjx'\n",
-        "\n",
-        "  def _pos_vel(\n",
-        "      self, data: mjx.Data\n",
-        "      ) -\u003e Tuple[Transform, Motion]:\n",
-        "    \"\"\"Returns 6d spatial transform and 6d velocity for all bodies.\"\"\"\n",
-        "    x = Transform(pos=data.xpos[1:, :], rot=data.xquat[1:, :])\n",
-        "    cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])\n",
-        "    offset = data.xpos[1:, :] - data.subtree_com[\n",
-        "        self.model.body_rootid[np.arange(1, self.model.nbody)]]\n",
-        "    xd = Transform.create(pos=offset).vmap().do(cvel)\n",
-        "    return x, xd\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "iPlFu4CiIgBN"
-      },
-      "source": [
-        "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."
+        "Below, we implement the classic Humanoid environment using MJX and Brax. We inherit from the `MjxEnv` implementation in Brax so that we can step the physics with MJX while training with Brax RL implementations.\n"
       ]
     },
     {
@@ -361,10 +436,10 @@
         "    mj_model.opt.ls_iterations = 6\n",
         "\n",
         "    physics_steps_per_control_step = 5\n",
-        "    kwargs['physics_steps_per_control_step'] = kwargs.get(\n",
-        "        'physics_steps_per_control_step', physics_steps_per_control_step)\n",
+        "    kwargs['n_frames'] = kwargs.get(\n",
+        "        'n_frames', physics_steps_per_control_step)\n",
         "\n",
-        "    super().__init__(mj_model=mj_model, **kwargs)\n",
+        "    super().__init__(model=mj_model, **kwargs)\n",
         "\n",
         "    self._forward_reward_weight = forward_reward_weight\n",
         "    self._ctrl_cost_weight = ctrl_cost_weight\n",
@@ -390,7 +465,7 @@
         "\n",
         "    data = self.pipeline_init(qpos, qvel)\n",
         "\n",
-        "    obs = self._get_obs(data, jp.zeros(self.sys.nu))\n",
+        "    obs = self._get_obs(data.data, jp.zeros(self.sys.nu))\n",
         "    reward, done, zero = jp.zeros(3)\n",
         "    metrics = {\n",
         "        'forward_reward': zero,\n",
@@ -410,14 +485,14 @@
         "    data0 = state.pipeline_state\n",
         "    data = self.pipeline_step(data0, action)\n",
         "\n",
-        "    com_before = data0.subtree_com[1]\n",
-        "    com_after = data.subtree_com[1]\n",
+        "    com_before = data0.data.subtree_com[1]\n",
+        "    com_after = data.data.subtree_com[1]\n",
         "    velocity = (com_after - com_before) / self.dt\n",
         "    forward_reward = self._forward_reward_weight * velocity[0]\n",
         "\n",
         "    min_z, max_z = self._healthy_z_range\n",
-        "    is_healthy = jp.where(data.qpos[2] \u003c min_z, 0.0, 1.0)\n",
-        "    is_healthy = jp.where(data.qpos[2] \u003e max_z, 0.0, is_healthy)\n",
+        "    is_healthy = jp.where(data.q[2] \u003c min_z, 0.0, 1.0)\n",
+        "    is_healthy = jp.where(data.q[2] \u003e max_z, 0.0, is_healthy)\n",
         "    if self._terminate_when_unhealthy:\n",
         "      healthy_reward = self._healthy_reward\n",
         "    else:\n",
@@ -425,7 +500,7 @@
         "\n",
         "    ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n",
         "\n",
-        "    obs = self._get_obs(data, action)\n",
+        "    obs = self._get_obs(data.data, action)\n",
         "    reward = forward_reward + healthy_reward - ctrl_cost\n",
         "    done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n",
         "    state.metrics.update(\n",
@@ -492,31 +567,7 @@
         "\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"
+        "jit_step = jax.jit(env.step)\n"
       ]
     },
     {
@@ -529,17 +580,15 @@
       "source": [
         "# initialize the state\n",
         "state = jit_reset(jax.random.PRNGKey(0))\n",
-        "rollout = [state]\n",
-        "images = [get_image(state, camera='side')]\n",
+        "rollout = [state.pipeline_state]\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",
+        "  rollout.append(state.pipeline_state)\n",
         "\n",
-        "media.show_video(images, fps=1.0 / env.dt)"
+        "media.show_video(env.render(rollout, camera='side'), fps=1.0 / env.dt)"
       ]
     },
     {
@@ -550,7 +599,7 @@
       "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."
+        "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 9-10 minutes on a Tesla A100 GPU."
       ]
     },
     {
@@ -604,7 +653,7 @@
         "id": "YYIch0HEApBx"
       },
       "source": [
-        "## Save and Load Policy\n",
+        "\u003c!-- ## Save and Load Policy --\u003e\n",
         "\n",
         "We can save and load the policy using the brax model API."
       ]
@@ -673,8 +722,7 @@
         "# 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",
+        "rollout = [state.pipeline_state]\n",
         "\n",
         "# grab a trajectory\n",
         "n_steps = 500\n",
@@ -684,14 +732,12 @@
         "  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",
+        "  rollout.append(state.pipeline_state)\n",
         "\n",
         "  if state.done:\n",
         "    break\n",
         "\n",
-        "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
+        "media.show_video(env.render(rollout[::render_every], camera='side'), fps=1.0 / env.dt / render_every)"
       ]
     },
     {
@@ -702,7 +748,7 @@
       "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."
+        "We can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo."
       ]
     },
     {
@@ -713,7 +759,7 @@
       },
       "outputs": [],
       "source": [
-        "mj_model = eval_env.model\n",
+        "mj_model = eval_env._model\n",
         "mj_data = mujoco.MjData(mj_model)\n",
         "\n",
         "renderer = mujoco.Renderer(mj_model)\n",
@@ -723,11 +769,11 @@
         "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",
+        "  obs = eval_env._get_obs(mjx.put_data(mj_model, 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",
+        "  for _ in range(eval_env._n_frames):\n",
         "    mujoco.mj_step(mj_model, mj_data)  # Physics step using MuJoCo mj_step.\n",
         "\n",
         "  if i % render_every == 0:\n",
@@ -743,7 +789,7 @@
         "id": "65mIPj6DQNNa"
       },
       "source": [
-        "# Domain Randomization\n",
+        "# Training a Policy with 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."
       ]
@@ -766,7 +812,7 @@
         "    friction = sys.geom_friction.at[:, 0].set(friction)\n",
         "    # actuator\n",
         "    _, key = jax.random.split(key, 2)\n",
-        "    gain_range = (-10, -5)\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",
@@ -798,7 +844,7 @@
         "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."
+        "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjx.Model` along with a dictionary specifying the axes that are batched."
       ]
     },
     {
@@ -828,7 +874,18 @@
       "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."
+        "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour vb Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_vb) from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie). We implement an environment that trains a joystick policy with Brax."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VfyK73gtRXid"
+      },
+      "outputs": [],
+      "source": [
+        "!git clone https://github.com/google-deepmind/mujoco_menagerie"
       ]
     },
     {
@@ -839,7 +896,7 @@
       },
       "outputs": [],
       "source": [
-        "#@title Barkour v0 Quadruped Env\n",
+        "#@title Barkour vb Quadruped Env\n",
         "\n",
         "def get_config():\n",
         "  \"\"\"Returns reward config for barkour quadruped environment.\"\"\"\n",
@@ -869,11 +926,10 @@
         "                    # 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.002,\n",
+        "                    torques=-0.0002,\n",
         "                    # Penalize the change in the action and encourage smooth\n",
         "                    # actions. L2 regularization |action - last_action|^2\n",
-        "                    action_rate=-0.1,\n",
+        "                    action_rate=-0.01,\n",
         "                    # Encourage long swing steps.  However, it does not\n",
         "                    # encourage high clearances.\n",
         "                    feet_air_time=0.2,\n",
@@ -893,7 +949,10 @@
         "    return default_config\n",
         "\n",
         "  default_config = config_dict.ConfigDict(\n",
-        "      dict(rewards=get_default_rewards_config(),))\n",
+        "      dict(\n",
+        "          rewards=get_default_rewards_config(),\n",
+        "      )\n",
+        "  )\n",
         "\n",
         "  return default_config\n",
         "\n",
@@ -904,46 +963,69 @@
         "  def __init__(\n",
         "      self,\n",
         "      obs_noise: float = 0.05,\n",
-        "      action_scale: float=0.3,\n",
+        "      action_scale: float = 0.3,\n",
+        "      kick_vel: float = 0.05,\n",
         "      **kwargs,\n",
         "  ):\n",
-        "    path = epath.Path(epath.resource_path('mujoco')) / (\n",
-        "        'mjx/benchmark/model/barkour_v0/assets'\n",
-        "    )\n",
-        "    mj_model = mujoco.MjModel.from_xml_path(\n",
-        "        (path / 'barkour_v0_mjx.xml').as_posix())\n",
-        "    mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
-        "    mj_model.opt.iterations = 4\n",
-        "    mj_model.opt.ls_iterations = 6\n",
+        "    path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n",
+        "    self._dt = 0.02  # this environment is 50 fps\n",
+        "    self.brax_sys = mjcf.load(path).replace(dt=self._dt)\n",
+        "    model = self.brax_sys.get_model()\n",
+        "    model.opt.timestep = 0.004\n",
         "\n",
-        "    physics_steps_per_control_step = 10\n",
-        "    kwargs['physics_steps_per_control_step'] = kwargs.get(\n",
-        "        'physics_steps_per_control_step', physics_steps_per_control_step)\n",
-        "    super().__init__(mj_model=mj_model, **kwargs)\n",
+        "    # override menagerie params for smoother policy\n",
+        "    model.dof_damping[6:] = 0.5239\n",
+        "    model.actuator_gainprm[:, 0] = 35.0\n",
+        "    model.actuator_biasprm[:, 1] = -35.0\n",
         "\n",
-        "    self.torso_idx = mujoco.mj_name2id(\n",
-        "        mj_model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
+        "    n_frames = kwargs.pop('n_frames', int(self._dt / model.opt.timestep))\n",
+        "    super().__init__(model=model, n_frames=n_frames)\n",
+        "\n",
+        "    self.reward_config = get_config()\n",
+        "    # set custom from kwargs\n",
+        "    for k, v in kwargs.items():\n",
+        "      if k.endswith('_scale'):\n",
+        "        self.reward_config.rewards.scales[k[:-6]] = v\n",
+        "\n",
+        "    self._torso_idx = mujoco.mj_name2id(\n",
+        "        model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
         "    )\n",
         "    self._action_scale = action_scale\n",
         "    self._obs_noise = obs_noise\n",
-        "    self._reset_horizon = 500\n",
-        "    self._feet_index = jp.array([3, 6, 9, 12])\n",
-        "    # local positions for each foot\n",
-        "    self._feet_pos = jp.array([\n",
-        "        [-0.191284, -0.0191638, 0.013],\n",
-        "        [-0.191284, -0.0191638, -0.013],\n",
-        "        [-0.191284, -0.0191638, 0.013],\n",
-        "        [-0.191284, -0.0191638, -0.013],\n",
-        "    ])\n",
-        "    self._init_q = mj_model.keyframe('standing').qpos\n",
-        "    self._default_ap_pose = mj_model.keyframe('standing').qpos[7:]\n",
-        "    self.reward_config = get_config()\n",
-        "    self.lowers = self._default_ap_pose - jp.array([0.2, 0.8, 0.8] * 4)\n",
-        "    self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n",
-        "    self._foot_radius = 0.014\n",
+        "    self._kick_vel = kick_vel\n",
+        "    self._init_q = jp.array(model.keyframe('home').qpos)\n",
+        "    self._default_pose = model.keyframe('home').qpos[7:]\n",
+        "    self.lowers = jp.array([-0.7, -1.0, 0.05] * 4)\n",
+        "    self.uppers = jp.array([0.52, 2.1, 2.1] * 4)\n",
+        "    feet_site = [\n",
+        "        'foot_front_left',\n",
+        "        'foot_hind_left',\n",
+        "        'foot_front_right',\n",
+        "        'foot_hind_right',\n",
+        "    ]\n",
+        "    feet_site_id = [\n",
+        "        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n",
+        "        for f in feet_site\n",
+        "    ]\n",
+        "    assert not any(id_ == -1 for id_ in feet_site_id), 'Site not found.'\n",
+        "    self._feet_site_id = np.array(feet_site_id)\n",
+        "    lower_leg_body = [\n",
+        "        'lower_leg_front_left',\n",
+        "        'lower_leg_hind_left',\n",
+        "        'lower_leg_front_right',\n",
+        "        'lower_leg_hind_right',\n",
+        "    ]\n",
+        "    lower_leg_body_id = [\n",
+        "        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n",
+        "        for l in lower_leg_body\n",
+        "    ]\n",
+        "    assert not any(id_ == -1 for id_ in lower_leg_body_id), 'Body not found.'\n",
+        "    self._lower_leg_body_id = np.array(lower_leg_body_id)\n",
+        "    self._foot_radius = 0.0175\n",
+        "    self._nv = model.nv\n",
         "\n",
         "  def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n",
-        "    lin_vel_x = [-0.6, 1.0]  # min max [m/s]\n",
+        "    lin_vel_x = [-0.6, 1.5]  # 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",
@@ -960,216 +1042,156 @@
         "    new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n",
         "    return new_cmd\n",
         "\n",
-        "  def reset(self, rng: jax.Array) -\u003e State:\n",
+        "  def reset(self, rng: jax.Array) -\u003e State:  # pytype: disable=signature-mismatch\n",
         "    rng, key = jax.random.split(rng)\n",
         "\n",
-        "    qpos = jp.array(self._init_q)\n",
-        "    qvel = jp.zeros(self.model.nv)\n",
-        "    new_cmd = self.sample_command(key)\n",
-        "    data = self.pipeline_init(qpos, qvel)\n",
+        "    pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n",
         "\n",
         "    state_info = {\n",
         "        'rng': rng,\n",
         "        'last_act': jp.zeros(12),\n",
         "        'last_vel': jp.zeros(12),\n",
-        "        'last_contact_buffer': jp.zeros((20, 4), dtype=bool),\n",
-        "        'command': new_cmd,\n",
+        "        'command': self.sample_command(key),\n",
         "        'last_contact': jp.zeros(4, dtype=bool),\n",
         "        'feet_air_time': jp.zeros(4),\n",
-        "        'obs_history': jp.zeros(15 * 31),\n",
-        "        'reward_tuple': {\n",
-        "            'tracking_lin_vel': 0.0,\n",
-        "            'tracking_ang_vel': 0.0,\n",
-        "            'lin_vel_z': 0.0,\n",
-        "            'ang_vel_xy': 0.0,\n",
-        "            'orientation': 0.0,\n",
-        "            'torque': 0.0,\n",
-        "            'action_rate': 0.0,\n",
-        "            'stand_still': 0.0,\n",
-        "            'feet_air_time': 0.0,\n",
-        "            'foot_slip': 0.0,\n",
-        "        },\n",
+        "        'rewards': {k: 0.0 for k in self.reward_config.rewards.scales.keys()},\n",
+        "        'kick': jp.array([0.0, 0.0]),\n",
         "        'step': 0,\n",
         "    }\n",
         "\n",
-        "    x, xd = self._pos_vel(data)\n",
-        "    obs = self._get_obs(data.qpos, x, xd, state_info)\n",
+        "    obs_history = jp.zeros(15 * 31)  # store 15 steps of history\n",
+        "    obs = self._get_obs(pipeline_state, state_info, obs_history)\n",
         "    reward, done = jp.zeros(2)\n",
         "    metrics = {'total_dist': 0.0}\n",
-        "    for k in state_info['reward_tuple']:\n",
-        "      metrics[k] = state_info['reward_tuple'][k]\n",
-        "    state = State(data, obs, reward, done, metrics, state_info)\n",
+        "    for k in state_info['rewards']:\n",
+        "      metrics[k] = state_info['rewards'][k]\n",
+        "    state = State(pipeline_state, obs, reward, done, metrics, state_info)  # pytype: disable=wrong-arg-types\n",
         "    return state\n",
         "\n",
-        "  def step(self, state: State, action: jax.Array) -\u003e State:\n",
-        "    rng, rng_noise, cmd_rng = jax.random.split(\n",
-        "        state.info['rng'], 3\n",
-        "    )\n",
+        "  def step(self, state: State, action: jax.Array) -\u003e State:  # pytype: disable=signature-mismatch\n",
+        "    rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n",
+        "\n",
+        "    # kick\n",
+        "    push_interval = 10\n",
+        "    kick_theta = jax.random.uniform(kick_noise_2, maxval=2 * jp.pi)\n",
+        "    kick = jp.array([jp.cos(kick_theta), jp.sin(kick_theta)])\n",
+        "    kick *= jp.mod(state.info['step'], push_interval) == 0\n",
+        "    qvel = state.pipeline_state.data.qvel  # pytype: disable=attribute-error\n",
+        "    qvel = qvel.at[:2].set(kick * self._kick_vel + qvel[:2])\n",
+        "    state = state.tree_replace({'pipeline_state.data.qvel': qvel})\n",
         "\n",
         "    # physics step\n",
-        "    cur_action = jp.array(action)\n",
-        "    action = action[:12] * self._action_scale\n",
-        "    motor_targets = jp.clip(\n",
-        "        action + self._default_ap_pose, self.lowers, self.uppers\n",
-        "    )\n",
-        "    data = self.pipeline_step(state.pipeline_state, motor_targets)\n",
+        "    motor_targets = self._default_pose + action * self._action_scale\n",
+        "    motor_targets = jp.clip(motor_targets, self.lowers, self.uppers)\n",
+        "    pipeline_state = self.pipeline_step(state.pipeline_state, motor_targets)\n",
+        "    x, xd = pipeline_state.x, pipeline_state.xd\n",
         "\n",
         "    # observation data\n",
-        "    x, xd = self._pos_vel(data)\n",
-        "    obs = self._get_obs(data.qpos, x, xd, state.info)\n",
-        "    obs_noise = self._obs_noise * jax.random.uniform(\n",
-        "        rng_noise, obs.shape, minval=-1, maxval=1)\n",
-        "    qpos, qvel = data.qpos, data.qvel\n",
-        "    joint_angles = qpos[7:]\n",
-        "    joint_vel = qvel[6:]\n",
+        "    obs = self._get_obs(pipeline_state, state.info, state.obs)\n",
+        "    joint_angles = pipeline_state.q[7:]\n",
+        "    joint_vel = pipeline_state.qd[6:]\n",
         "\n",
         "    # foot contact data based on z-position\n",
-        "    foot_contact_pos = (\n",
-        "        self._get_feet_pos_vel(x, xd)[0][:, 2]\n",
-        "        - self._foot_radius\n",
-        "    )\n",
-        "    contact = foot_contact_pos \u003c 1e-3  # a mm or less off the floor\n",
-        "    contact_filt_mm = jp.logical_or(contact, state.info['last_contact'])\n",
-        "    contact_filt_cm = jp.logical_or(\n",
-        "        foot_contact_pos \u003c 3e-2, state.info['last_contact']\n",
-        "    )  # 3cm or less off the floor\n",
-        "    first_contact = (state.info['feet_air_time'] \u003e 0) * (contact_filt_mm)\n",
+        "    foot_pos = pipeline_state.data.site_xpos[self._feet_site_id]  # pytype: disable=attribute-error\n",
+        "    foot_contact_z = foot_pos[:, 2] - self._foot_radius\n",
+        "    contact = foot_contact_z \u003c 1e-3  # a mm or less off the floor\n",
+        "    contact_filt_mm = contact | state.info['last_contact']\n",
+        "    contact_filt_cm = (foot_contact_z \u003c 3e-2) | state.info['last_contact']\n",
+        "    first_contact = (state.info['feet_air_time'] \u003e 0) * contact_filt_mm\n",
         "    state.info['feet_air_time'] += self.dt\n",
         "\n",
+        "    # done if joint limits are reached or robot is falling\n",
+        "    up = jp.array([0.0, 0.0, 1.0])\n",
+        "    done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) \u003c 0\n",
+        "    done |= jp.any(joint_angles \u003c self.lowers)\n",
+        "    done |= jp.any(joint_angles \u003e self.uppers)\n",
+        "    done |= pipeline_state.x.pos[self._torso_idx - 1, 2] \u003c 0.18\n",
+        "\n",
         "    # reward\n",
-        "    reward_tuple = {\n",
+        "    rewards = {\n",
         "        'tracking_lin_vel': (\n",
         "            self._reward_tracking_lin_vel(state.info['command'], x, xd)\n",
-        "            * self.reward_config.rewards.scales.tracking_lin_vel\n",
         "        ),\n",
         "        'tracking_ang_vel': (\n",
         "            self._reward_tracking_ang_vel(state.info['command'], x, xd)\n",
-        "            * self.reward_config.rewards.scales.tracking_ang_vel\n",
         "        ),\n",
-        "        'lin_vel_z': (\n",
-        "            self._reward_lin_vel_z(xd)\n",
-        "            * self.reward_config.rewards.scales.lin_vel_z\n",
+        "        'lin_vel_z': self._reward_lin_vel_z(xd),\n",
+        "        'ang_vel_xy': self._reward_ang_vel_xy(xd),\n",
+        "        'orientation': self._reward_orientation(x),\n",
+        "        'torques': self._reward_torques(pipeline_state.data.qfrc_actuator),  # pytype: disable=attribute-error\n",
+        "        'action_rate': self._reward_action_rate(action, state.info['last_act']),\n",
+        "        'stand_still': self._reward_stand_still(\n",
+        "            state.info['command'], joint_angles,\n",
         "        ),\n",
-        "        'ang_vel_xy': (\n",
-        "            self._reward_ang_vel_xy(xd)\n",
-        "            * self.reward_config.rewards.scales.ang_vel_xy\n",
-        "        ),\n",
-        "        'orientation': (\n",
-        "            self._reward_orientation(x)\n",
-        "            * self.reward_config.rewards.scales.orientation\n",
-        "        ),\n",
-        "        'torque': (\n",
-        "            self._reward_torques(data.qfrc_actuator)\n",
-        "            * self.reward_config.rewards.scales.torques\n",
-        "        ),\n",
-        "        'action_rate': (\n",
-        "            self._reward_action_rate(cur_action, state.info['last_act'])\n",
-        "            * self.reward_config.rewards.scales.action_rate\n",
-        "        ),\n",
-        "        'stand_still': (\n",
-        "            self._reward_stand_still(\n",
-        "                state.info['command'], joint_angles, self._default_ap_pose\n",
-        "            )\n",
-        "            * self.reward_config.rewards.scales.stand_still\n",
-        "        ),\n",
-        "        'feet_air_time': (\n",
-        "            self._reward_feet_air_time(\n",
-        "                state.info['feet_air_time'],\n",
-        "                first_contact,\n",
-        "                state.info['command'],\n",
-        "            )\n",
-        "            * self.reward_config.rewards.scales.feet_air_time\n",
-        "        ),\n",
-        "        'foot_slip': (\n",
-        "            self._reward_foot_slip(x, xd, contact_filt_cm)\n",
-        "            * self.reward_config.rewards.scales.foot_slip\n",
+        "        'feet_air_time': self._reward_feet_air_time(\n",
+        "            state.info['feet_air_time'],\n",
+        "            first_contact,\n",
+        "            state.info['command'],\n",
         "        ),\n",
+        "        'foot_slip': self._reward_foot_slip(pipeline_state, contact_filt_cm),\n",
+        "        'termination': self._reward_termination(done, state.info['step']),\n",
         "    }\n",
-        "    reward = sum(reward_tuple.values())\n",
-        "    reward = jp.clip(reward * self.dt, 0.0, 10000.0)\n",
+        "    rewards = {\n",
+        "        k: v * self.reward_config.rewards.scales[k] for k, v in rewards.items()\n",
+        "    }\n",
+        "    reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)\n",
         "\n",
         "    # state management\n",
-        "    state.info['last_act'] = cur_action\n",
+        "    state.info['kick'] = kick\n",
+        "    state.info['last_act'] = action\n",
         "    state.info['last_vel'] = joint_vel\n",
         "    state.info['feet_air_time'] *= ~contact_filt_mm\n",
         "    state.info['last_contact'] = contact\n",
-        "    state.info['last_contact_buffer'] = jp.roll(\n",
-        "        state.info['last_contact_buffer'], 1, axis=0\n",
-        "    )\n",
-        "    state.info['last_contact_buffer'] = (\n",
-        "        state.info['last_contact_buffer'].at[0].set(contact)\n",
-        "    )\n",
-        "    state.info['reward_tuple'] = reward_tuple\n",
+        "    state.info['rewards'] = rewards\n",
         "    state.info['step'] += 1\n",
-        "    state.info.update(rng=rng)\n",
+        "    state.info['rng'] = rng\n",
         "\n",
-        "    # resetting logic if joint limits are reached or robot is falling\n",
-        "    up = jp.array([0.0, 0.0, 1.0])\n",
-        "    done = jp.dot(math.rotate(up, x.rot[0]), up) \u003c 0\n",
-        "    done |= jp.any(joint_angles \u003c 0.98 * self.lowers)\n",
-        "    done |= jp.any(joint_angles \u003e 0.98 * self.uppers)\n",
-        "    done |= x.pos[0, 2] \u003c 0.18\n",
-        "\n",
-        "    # termination reward\n",
-        "    reward += (\n",
-        "        done  * (state.info['step'] \u003c self._reset_horizon) *\n",
-        "        self.reward_config.rewards.scales.termination\n",
-        "    )\n",
-        "\n",
-        "    # when done, sample new command if more than _reset_horizon timesteps\n",
-        "    # achieved\n",
+        "    # sample new command if more than 500 timesteps achieved\n",
         "    state.info['command'] = jp.where(\n",
-        "        done \u0026 (state.info['step'] \u003e self._reset_horizon),\n",
-        "        self.sample_command(cmd_rng), state.info['command'])\n",
+        "        state.info['step'] \u003e 500,\n",
+        "        self.sample_command(cmd_rng),\n",
+        "        state.info['command'],\n",
+        "    )\n",
         "    # reset the step counter when done\n",
         "    state.info['step'] = jp.where(\n",
-        "        done | (state.info['step'] \u003e self._reset_horizon), 0,\n",
-        "        state.info['step']\n",
+        "        done | (state.info['step'] \u003e 500), 0, state.info['step']\n",
         "    )\n",
         "\n",
         "    # log total displacement as a proxy metric\n",
-        "    state.metrics['total_dist'] = math.normalize(x.pos[self.torso_idx])[1]\n",
-        "    for k in state.info['reward_tuple'].keys():\n",
-        "      state.metrics[k] = state.info['reward_tuple'][k]\n",
+        "    state.metrics['total_dist'] = math.normalize(x.pos[self._torso_idx - 1])[1]\n",
+        "    state.metrics.update(state.info['rewards'])\n",
         "\n",
+        "    done = jp.float32(done)\n",
         "    state = state.replace(\n",
-        "        pipeline_state=data, obs=obs + obs_noise, reward=reward,\n",
-        "        done=done * 1.0)\n",
+        "        pipeline_state=pipeline_state, obs=obs, reward=reward, done=done\n",
+        "    )\n",
         "    return state\n",
         "\n",
-        "  def _get_obs(self, qpos: jax.Array, x: Transform, xd: Motion,\n",
-        "               state_info: Dict[str, Any]) -\u003e jax.Array:\n",
-        "    # Get observations:\n",
-        "    # yaw_rate,  projected_gravity, command,  motor_angles, last_action\n",
+        "  def _get_obs(\n",
+        "      self,\n",
+        "      pipeline_state: base.State,\n",
+        "      state_info: dict[str, Any],\n",
+        "      obs_history: jax.Array,\n",
+        "  ) -\u003e jax.Array:\n",
+        "    inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n",
+        "    local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n",
         "\n",
-        "    inv_base_orientation = math.quat_inv(x.rot[0])\n",
-        "    local_rpyrate = math.rotate(xd.ang[0], inv_base_orientation)\n",
-        "    cmd = state_info['command']\n",
+        "    obs = jp.concatenate([\n",
+        "        jp.array([local_rpyrate[2]]) * 0.25,                 # yaw rate\n",
+        "        math.rotate(jp.array([0, 0, -1]), inv_torso_rot),    # projected gravity\n",
+        "        state_info['command'] * jp.array([2.0, 2.0, 0.25]),  # command\n",
+        "        pipeline_state.q[7:] - self._default_pose,           # motor angles\n",
+        "        state_info['last_act'],                              # last action\n",
+        "    ])\n",
         "\n",
-        "    obs_list = []\n",
-        "    # yaw rate\n",
-        "    obs_list.append(jp.array([local_rpyrate[2]]) * 0.25)\n",
-        "    # projected gravity\n",
-        "    obs_list.append(\n",
-        "        math.rotate(jp.array([0.0, 0.0, -1.0]), inv_base_orientation))\n",
-        "    # command\n",
-        "    obs_list.append(cmd * jp.array([2.0, 2.0, 0.25]))\n",
-        "    # motor angles\n",
-        "    angles = qpos[7:19]\n",
-        "    obs_list.append(angles - self._default_ap_pose)\n",
-        "    # last action\n",
-        "    obs_list.append(state_info['last_act'])\n",
-        "\n",
-        "    obs = jp.clip(jp.concatenate(obs_list), -100.0, 100.0)\n",
-        "\n",
-        "    # stack observations through time\n",
-        "    single_obs_size = len(obs)\n",
-        "    state_info['obs_history'] = jp.roll(\n",
-        "        state_info['obs_history'], single_obs_size\n",
+        "    # clip, noise\n",
+        "    obs = jp.clip(obs, -100.0, 100.0) + self._obs_noise * jax.random.uniform(\n",
+        "        state_info['rng'], obs.shape, minval=-1, maxval=1\n",
         "    )\n",
-        "    state_info['obs_history'] = jp.array(\n",
-        "        state_info['obs_history']).at[:single_obs_size].set(obs)\n",
-        "    return state_info['obs_history']\n",
+        "    # stack observations through time\n",
+        "    obs = jp.roll(obs_history, obs.size).at[:obs.size].set(obs)\n",
+        "\n",
+        "    return obs\n",
         "\n",
         "  # ------------ reward functions----------------\n",
         "  def _reward_lin_vel_z(self, xd: Motion) -\u003e jax.Array:\n",
@@ -1191,12 +1213,14 @@
         "    return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n",
         "\n",
         "  def _reward_action_rate(\n",
-        "      self, act: jax.Array, last_act: jax.Array) -\u003e jax.Array:\n",
+        "      self, act: jax.Array, last_act: jax.Array\n",
+        "  ) -\u003e jax.Array:\n",
         "    # Penalize changes in actions\n",
         "    return jp.sum(jp.square(act - last_act))\n",
         "\n",
         "  def _reward_tracking_lin_vel(\n",
-        "      self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n",
+        "      self, commands: jax.Array, x: Transform, xd: Motion\n",
+        "  ) -\u003e jax.Array:\n",
         "    # Tracking of linear velocity commands (xy axes)\n",
         "    local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\n",
         "    lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n",
@@ -1206,15 +1230,16 @@
         "    return lin_vel_reward\n",
         "\n",
         "  def _reward_tracking_ang_vel(\n",
-        "      self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n",
+        "      self, commands: jax.Array, x: Transform, xd: Motion\n",
+        "  ) -\u003e jax.Array:\n",
         "    # Tracking of angular velocity commands (yaw)\n",
         "    base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\n",
         "    ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n",
-        "    return jp.exp(-ang_vel_error/self.reward_config.rewards.tracking_sigma)\n",
+        "    return jp.exp(-ang_vel_error / self.reward_config.rewards.tracking_sigma)\n",
         "\n",
         "  def _reward_feet_air_time(\n",
-        "      self, air_time: jax.Array, first_contact: jax.Array,\n",
-        "      commands: jax.Array) -\u003e jax.Array:\n",
+        "      self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n",
+        "  ) -\u003e jax.Array:\n",
         "    # Reward air time.\n",
         "    rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n",
         "    rew_air_time *= (\n",
@@ -1223,30 +1248,38 @@
         "    return rew_air_time\n",
         "\n",
         "  def _reward_stand_still(\n",
-        "      self, commands: jax.Array, joint_angles: jax.Array,\n",
-        "      default_angles: jax.Array) -\u003e jax.Array:\n",
+        "      self,\n",
+        "      commands: jax.Array,\n",
+        "      joint_angles: jax.Array,\n",
+        "  ) -\u003e jax.Array:\n",
         "    # Penalize motion at zero commands\n",
-        "    return jp.sum(jp.abs(joint_angles - default_angles)) * (\n",
+        "    return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n",
         "        math.normalize(commands[:2])[1] \u003c 0.1\n",
         "    )\n",
         "\n",
-        "  def _get_feet_pos_vel(\n",
-        "      self, x: Transform, xd: Motion) -\u003e Tuple[jax.Array, jax.Array]:\n",
-        "    offset = Transform.create(pos=self._feet_pos)\n",
-        "    pos = x.take(self._feet_index).vmap().do(offset).pos\n",
-        "    world_offset = Transform.create(pos=pos - x.take(self._feet_index).pos)\n",
-        "    vel = world_offset.vmap().do(xd.take(self._feet_index)).vel\n",
-        "    return pos, vel\n",
-        "\n",
         "  def _reward_foot_slip(\n",
-        "      self, x: Transform, xd: Motion, contact_filt: jax.Array) -\u003e jax.Array:\n",
-        "    # Get feet velocities\n",
-        "    _, foot_world_vel = self._get_feet_pos_vel(x, xd)\n",
-        "    # Penalize large feet velocity for feet that are in contact with the ground.\n",
-        "    return jp.sum(\n",
-        "        jp.square(foot_world_vel[:, :2]) * contact_filt.reshape((-1, 1))\n",
-        "    )\n",
+        "      self, pipeline_state: base.State, contact_filt: jax.Array\n",
+        "  ) -\u003e jax.Array:\n",
+        "    # get velocities at feet which are offset from lower legs\n",
+        "    # pytype: disable=attribute-error\n",
+        "    pos = pipeline_state.data.site_xpos[self._feet_site_id]  # feet position\n",
+        "    feet_offset = pos - pipeline_state.data.xpos[self._lower_leg_body_id]\n",
+        "    # pytype: enable=attribute-error\n",
+        "    offset = base.Transform.create(pos=feet_offset)\n",
+        "    foot_indices = self._lower_leg_body_id - 1  # we got rid of the world body\n",
+        "    foot_vel = offset.vmap().do(pipeline_state.xd.take(foot_indices)).vel\n",
         "\n",
+        "    # Penalize large feet velocity for feet that are in contact with the ground.\n",
+        "    return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n",
+        "\n",
+        "  def _reward_termination(self, done: jax.Array, step: jax.Array) -\u003e jax.Array:\n",
+        "    return done \u0026 (step \u003c 500)\n",
+        "\n",
+        "  def render(\n",
+        "      self, trajectory: List[base.State], camera: str | None = None\n",
+        "  ) -\u003e Sequence[np.ndarray]:\n",
+        "    camera = camera or 'track'\n",
+        "    return super().render(trajectory, camera)\n",
         "\n",
         "envs.register_environment('barkour', BarkourEnv)"
       ]
@@ -1260,10 +1293,7 @@
       "outputs": [],
       "source": [
         "env_name = 'barkour'\n",
-        "env = envs.get_environment(env_name)\n",
-        "\n",
-        "# re-instantiate the renderer\n",
-        "renderer = mujoco.Renderer(env.model)"
+        "env = envs.get_environment(env_name)"
       ]
     },
     {
@@ -1274,7 +1304,7 @@
       "source": [
         "## Train Policy\n",
         "\n",
-        "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."
+        "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  8-9 minutes on a Tesla A100 GPU."
       ]
     },
     {
@@ -1289,22 +1319,19 @@
         "    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",
+        "      ppo.train, num_timesteps=100_000_000, num_evals=10,\n",
+        "      reward_scaling=1, episode_length=1000, normalize_observations=True,\n",
+        "      action_repeat=1, unroll_length=20, num_minibatches=32,\n",
+        "      num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n",
+        "      entropy_cost=1e-2, num_envs=8192, batch_size=256,\n",
         "      network_factory=make_networks_factory,\n",
-        "      num_resets_per_eval=10,\n",
         "      randomization_fn=domain_randomize, seed=0)\n",
         "\n",
-        "\n",
         "x_data = []\n",
         "y_data = []\n",
         "ydataerr = []\n",
         "times = [datetime.now()]\n",
-        "max_y, min_y = 30, 0\n",
+        "max_y, min_y = 40, 0\n",
         "\n",
         "# Reset environments since internals may be overwritten by tracers from the\n",
         "# domain randomization function.\n",
@@ -1368,7 +1395,6 @@
       },
       "outputs": [],
       "source": [
-        "\n",
         "# @markdown Commands **only used for Barkour Env**:\n",
         "x_vel = 1.0  #@param {type: \"number\"}\n",
         "y_vel = 0.0  #@param {type: \"number\"}\n",
@@ -1380,8 +1406,7 @@
         "rng = jax.random.PRNGKey(0)\n",
         "state = jit_reset(rng)\n",
         "state.info['command'] = the_command\n",
-        "rollout = [state]\n",
-        "images = [get_image(state, camera='track')]\n",
+        "rollout = [state.pipeline_state]\n",
         "\n",
         "# grab a trajectory\n",
         "n_steps = 500\n",
@@ -1391,11 +1416,31 @@
         "  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='track'))\n",
+        "  rollout.append(state.pipeline_state)\n",
         "\n",
-        "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
+        "media.show_video(\n",
+        "    eval_env.render(rollout[::render_every], camera='track'),\n",
+        "    fps=1.0 / eval_env.dt / render_every)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aD6H6WD0915X"
+      },
+      "source": [
+        "We can also render the rollout using the Brax renderer."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "V7jqv08X95u4"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(html.render(eval_env.brax_sys, rollout))"
       ]
     }
   ],
@@ -1403,12 +1448,13 @@
     "accelerator": "GPU",
     "colab": {
       "gpuClass": "premium",
-      "gpuType": "V100",
+      "gpuType": "A100",
+      "machine_shape": "hm",
       "private_outputs": true,
       "provenance": [
         {
-          "file_id": "1QsuS7EJhdPEHxxAu9XwozvA7eb4ZnlAb",
-          "timestamp": 1701993737024
+          "file_id": "11cFRVCJ8Kn71tlQFbFcw4JzQZ00F8BRG",
+          "timestamp": 1704355889284
         }
       ],
       "toc_visible": true