Add VBD predicted outputs to notebooks for inspection

Emerge-Lab · Nov 4, 2024 · 60abc1b · 60abc1b
1 parent 0e9b33b
commit 60abc1b
Show file tree

Hide file tree

Showing 6 changed files with 276 additions and 50 deletions.
diff --git a/notebooks/00_align_simulators_vbd.ipynb b/notebooks/00_align_simulators_vbd.ipynb
diff --git a/notebooks/01_features_deepdive.ipynb b/notebooks/01_features_deepdive.ipynb
diff --git a/notebooks/gpudrive_vbd_sample_11671609ebfa3185.pkl b/notebooks/gpudrive_vbd_sample_11671609ebfa3185.pkl
diff --git a/notebooks/waymax_vbd_sample_11671609ebfa3185.pkl b/notebooks/waymax_vbd_sample_11671609ebfa3185.pkl
diff --git a/vbd/model/VBD.py b/vbd/model/VBD.py
@@ -1,5 +1,6 @@
 import torch
-#import lightning.pytorch as pl
+
+# import lightning.pytorch as pl
 import pytorch_lightning as pl
 
 from .modules import Encoder, Denoiser, GoalPredictor
@@ -155,7 +156,11 @@ def forward(self, inputs, noised_actions_normalized, diffusion_step):
         return output_dict
 
     def forward_denoiser(
-        self, encoder_outputs, noised_actions_normalized, diffusion_step
+        self,
+        encoder_outputs,
+        noised_actions_normalized,
+        diffusion_step,
+        global_frame=True,
     ):
         """
         Forward pass of the denoiser module.
@@ -185,7 +190,7 @@ def forward_denoiser(
             current_states,
             denoised_actions,
             action_len=self.denoiser._action_len,
-            global_frame=True,
+            global_frame=global_frame,
         )
 
         return {

diff --git a/vbd/sim_agent/sim_actor.py b/vbd/sim_agent/sim_actor.py
@@ -519,7 +519,9 @@ def ibr_guidance(
         return denoiser_output, x_t_prev, guide_history
 
     ################### Denoising ###################
-    def step_denoiser(self, x_t: torch.Tensor, c: dict, t: int):
+    def step_denoiser(
+        self, x_t: torch.Tensor, c: dict, t: int, global_frame: bool = True
+    ):
         """
         Perform a denoising step to sample x_{t-1} ~ P[x_{t-1} | x_t, D(x_t, c, t)].
 
@@ -538,7 +540,10 @@ def step_denoiser(self, x_t: torch.Tensor, c: dict, t: int):
 
         # Denoise to reconstruct x_0 ~ D(x_t, c, t)
         denoiser_output = self.forward_denoiser(
-            encoder_outputs=c, noised_actions_normalized=x_t, diffusion_step=t
+            encoder_outputs=c,
+            noised_actions_normalized=x_t,
+            diffusion_step=t,
+            global_frame=global_frame,
         )
 
         x_0 = denoiser_output["denoised_actions_normalized"]
@@ -554,7 +559,13 @@ def step_denoiser(self, x_t: torch.Tensor, c: dict, t: int):
 
     @torch.no_grad()
     def sample_denoiser(
-        self, batch, num_samples=1, x_t=None, use_tqdm=True, **kwargs
+        self,
+        batch,
+        num_samples=1,
+        x_t=None,
+        use_tqdm=True,
+        global_frame=True,
+        **kwargs
     ):
         """
         Perform denoising inference on the given batch of data.
@@ -626,6 +637,7 @@ def sample_denoiser(
                     x_t=x_t,
                     c=encoder_outputs,
                     t=t,
+                    global_frame=global_frame,
                 )
                 guide = None