make vbd 91 steps

gpudrive/datatypes/trajectory.py

@@ -216,20 +216,15 @@ class VBDTrajectory:
 
     def __init__(self, vbd_traj_tensor: torch.Tensor):
         """Initializes the VBD trajectory with a tensor."""
-        # Pad to make sure the trajectory is of length 91
-        padding = (0, 0, 1, 0)
-
-        vbd_traj_tensor_pad = torch.nn.functional.pad(vbd_traj_tensor, padding, mode='constant', value=-1.0)
-
-        self.pos_x = vbd_traj_tensor_pad[:, :, :, 0].unsqueeze(-1)
-        self.pos_y = vbd_traj_tensor_pad[:, :, :, 1].unsqueeze(-1)
-        self.pos_xy = vbd_traj_tensor_pad[:, :, :, :2]
-        self.yaw = wrap_yaws(vbd_traj_tensor_pad[:, :, :, 2].unsqueeze(-1))
-        self.vel_x = vbd_traj_tensor_pad[:, :, :, 3].unsqueeze(-1)
-        self.vel_y = vbd_traj_tensor_pad[:, :, :, 4].unsqueeze(-1)
-        self.vel_xy = vbd_traj_tensor_pad[:, :, :, 3:5]
+        self.pos_x = vbd_traj_tensor[:, :, :, 0].unsqueeze(-1)
+        self.pos_y = vbd_traj_tensor[:, :, :, 1].unsqueeze(-1)
+        self.pos_xy = vbd_traj_tensor[:, :, :, :2]
+        self.yaw = wrap_yaws(vbd_traj_tensor[:, :, :, 2].unsqueeze(-1))
+        self.vel_x = vbd_traj_tensor[:, :, :, 3].unsqueeze(-1)
+        self.vel_y = vbd_traj_tensor[:, :, :, 4].unsqueeze(-1)
+        self.vel_xy = vbd_traj_tensor[:, :, :, 3:5]
         self.ref_speed = self.comp_reference_speed()
-        self.valids = vbd_traj_tensor_pad[:, :, :, 5].unsqueeze(-1).to(
+        self.valids = vbd_traj_tensor[:, :, :, 5].unsqueeze(-1).to(
             torch.int32
         )
 

src/init.hpp

@@ -37,7 +37,7 @@ namespace madrona_gpudrive
         uint32_t id;
         MapVector2 mean;
         bool markAsExpert{false};
-        float vbd_trajectories[consts::episodeLen][5];  // x, y, yaw, vx, vy
+        float vbd_trajectories[consts::kTrajectoryLength][6];  // x, y, yaw, vx, vy, valid
     };
 
     struct MapRoad

src/level_gen.cpp

@@ -109,8 +109,8 @@ static inline void populateVBDTrajectory(Engine &ctx, const Entity &agent, const
     auto &vbd_trajectory = ctx.get<VBDTrajectory>(agent_iface);
 
     // Copy VBD trajectories from the map object
-    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++) {
             vbd_trajectory.trajectories[i][j] = agentInit.vbd_trajectories[i][j];
         }
     }

src/mgr.cpp

@@ -907,7 +907,7 @@ Tensor Manager::vbdTrajectoryTensor() const {
     {
         impl_->numWorlds,
         consts::kMaxAgentCount,
-        consts::episodeLen,
+        consts::kTrajectoryLength,
         6,
     });
 }

src/types.hpp

@@ -368,11 +368,11 @@ namespace madrona_gpudrive
     {
         // 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];
+        float trajectories[consts::kTrajectoryLength][6];
 
         static inline void zero(VBDTrajectory& vbd_traj)
         {
-            for (int i = 0; i < consts::episodeLen; i++)
+            for (int i = 0; i < consts::kTrajectoryLength; i++)
             {
                 for (int j = 0; j < 6; j++)
                 {
@@ -382,7 +382,7 @@ namespace madrona_gpudrive
         }
     };
 
-    const size_t VBDTrajectoryExportSize = consts::episodeLen * 6;
+    const size_t VBDTrajectoryExportSize = consts::kTrajectoryLength * 6;
 
     static_assert(sizeof(VBDTrajectory) == sizeof(float) * VBDTrajectoryExportSize);