Merge pull request #9 from mohammadzainabbas/dev

Dev

Author: mohammadzainabbas

notebooks/demo_ppo_train.ipynb

@@ -159,18 +159,13 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "training_num_timesteps = [1_000, 1_000_000, 100_000_000]\n",
-        "vis_steps = [100, 150, 300]\n",
+        "training_num_timesteps = [1_000, 5_000_000, 400_000_000]\n",
         "\n",
-        "env_sys = []\n",
-        "rollouts = []\n",
+        "inference_fns = []\n",
         "\n",
         "for idx, num_timesteps in enumerate(training_num_timesteps):\n",
         "\tmake_inference_fn, params, times, xdata, ydata = train_ppo(num_timesteps, env_name)\n",
-        "\tinference_fn = make_inference_fn(params)\n",
-        "\tsys, rollout = visual_rollout(inference_fn, env_name, steps=vis_steps[idx], seed=SEED)\n",
-        "\tenv_sys.append(sys)\n",
-        "\trollouts.append(rollout)"
+        "\tinference_fns.append(make_inference_fn(params))"
       ]
     },
     {
@@ -181,6 +176,23 @@
         "#### Visualise learning"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "vis_steps = [300, 500, 750]\n",
+        "\n",
+        "env_sys = []\n",
+        "rollouts = []\n",
+        "\n",
+        "for idx, inference_fn in enumerate(inference_fns):\n",
+        "\tsys, rollout = visual_rollout(inference_fn, env_name, steps=vis_steps[idx], seed=SEED)\n",
+        "\tenv_sys.append(sys)\n",
+        "\trollouts.append(rollout)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,