From 651497add593a90ac40fddea7a54f960c382dda2 Mon Sep 17 00:00:00 2001
From: kevin <kevinwinston184@gmail.com>
Date: Tue, 13 May 2025 13:10:08 -0400
Subject: [PATCH] make vbd 91 steps

---
 gpudrive/datatypes/trajectory.py | 21 ++++++++-------------
 src/init.hpp                     |  2 +-
 src/level_gen.cpp                |  4 ++--
 src/mgr.cpp                      |  2 +-
 src/types.hpp                    |  6 +++---
 5 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/gpudrive/datatypes/trajectory.py b/gpudrive/datatypes/trajectory.py
index 15bd58a89..30881eed5 100644
--- a/gpudrive/datatypes/trajectory.py
+++ b/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
         )
         
diff --git a/src/init.hpp b/src/init.hpp
index 8d4ff60dc..15b19dc11 100755
--- a/src/init.hpp
+++ b/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
diff --git a/src/level_gen.cpp b/src/level_gen.cpp
index c43ad403f..06da6cb52 100755
--- a/src/level_gen.cpp
+++ b/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];
         }
     }
diff --git a/src/mgr.cpp b/src/mgr.cpp
index ae6cbf2c9..55227b95a 100755
--- a/src/mgr.cpp
+++ b/src/mgr.cpp
@@ -907,7 +907,7 @@ Tensor Manager::vbdTrajectoryTensor() const {
     {
         impl_->numWorlds,
         consts::kMaxAgentCount,
-        consts::episodeLen,
+        consts::kTrajectoryLength,
         6,
     });
 }
diff --git a/src/types.hpp b/src/types.hpp
index b0fe159ab..c4823d66d 100755
--- a/src/types.hpp
+++ b/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);