Fix replace just computed stochastic state (#275)

Eclectic-Sheep · May 2, 2024 · 96040b1 · 96040b1
1 parent 2bae379
commit 96040b1
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/notebooks/dreamer_v3_imagination.ipynb b/notebooks/dreamer_v3_imagination.ipynb
@@ -337,11 +337,9 @@
     "            )[None]\n",
     "\n",
     "        # imagination step\n",
-    "        imagined_stochastic_state, recurrent_state = world_model.rssm.imagination(\n",
-    "            stochastic_state, recurrent_state, actions\n",
-    "        )\n",
+    "        stochastic_state, recurrent_state = world_model.rssm.imagination(stochastic_state, recurrent_state, actions)\n",
     "        # update current state\n",
-    "        imagined_latent_states = torch.cat((imagined_stochastic_state.view(1, 1, -1), recurrent_state), -1)\n",
+    "        imagined_latent_states = torch.cat((stochastic_state.view(1, 1, -1), recurrent_state), -1)\n",
     "        rec_obs = world_model.observation_model(imagined_latent_states)\n",
     "        step_data[\"rgb\"] = rec_obs[\"rgb\"].unsqueeze(0).detach().cpu().numpy()\n",
     "        step_data[\"actions\"] = actions.unsqueeze(0).detach().cpu().numpy()\n",