Fix: Add missing warp-lang dependency to vision notebook

PiperOrigin-RevId: 884607018
Change-Id: I7afcf4b715430a46470e1b8c35da4d9e505d0a86

Author: btaba

learning/notebooks/vision.ipynb

@@ -42,6 +42,7 @@
       "outputs": [],
       "source": [
         "#@title Install pre-requisites\n",
+        "!pip install warp-lang\n",
         "!pip install mujoco\n",
         "!pip install mujoco_mjx\n",
         "!pip install brax"
@@ -640,11 +641,11 @@
         "make_inference_fn, params, metrics = train_fn(\n",
         "    environment=wrapped_env,\n",
         "    eval_env=wrapper.wrap_for_brax_training(\n",
-    "        eval_env,\n",
-    "        episode_length=episode_length,\n",
-    "        action_repeat=1\n",
-    "    ),\n",
-    "    wrap_env=False,\n",
+        "        eval_env,\n",
+        "        episode_length=episode_length,\n",
+        "        action_repeat=1\n",
+        "    ),\n",
+        "    wrap_env=False,\n",
         ")\n",
         "print(f\"time to jit: {times[1] - times[0]}\")\n",
         "print(f\"time to train: {times[-1] - times[1]}\")"