Bug fixes to get rid of large position values with vbd_online

examples/eval/notebooks/obtain_guidance_data.py

@@ -11,7 +11,7 @@
 
 if __name__ == "__main__":
 
-    GUIDANCE_MODE = "log_replay"
+    GUIDANCE_MODE = "vbd_online"
     DATASET = "data/processed/wosac/validation_json_100"  # Ensure VBD trajectory structures are in here
     SAVE_PATH = "examples/eval/figures_data/"
 
@@ -24,13 +24,13 @@
         add_reference_speed=True,
         add_reference_pos_xy=True,
         init_mode="wosac_train",
-        smoothen_trajectory=True,
+        smoothen_trajectory=False,
     )
     render_config = RenderConfig()
 
     train_loader = SceneDataLoader(
         root=DATASET,
-        batch_size=10,
+        batch_size=2,
         dataset_size=100,
         sample_with_replacement=False,
         shuffle=False,
@@ -61,6 +61,6 @@
     )
 
     np.save(
-        f"{SAVE_PATH}reference_{GUIDANCE_MODE}_smooth.npy", reference_traj_np
+        f"{SAVE_PATH}reference_{GUIDANCE_MODE}.npy", reference_traj_np
     )
     print(f"Saved reference trajectory for {GUIDANCE_MODE} mode.")

gpudrive/datatypes/trajectory.py

@@ -151,14 +151,14 @@ def __init__(
         self.vel_xy = vbd_traj_tensor[:, :, :, 3:5]
         self.ref_speed = self.comp_reference_speed()
         # Assumption: All timesteps are valid (correct)
-        self.valids = torch.ones_like(self.pos_x, dtype=torch.int32)
+        self.valids = vbd_traj_tensor[:, :, :, 5].unsqueeze(-1).to(
+            torch.int32
+        )
 
         self.mean_pos_xy = mean_pos_xy
         self.mean_x = mean_pos_xy[:, 0]
         self.mean_y = mean_pos_xy[:, 1]
 
-        self.demean_positions()
-
     @classmethod
     def from_tensor(
         cls,
@@ -195,8 +195,8 @@ def demean_positions(self):
         mean_y_reshaped = self.mean_y.view(-1, 1, 1)
 
         # Apply to x and y coordinates
-        self.pos_xy[..., 0] -= mean_x_reshaped
-        self.pos_xy[..., 1] -= mean_y_reshaped
+        self.pos_xy[..., 10:, 0] -= mean_x_reshaped
+        self.pos_xy[..., 10:, 1] -= mean_y_reshaped
 
         self.pos_x = self.pos_xy[..., 0].unsqueeze(-1)
         self.pos_y = self.pos_xy[..., 1].unsqueeze(-1)
@@ -224,9 +224,9 @@ def __init__(self, vbd_traj_tensor: torch.Tensor):
         self.vel_y = vbd_traj_tensor[:, :, :, 4].unsqueeze(-1)
         self.vel_xy = vbd_traj_tensor[:, :, :, 3:5]
         self.ref_speed = self.comp_reference_speed()
-        # Assumption: All timesteps are valid
-        # TODO: The first 10 timesteps come from the logs, so there may be invalid steps
-        self.valids = torch.ones_like(self.pos_x, dtype=torch.int32)
+        self.valids = vbd_traj_tensor[:, :, :, 5].unsqueeze(-1).to(
+            torch.int32
+        )
 
     @classmethod
     def from_tensor(

gpudrive/env/base_env.py

@@ -165,7 +165,15 @@ def _setup_environment_parameters(self):
         params = self._set_collision_behavior(params)
         params = self._set_road_reduction_params(params)
         params = self._set_goal_behavior(params)
-        self.init_steps = getattr(self.config, "init_steps", 0)
+
+        if self.config.guidance and self.config.guidance_mode == "vbd_online":
+            self.init_steps = max(self.config.init_steps, 10)
+            print(
+                f"\n[Note] Guidance mode '{self.config.guidance_mode}' requires at least 10 initialization steps to provide sufficient scene context for the diffusion model. Automatically setting simulator time to t = {self.init_steps}. \n"
+            )
+        else:
+            self.init_steps = getattr(self.config, "init_steps", 0)
+
         params.initSteps = self.init_steps
 
         return params

gpudrive/env/config.py

@@ -40,7 +40,7 @@ class EnvConfig:
     add_reference_pos_xy: bool = True  # (x, y) position time series
     add_reference_speed: bool = False  # speed time series
     add_reference_heading: bool = False  # heading time series
-    smoothen_trajectory: bool = True  # Filters out the trajectory
+    smoothen_trajectory: bool = False  # Filters out the trajectory
     guidance_pos_xy_radius: float = 1.0  # Tightness of the positions guidance
 
     # Maximum number of controlled agents in the scene

gpudrive/env/env_torch.py

@@ -36,6 +36,7 @@
 from gpudrive.integrations.vbd.data.utils import process_scenario_data
 
 from gpudrive.utils.preprocess import smooth_scenario
+import madrona_gpudrive
 
 
 class GPUDriveTorchEnv(GPUDriveGymEnv):
@@ -72,18 +73,16 @@ def __init__(
 
         # Get the initial data batch (set of traffic scenarios)
         self.data_batch = next(self.data_iterator)
-
+        
         assert self.num_worlds == len(
             self.data_batch
         ), f"Number of scenarios in data_batch ({len(self.data_batch)}) \
         should equal number of worlds ({self.num_worlds}). \
         \n Please check your data loader configuration."
-
+        
         # Initialize simulator
         self.sim = self._initialize_simulator(params, self.data_batch)
 
-        self.init_steps = self.config.init_steps
-
         # Controlled agents setup
         self.cont_agent_mask = self.get_controlled_agents_mask()
         self.max_agent_count = self.cont_agent_mask.shape[1]
@@ -157,16 +156,16 @@ def setup_guidance(self):
                 trajectory_tensor, self.backend, self.device
             )
         elif self.guidance_mode == "vbd_online":
-
-            self.init_steps = max(self.init_steps, 10)
-            print(
-                f"\n[Note] Guidance mode '{self.guidance_mode}' requires at least {self.init_steps} initialization steps to provide sufficient scene context for the diffusion model. Automatically setting simulator time to t = {self.init_steps}. \n"
-            )
-
             # Load pre-trained Versatile Behavior Diffusion (VBD) model
             self.vbd_model = self._load_vbd_model(
                 model_path=self.config.vbd_model_path
             )
+
+            self.init_steps = max(self.init_steps, 10)
+            print(
+                f"\n[Note] Guidance mode '{self.guidance_mode}' requires at least {self.init_steps} initialization steps to provide sufficient scene context for the diffusion model. Automatically setting simulator time to t = {self.init_steps}. \n"
+            )
+
             # Construct scene context dict for the VBD model
             scene_context = self.construct_context(init_steps=self.init_steps)
 
@@ -175,9 +174,66 @@ def setup_guidance(self):
             # Query the model online for the reference trajectory
             predicted = self.vbd_model.sample_denoiser(scene_context)
 
+            # Pad first 10 steps with logs
+            trajectory_tensor = self.sim.expert_trajectory_tensor()
+            log_trajectory = LogTrajectory.from_tensor(
+                trajectory_tensor,
+                self.num_worlds,
+                self.max_agent_count,
+                self.backend,
+                self.device,
+            )
+
+            reference_trajectory = torch.zeros(
+                self.num_worlds,
+                self.max_agent_count,
+                madrona_gpudrive.kTrajectoryLength,
+                6
+            )
+            reference_trajectory[:, :, :self.init_steps + 1, :2] = log_trajectory.pos_xy[
+                :, :, :self.init_steps + 1
+            ]
+            reference_trajectory[:, :, :self.init_steps + 1, 2] = log_trajectory.yaw[
+                :, :, :self.init_steps + 1, 0
+            ]
+            reference_trajectory[:, :, :self.init_steps + 1, 3:5] = log_trajectory.vel_xy[
+                :, :, :self.init_steps + 1
+            ]
+            reference_trajectory[:, :, :self.init_steps + 1, 5] = log_trajectory.valids[
+                :, :, :self.init_steps + 1, 0
+            ]
+
+            vbd_predictions=predicted["denoised_trajs"].to(self.device).detach()
+
+            # Get the world means
+            world_means = (
+                self.sim.world_means_tensor().to_torch()[:, :2].to(self.device)
+            )
+
+            # Add vbd predictions to the reference trajectory
+            for i in range(self.num_worlds):
+                # Get controlled agent indices for this world
+                valid_mask = (
+                    scene_context["agents_id"][i] >= 0
+                )
+                valid_world_indices = scene_context["agents_id"][i][valid_mask]
+
+                reference_trajectory[
+                    i, valid_world_indices, self.init_steps + 1:, :2
+                ] = vbd_predictions[i, valid_world_indices, :, :2] - world_means[
+                    i
+                ].view(1, 1, 2)
+                reference_trajectory[
+                    i, valid_world_indices, self.init_steps + 1:, 2:5
+                ] = vbd_predictions[i, valid_world_indices, :, 2:5]
+                reference_trajectory[
+                    i, valid_world_indices, self.init_steps + 1:, 5
+                ] = 1
+
+
             # Wrap predictions into a VBDTrajectoryOnline object
             self.reference_trajectory = VBDTrajectoryOnline.from_tensor(
-                vbd_predictions=predicted["denoised_trajs"],
+                vbd_predictions=reference_trajectory,
                 mean_pos_xy=self.sim.world_means_tensor().to_torch()[:, :2],
                 backend=self.backend,
                 device=self.device,

gpudrive/integrations/vbd/data/amortize.py

@@ -27,6 +27,7 @@
 from gpudrive.datatypes.trajectory import LogTrajectory
 from gpudrive.datatypes.observation import GlobalEgoState
 from gpudrive.integrations.vbd.sim_agent.sim_actor import VBDTest
+import madrona_gpudrive
 
 
 def load_vbd_model(model_path, device="cpu", max_cont_agents=64):
@@ -77,7 +78,7 @@ def main():
 
     # Always use CPU device
     device = "cpu"
-    MAX_CONTROLLED_AGENTS = 64
+    MAX_CONTROLLED_AGENTS = 32
     print(f"Loading VBD model on {device}...")
     vbd_model = load_vbd_model(args.model_path, device, MAX_CONTROLLED_AGENTS)
 
@@ -106,7 +107,7 @@ def main():
             dataset_size=args.num_scenes,
             file_prefix="",
         ),
-        render_config=RenderConfig(render_3d=True),
+        render_config=RenderConfig(),
         max_cont_agents=MAX_CONTROLLED_AGENTS,  # Maximum number of agents to control per scene
         device=device,
     )
@@ -115,8 +116,8 @@ def main():
     output_trajectories = torch.zeros(
         gpudrive_env.num_worlds,
         gpudrive_env.max_agent_count,
-        env_config.episode_len,
-        5,
+        madrona_gpudrive.kTrajectoryLength,
+        6,
     )
 
     # Save init steps from logs
@@ -127,14 +128,17 @@ def main():
         device=device,
     )
 
-    output_trajectories[:, :, :INIT_STEPS, :2] = log_trajectory.pos_xy[
-        :, :, :INIT_STEPS
+    output_trajectories[:, :, : INIT_STEPS + 1, :2] = log_trajectory.pos_xy[
+        :, :, : INIT_STEPS + 1
     ]
-    output_trajectories[:, :, :INIT_STEPS, 2] = log_trajectory.yaw[
-        :, :, :INIT_STEPS, 0
+    output_trajectories[:, :, : INIT_STEPS + 1, 2] = log_trajectory.yaw[
+        :, :, : INIT_STEPS + 1, 0
     ]
-    output_trajectories[:, :, :INIT_STEPS, 3:] = log_trajectory.vel_xy[
-        :, :, :INIT_STEPS
+    output_trajectories[:, :, : INIT_STEPS + 1, 3:5] = log_trajectory.vel_xy[
+        :, :, : INIT_STEPS + 1
+    ]
+    output_trajectories[:, :, : INIT_STEPS + 1, 5] = log_trajectory.valids[
+        :, :, : INIT_STEPS + 1, 0
     ]
 
     # Action tensor to step through simulation
@@ -181,20 +185,23 @@ def main():
                 i, valid_world_indices, :, 2
             ]
             predicted_actions[i, valid_world_indices, :, 4:6] = vbd_output[
-                i, valid_world_indices, :, 3:
+                i, valid_world_indices, :, 3:5
             ]
 
             # Populate output trajectories
             output_trajectories[
-                i, valid_world_indices, INIT_STEPS:, :2
+                i, valid_world_indices, INIT_STEPS + 1 :, :2
             ] = vbd_output[i, valid_world_indices, :, :2] - world_means[
                 i
             ].view(
                 1, 1, 2
             )
             output_trajectories[
-                i, valid_world_indices, INIT_STEPS:, 2:
+                i, valid_world_indices, INIT_STEPS + 1 :, 2:5
             ] = vbd_output[i, valid_world_indices, :, 2:]
+            output_trajectories[
+                i, valid_world_indices, INIT_STEPS + 1 :, 5
+            ] = 1.0
 
         # Save to each file's json
         index_to_id = GlobalEgoState.from_tensor(

src/bindings.cpp

@@ -23,6 +23,7 @@ namespace madrona_gpudrive
         m.attr("kMaxAgentCount") = consts::kMaxAgentCount;
         m.attr("kMaxRoadEntityCount") = consts::kMaxRoadEntityCount;
         m.attr("kMaxAgentMapObservationsCount") = consts::kMaxAgentMapObservationsCount;
+        m.attr("kTrajectoryLength") = consts::kTrajectoryLength;
         m.attr("episodeLen") = consts::episodeLen;
         m.attr("numLidarSamples") = consts::numLidarSamples;
         m.attr("vehicleScale") = consts::vehicleLengthScale;

src/consts.hpp

@@ -8,7 +8,7 @@ namespace madrona_gpudrive {
 
 namespace consts {
 
-inline constexpr madrona::CountT kMaxAgentCount = 64;
+inline constexpr madrona::CountT kMaxAgentCount = 32;
 inline constexpr madrona::CountT kMaxRoadEntityCount = 10000;
 inline constexpr madrona::CountT kMaxAgentMapObservationsCount = 128;
 

src/json_serialization.hpp

@@ -145,8 +145,8 @@ namespace madrona_gpudrive
         obj.metadata.avgZ = avgZ;
 
         // Initialize VBD trajectories to zeros
-        for (int i = 0; i < consts::episodeLen; i++) {
-            for (int j = 0; j < 5; j++) {
+        for (int i = 0; i < consts::kTrajectoryLength; i++) {
+            for (int j = 0; j < 6; j++) {
                 obj.vbd_trajectories[i][j] = 0.0f;
             }
         }
@@ -155,13 +155,14 @@ namespace madrona_gpudrive
         if (j.contains("vbd_trajectory")) {
             int vbd_idx = 0;
             for (const auto &vbd_traj : j.at("vbd_trajectory")) {
-                if (vbd_idx < consts::episodeLen) {
+                if (vbd_idx < consts::kTrajectoryLength) {
                     if (!vbd_traj.is_null()) {
                         obj.vbd_trajectories[vbd_idx][0] = vbd_traj.at(0).get<float>();
                         obj.vbd_trajectories[vbd_idx][1] = vbd_traj.at(1).get<float>();
                         obj.vbd_trajectories[vbd_idx][2] = vbd_traj.at(2).get<float>();
                         obj.vbd_trajectories[vbd_idx][3] = vbd_traj.at(3).get<float>();
                         obj.vbd_trajectories[vbd_idx][4] = vbd_traj.at(4).get<float>();
+                        obj.vbd_trajectories[vbd_idx][5] = vbd_traj.at(5).get<float>();
                     }
                 }
                 vbd_idx++;

src/types.hpp

@@ -364,23 +364,23 @@ namespace madrona_gpudrive
 
     struct VBDTrajectory
     {
-        // For each agent, store the full VBD trajectory (x, y, yaw, vx, vy)
-        // The tensor has shape [traj_length, 5]
-        float trajectories[consts::episodeLen][5];
+        // For each agent, store the full VBD trajectory (x, y, yaw, vx, vy, valid)
+        // The tensor has shape [traj_length, 6]
+        float trajectories[consts::episodeLen][6];
 
         static inline void zero(VBDTrajectory& vbd_traj)
         {
             for (int i = 0; i < consts::episodeLen; i++)
             {
-                for (int j = 0; j < 5; j++)
+                for (int j = 0; j < 6; j++)
                 {
                     vbd_traj.trajectories[i][j] = 0.0f;
                 }
             }
         }
     };
 
-    const size_t VBDTrajectoryExportSize = consts::episodeLen * 5;
+    const size_t VBDTrajectoryExportSize = consts::episodeLen * 6;
 
     static_assert(sizeof(VBDTrajectory) == sizeof(float) * VBDTrajectoryExportSize);