Minor update to colabs.

Unity-Technologies · Oct 4, 2023 · dfe9716 · dfe9716
1 parent 2a47595
commit dfe9716
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/colab/Colab_UnityEnvironment_1_Run.ipynb b/colab/Colab_UnityEnvironment_1_Run.ipynb
@@ -536,7 +536,7 @@
     "for index, obs_spec in enumerate(spec.observation_specs):\n",
     "  if len(obs_spec.shape) == 3:\n",
     "    print(\"Here is the first visual observation\")\n",
-    "    plt.imshow(np.rollaxis(decision_steps.obs[index][0, :, :, :], 0, 3))\n",
+    "    plt.imshow(np.moveaxis(decision_steps.obs[index][0, :, :, :], 0, -1))\n",
     "    plt.show()\n",
     "\n",
     "for index, obs_spec in enumerate(spec.observation_specs):\n",

diff --git a/colab/Colab_UnityEnvironment_3_SideChannel.ipynb b/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
@@ -326,7 +326,7 @@
     "\n",
     "print(\"Observation without changing the environment parameters\")\n",
     "decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
-    "plt.imshow(np.rollaxis(decision_steps.obs[0][0,:,:,:], 0, 3))\n",
+    "plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))\n",
     "plt.show()\n",
     "\n",
     "print(\"Increasing the dimensions of the grid from 5 to 7\")\n",
@@ -340,7 +340,7 @@
     "\n",
     "\n",
     "decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
-    "plt.imshow(np.rollaxis(decision_steps.obs[0][0,:,:,:], 0, 3))\n",
+    "plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))\n",
     "plt.show()\n",
     "\n",
     "\n",