Fix external import in colab.

PiperOrigin-RevId: 682514874
Change-Id: I6157c2613464d2f5d3c23a73248913f1a240a9f1

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": [],