From 0dde0ba4ba3986039502928b4cdc381a1f7b5ac3 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 4 Oct 2024 18:11:39 -0700
Subject: [PATCH] Fix external import in colab.

PiperOrigin-RevId: 682514874
Change-Id: I6157c2613464d2f5d3c23a73248913f1a240a9f1
---
 mjx/tutorial.ipynb | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index d462fa5e..e9b983de 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -186,6 +186,7 @@
     "from brax import envs\n",
     "from brax import math\n",
     "from brax.base import Base, Motion, Transform\n",
+    "from brax.base import State as PipelineState\n",
     "from brax.envs.base import Env, PipelineEnv, State\n",
     "from brax.mjx.base import State as MjxState\n",
     "from brax.training.agents.ppo import train as ppo\n",
@@ -1636,6 +1637,7 @@
    "cell_type": "code",
    "execution_count": 0,
    "metadata": {
+    "cellView": "form",
     "id": "AuU-9nUquEu4"
    },
    "outputs": [],