Fix basics and training_torch notebooks (#490)

* Update basics notebook to use opt.timestep field * Upgrade gym in training_torch notebook to fix env vectorization
google · Jun 7, 2024 · c89a3ad · c89a3ad
1 parent 2b95fe9
commit c89a3ad
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/notebooks/basics.ipynb b/notebooks/basics.ipynb
@@ -542,7 +542,8 @@
         "link = pendulum.link.replace(\n",
         "    constraint_stiffness=jp.repeat(stiffness, 3),\n",
         "    constraint_vel_damping=jp.zeros((3,)))\n",
-        "pendulum = pendulum.replace(dt=0.01 / substeps, link=link)\n",
+        "pendulum = pendulum.replace(link=link)\n",
+        "pendulum = pendulum.tree_replace({'opt.timestep': 0.01 / substeps})\n",
         "\n",
         "init_q = jp.array([initial_angle * jp.pi / 180, 0, 0])\n",
         "state = jax.jit(pipeline.init)(pendulum, init_q, jp.zeros(pendulum.qd_size()))\n",

diff --git a/notebooks/training_torch.ipynb b/notebooks/training_torch.ipynb
@@ -27,9 +27,20 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade gym --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
       "metadata": {
         "id": "GJhPpM5ZPrpq"
       },
+      "outputs": [],
       "source": [
         "#@title Import Brax and some helper modules\n",
         "from IPython.display import clear_output\n",