address comments

Author: xinshuoweng

src/trajdata/dataset_specific/mads/__init__.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
 
 """MADS dataset package exports."""
 

src/trajdata/dataset_specific/mads/constant.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
 
 """Constants and enums for the MADS dataset loader."""
 
@@ -25,7 +26,7 @@
 EGO_WIDTH: Final[float] = 2.11311007
 EGO_HEIGHT: Final[float] = 1.34435794
 
-# by default, we support internal production v2 dataset, with clipgt-2.0.0 and above.
+# by default, we support v2 data, version clipgt-2.0.0, clipgt-2.3.0 and above.
 DATA_SRC: Final[str] = "v2"
 
 # Allowed MADS data source tags.
@@ -53,7 +54,7 @@ def resolve_data_src(cli_override: Optional[str] = None) -> str:
 
 
 class ObstacleClassV1(Enum):
-    """Obstacle classes for MADS based on NDAS `obstacle_types.proto`."""
+    """Obstacle classes for MADS."""
 
     # buf:lint:ignore ENUM_ZERO_VALUE_SUFFIX
     OBSTACLE_CLASS_INVALID = 0

src/trajdata/dataset_specific/mads/mads_dataset.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
 
 import argparse
 import glob
@@ -124,7 +125,7 @@ def compute_metadata(self, env_name: str, data_dir: str) -> EnvMetadata:
         # get all clip ids
         self.clip_duration = clip_duration
         clip_items = list(clip_dir.items())
-        random.shuffle(clip_items)
+        random.Random(42).shuffle(clip_items)
         self.clip_dir = dict(clip_items)
         clip_ids = list(clip_dir.keys())
         all_clips = [clip_id for clip_id, _ in clip_items[:]]
@@ -278,7 +279,7 @@ def get_scene(self, scene_info: SceneMetadata) -> Scene:
         )
 
     @staticmethod
-    def get_df_from_path(
+    def get_ego_df_from_path(
         scene_path: str,
         scene_name: str,
         verbose: bool = False,
@@ -598,7 +599,7 @@ def get_row_first_seen(df: pd.DataFrame) -> pd.DataFrame:
     def get_agent_info(
         self, scene: Scene, cache_path: Path, cache_class: Type[SceneCache]
     ) -> Tuple[List[AgentMetadata], List[List[AgentMetadata]]]:
-        sorted_df = self.get_df_from_path(self.clip_dir[scene.name], scene.name)
+        sorted_df = self.get_ego_df_from_path(self.clip_dir[scene.name], scene.name)
 
         contain_obstacles: bool = False
         agent_list: List[AgentMetadata] = []
@@ -761,12 +762,10 @@ def cache_maps(
 
 def _debug_dump_scene_df(data_src: Optional[str] = None) -> None:
     """Debug helper to inspect one scene dataframe when run as a script."""
-    scene_path = (
-        "/lustre/fsw/portfolios/nvr/users/xweng/agentdriver_alpamayo/data/new_data"
-    )
+    scene_path = 'path/to/source/data'
     scene_name = "762e063d-6eb9-43ae-959c-e53af10b53f9"
     scene_path = os.path.join(scene_path, scene_name)
-    ego_df: pd.DataFrame = MADSDataset.get_df_from_path(
+    ego_df: pd.DataFrame = MADSDataset.get_ego_df_from_path(
         scene_path, scene_name, verbose=True, data_src=data_src
     )
 

src/trajdata/dataset_specific/mads/mads_utils.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
 
 import multiprocessing
 import os

src/trajdata/dataset_specific/mads/tar_extractor.py

@@ -1,4 +1,19 @@
-# Copyright (c) 2025 NVIDIA CORPORATION.  All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
 """Tar extraction utilities for component data files.
 
 Provides TarExtractor class for extracting tar archives to RAM filesystem

src/trajdata/dataset_specific/pai/pai_dataset.py

@@ -51,7 +51,7 @@ def compute_metadata(self, env_name: str, data_dir: str) -> EnvMetadata:
         egomotion_dir = Path(self.data_dir) / "labels" / "egomotion"
         if not egomotion_dir.exists():
             raise FileNotFoundError(
-                f"Expected PAI egomotion directory at {egomotion_dir}."
+                f"Expected PhysicalAI-AV egomotion directory at {egomotion_dir}."
             )
 
         clip_dir: Dict[str, str] = {}
@@ -85,7 +85,7 @@ def compute_metadata(self, env_name: str, data_dir: str) -> EnvMetadata:
         self.clip_duration = clip_duration
 
         clip_items = list(clip_dir.items())
-        random.shuffle(clip_items)
+        random.Random(42).shuffle(clip_items)
         self.clip_dir = dict(clip_items)
 
         all_clips = [clip_id for clip_id, _ in clip_items]
@@ -199,7 +199,7 @@ def _load_ego_df(scene_path: str) -> pd.DataFrame:
         return pd.read_parquet(scene_path)
 
     @staticmethod
-    def get_df_from_path(
+    def get_ego_df_from_path(
         scene_path: str,
         scene_name: str,
         verbose: bool = False,
@@ -243,24 +243,10 @@ def get_df_from_path(
         elif all(key in ego_df.columns for key in prefixed_map.keys()):
             normalized = ego_df.rename(columns=prefixed_map).copy()
         else:
+            expected_cols = [*col_map.keys(), *prefixed_map.keys()]
             missing = [
                 key
-                for key in [
-                    "x",
-                    "y",
-                    "z",
-                    "qx",
-                    "qy",
-                    "qz",
-                    "qw",
-                    "EgomotionEstimate.location.x",
-                    "EgomotionEstimate.location.y",
-                    "EgomotionEstimate.location.z",
-                    "EgomotionEstimate.orientation.x",
-                    "EgomotionEstimate.orientation.y",
-                    "EgomotionEstimate.orientation.z",
-                    "EgomotionEstimate.orientation.w",
-                ]
+                for key in expected_cols
                 if key not in ego_df.columns
             ]
             raise KeyError(
@@ -359,7 +345,7 @@ def _interp(col_name: str) -> np.ndarray:
     def get_agent_info(
         self, scene: Scene, cache_path: Path, cache_class: Type[SceneCache]
     ) -> Tuple[List[AgentMetadata], List[List[AgentMetadata]]]:
-        ego_df = self.get_df_from_path(self.clip_dir[scene.name], scene.name)
+        ego_df = self.get_ego_df_from_path(self.clip_dir[scene.name], scene.name)
         ego_df.set_index(["agent_id", "scene_ts"], inplace=True)
 
         ego_metadata = AgentMetadata(

src/trajdata/utils/env_utils.py

@@ -58,10 +58,11 @@ def get_raw_dataset(dataset_name: str, data_dir: str) -> RawDataset:
     if "mads" in dataset_name.lower():
         from trajdata.dataset_specific.mads import MADSDataset
 
-        return MADSDataset(dataset_name, data_dir, parallelizable=True, has_maps=False)
+        return MADSDataset(dataset_name, data_dir, parallelizable=True, has_maps=True)
 
     if "pai" in dataset_name.lower():
         from trajdata.dataset_specific.pai import PAIDataset
+
         return PAIDataset(dataset_name, data_dir, parallelizable=True, has_maps=False)
 
     raise ValueError(f"Dataset with name '{dataset_name}' is not supported")