Convert GPUDrive yaw to be between [-pi, pi]

data_utils/vbd_data.py

@@ -1,10 +1,29 @@
 """Functions used for processing roadgraph data and other features for VBD."""
 import torch
-
+import numpy as np
 import gpudrive
 from vbd.data.data_utils import calculate_relations
 
 
+def convert_yaw(yaw_tensor):
+    """
+    Convert yaw values from [0, 2π] (GPUDrive) range to [-π, π] range (Waymax).
+
+    Parameters:
+        yaw_tensor (torch.Tensor): A tensor containing yaw values in radians.
+
+    Returns:
+        torch.Tensor: A tensor with yaw values converted to [-π, π] range.
+    """
+    # Ensure the tensor is in the correct range
+    yaw_tensor = yaw_tensor % (2 * torch.pi)
+
+    # Convert to the range [-π, π]
+    yaw_tensor[yaw_tensor > torch.pi] -= 2 * torch.pi
+
+    return yaw_tensor
+
+
 def process_scenario_data(
     num_envs,
     max_controlled_agents,
@@ -39,8 +58,11 @@ def process_scenario_data(
     agents_future = torch.zeros(
         (num_envs, max_controlled_agents, episode_len - init_steps, 5)
     )
+
+    converted_yaw = convert_yaw(yaws)
+
     agents_future[:, :, :, 0:2] = positions[:, :, init_steps:, :]
-    agents_future[:, :, :, 2] = yaws[:, :, init_steps:, :].squeeze(-1)
+    agents_future[:, :, :, 2] = converted_yaw[:, :, init_steps:, :].squeeze(-1)
     agents_future[:, :, :, 3:5] = velocities[:, :, init_steps:, :]
 
     # Set all invalid agent values to zero
@@ -117,11 +139,13 @@ def construct_agent_history(
     global_traj[~controlled_agent_mask] = 0.0
     global_traj = global_traj[:, :max_cont_agents, :]
 
+    converted_yaw = convert_yaw(yaw)
+
     # x, y, heading, vel_x, vel_y, len, width, height
     agents_history = torch.cat(
         [
             pos_xy[:, :, : init_steps + 1, :],  # x, y
-            yaw[:, :, : init_steps + 1, :],  # heading
+            converted_yaw[:, :, : init_steps + 1, :],  # heading
             vel_xy[:, :, : init_steps + 1, :],  # vel_x, vel_y
             global_traj[:, :, 10]
             .unsqueeze(-1)
@@ -251,8 +275,6 @@ def construct_polylines(
     return polylines, polylines_valid
 
 
-
 def sample_to_action():
-    """Todo: Implement this function.
-    """
-    pass
+    """Todo: Implement this function."""
+    pass