From 8b8051216254093255a8f4809766ac37f2e8a4f2 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:27:43 +0000 Subject: [PATCH 01/17] feat: add SceneDataSource protocol and DataSourceConfig Introduce SceneDataSource protocol to abstract scene data loading, enabling support for multiple data sources (Artifact, trajdata, etc). Add DataSourceConfig to runtime config for unified data source configuration via trajdata's UnifiedDataset. Changes: - Add SceneDataSource protocol with standard interface (rig, traffic_objects, map, metadata) - Add DataSourceConfig with trajdata UnifiedDataset parameters - Support for both USDZ files and standard trajdata datasets --- src/runtime/alpasim_runtime/config.py | 40 ++++++++++ src/utils/alpasim_utils/scene_data_source.py | 78 ++++++++++++++++++++ 2 files changed, 118 insertions(+) create mode 100644 src/utils/alpasim_utils/scene_data_source.py diff --git a/src/runtime/alpasim_runtime/config.py b/src/runtime/alpasim_runtime/config.py index e6f883c9..47896f97 100644 --- a/src/runtime/alpasim_runtime/config.py +++ b/src/runtime/alpasim_runtime/config.py @@ -17,6 +17,42 @@ C = TypeVar("C") +@dataclass +class DataSourceConfig: + """Configuration for trajdata's UnifiedDataset. + + This config unifies all scene data loading through trajdata's UnifiedDataset, + supporting both USDZ files and standard trajdata datasets like NuPlan, Waymo, etc. + + Attributes: + desired_data: List of dataset names to load (e.g., ["nuplan_test", "usdz"]) + data_dirs: Dict mapping dataset names to their data directories + cache_location: Path to trajdata cache directory + config_dir: Optional directory containing YAML scene config files for batch mode + asset_base_path: Base path for rendering assets (e.g., MTGS assets) + incl_vector_map: Whether to load vector maps + rebuild_cache: Whether to force rebuild the trajdata cache + rebuild_maps: Whether to force rebuild maps + desired_dt: Desired time delta between frames in seconds + num_workers: Number of workers for data loading + num_timesteps_before: Number of timesteps before central token (batch mode) + num_timesteps_after: Number of timesteps after central token (batch mode) + """ + + desired_data: list[str] = MISSING + data_dirs: dict[str, str] = MISSING + cache_location: str = MISSING + config_dir: Optional[str] = None # For YAML batch mode + asset_base_path: Optional[str] = None + incl_vector_map: bool = True + rebuild_cache: bool = False + rebuild_maps: bool = False + desired_dt: float = 0.1 # 10 Hz default + num_workers: int = 1 + num_timesteps_before: int = 30 # For batch mode + num_timesteps_after: int = 80 # For batch mode + + def typed_parse_config(path: str | Path, config_type: Type[C]) -> C: """Reads a yaml file at `path` and parses it into a provided type using omegaconf.""" yaml_config = OmegaConf.create(load_yaml_dict(path)) @@ -253,6 +289,10 @@ class UserSimulatorConfig: # >1 = multi-worker mode with subprocess-based parallelism nr_workers: int = MISSING + # Unified data source configuration (optional for backward compatibility) + # When provided, data loading goes through trajdata's UnifiedDataset + data_source: Optional[DataSourceConfig] = None + @dataclass class SimulatorConfig: diff --git a/src/utils/alpasim_utils/scene_data_source.py b/src/utils/alpasim_utils/scene_data_source.py new file mode 100644 index 00000000..0e4ca073 --- /dev/null +++ b/src/utils/alpasim_utils/scene_data_source.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025-2026 NVIDIA Corporation + +""" +SceneDataSource Protocol for abstracting scene data loading. + +This Protocol allows Runtime to work with different data sources (USDZ, Nuplan, Waymo, etc.) +without being tied to a specific implementation. Any class that implements this Protocol +can be used as a data source for alpasim Runtime. +""" + +from __future__ import annotations + +from typing import Optional, Protocol, runtime_checkable + +try: + from trajdata.maps import VectorMap +except ImportError: + VectorMap = None # type: ignore + +from alpasim_utils.artifact import Metadata +from alpasim_utils.scenario import Rig, TrafficObjects + + +@runtime_checkable +class SceneDataSource(Protocol): + """ + Protocol defining the interface for scene data sources. + + Any class implementing this protocol can be used as a data source for alpasim Runtime. + This allows supporting multiple data formats (USDZ, Nuplan, Waymo, etc.) without + modifying Runtime code. + + Attributes: + scene_id: Unique identifier for the scene + """ + + scene_id: str + + @property + def rig(self) -> Rig: + """ + Get the rig (ego vehicle) trajectory and configuration. + + Returns: + Rig object containing trajectory, camera IDs, and vehicle config + """ + ... + + @property + def traffic_objects(self) -> TrafficObjects: + """ + Get traffic objects (vehicles, pedestrians, etc.) in the scene. + + Returns: + TrafficObjects dictionary mapping track_id to TrafficObject + """ + ... + + @property + def map(self) -> Optional[VectorMap]: + """ + Get the vector map for the scene. + + Returns: + VectorMap object or None if map data is not available + """ + ... + + @property + def metadata(self) -> Metadata: + """ + Get metadata about the scene. + + Returns: + Metadata object containing scene information + """ + ... From d2bf456aeee5ed5e7814ff34de894f7bb8af3e06 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:28:02 +0000 Subject: [PATCH 02/17] feat: implement TrajdataDataSource for direct trajdata loading Implement TrajdataDataSource as a SceneDataSource that loads data directly from trajdata datasets without requiring USDZ conversion. Features: - Lazy loading of rig, traffic_objects, map, and metadata - Pre-created scene_cache to avoid pickle errors in multiprocessing - Support for coordinate frame transformations (world to local NRE) - Trajectory smoothing with cubic splines - Camera calibration extraction from scene metadata - Map loading and transformation to local coordinates Benefits: - Eliminate USDZ conversion overhead - Reduce startup time with on-demand loading - Lower memory usage per worker --- .../alpasim_utils/trajdata_data_source.py | 1051 +++++++++++++++++ 1 file changed, 1051 insertions(+) create mode 100644 src/utils/alpasim_utils/trajdata_data_source.py diff --git a/src/utils/alpasim_utils/trajdata_data_source.py b/src/utils/alpasim_utils/trajdata_data_source.py new file mode 100644 index 00000000..4cd4146f --- /dev/null +++ b/src/utils/alpasim_utils/trajdata_data_source.py @@ -0,0 +1,1051 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025-2026 NVIDIA Corporation + +""" +TrajdataDataSource: Implementation for loading scene data directly from trajdata + +This class demonstrates how to create a SceneDataSource implementation that loads +data directly from trajdata converted data without requiring USDZ format. This is +useful for researchers using trajdata datasets. + +Usage example: + from trajdata import UnifiedDataset + from alpasim_utils.trajdata_data_source import TrajdataDataSource + + # Load trajdata dataset + dataset = UnifiedDataset( + desired_data=["nusc_mini"], + data_dirs={"/path/to/trajdata/data"}, + ... + ) + + # Get a scene + scene = dataset.get_scene("nusc_mini", "scene-0001") + + # Create data source + data_source = TrajdataDataSource.from_trajdata_scene(scene) + + # Now can be used in Runtime + # artifacts = {data_source.scene_id: data_source} +""" + +from __future__ import annotations + +import copy +import logging +import os +from dataclasses import dataclass +from typing import Optional + +import numpy as np +from scipy.spatial.transform import Rotation as R + +try: + from trajdata import AgentBatch + from trajdata.caching import EnvCache + from trajdata.data_structures.agent import AgentMetadata + from trajdata.data_structures.scene_metadata import Scene + from trajdata.dataset import UnifiedDataset + from trajdata.maps import VectorMap +except ImportError: + AgentBatch = None + Scene = None + UnifiedDataset = None + EnvCache = None + AgentMetadata = None + VectorMap = None + +from alpasim_utils.artifact import Metadata +from alpasim_utils.geometry import Trajectory +from alpasim_utils.scenario import ( + AABB, + CameraId, + Rig, + TrafficObject, + TrafficObjects, + VehicleConfig, +) +from alpasim_utils.scene_data_source import SceneDataSource + +logger = logging.getLogger(__name__) + + +@dataclass +class TrajdataDataSource(SceneDataSource): + """ + Implementation for loading scene data directly from trajdata. + + This class implements the SceneDataSource protocol, allowing direct loading + from trajdata Scene or AgentBatch objects without requiring USDZ format. + """ + + _scene: Scene | None = None + _scene_cache: EnvCache | None = None + _dataset: UnifiedDataset | None = None + _rig: Rig | None = None + _traffic_objects: TrafficObjects | None = None + _map: VectorMap | None = None + _metadata: Metadata | None = None + _smooth_trajectories: bool = True + _scene_id: str = "" + _asset_base_path: str | None = None # Base path for rendering assets + + @classmethod + def from_trajdata_scene( + cls, + scene: Scene, + dataset: Optional[UnifiedDataset] = None, + scene_cache: Optional[EnvCache] = None, + scene_id: Optional[str] = None, + smooth_trajectories: bool = True, + base_timestamp_us: int = 0, + asset_base_path: Optional[str] = None, + ) -> TrajdataDataSource: + """ + Create TrajdataDataSource from trajdata Scene object. + + Args: + scene: trajdata Scene object + dataset: UnifiedDataset instance (for getting scene_cache and map) + scene_cache: Optional EnvCache instance (if not provided, will be created from dataset) + scene_id: Optional scene ID (if not provided, uses scene.name) + smooth_trajectories: Whether to smooth trajectories + base_timestamp_us: Base timestamp in microseconds, starts from 0 if None + asset_base_path: Base path for rendering assets (e.g., MTGS assets) + + Returns: + TrajdataDataSource instance + """ + if Scene is None: + raise ImportError( + "trajdata is not installed. Please install it to use TrajdataDataSource." + ) + + data_source = cls( + _scene=scene, + _dataset=dataset, + _scene_cache=scene_cache, + _scene_id=scene_id or scene.name, + _smooth_trajectories=smooth_trajectories, + _asset_base_path=asset_base_path, + ) + data_source._base_timestamp_us = base_timestamp_us + return data_source + + @classmethod + def from_agent_batch( + cls, + batch: AgentBatch, + scene_id: str, + smooth_trajectories: bool = True, + ) -> TrajdataDataSource: + """ + Create TrajdataDataSource from trajdata AgentBatch object. + + Note: This method requires the batch to contain complete scene information. + It is generally recommended to use from_trajdata_scene instead. + + Args: + batch: trajdata AgentBatch object + scene_id: Scene ID + smooth_trajectories: Whether to smooth trajectories + + Returns: + TrajdataDataSource instance + """ + if AgentBatch is None: + raise ImportError( + "trajdata is not installed. Please install it to use TrajdataDataSource." + ) + + data_source = cls(scene_id=scene_id) + data_source._smooth_trajectories = smooth_trajectories + # Extract data from batch + data_source._load_from_batch(batch) + return data_source + + def _load_from_batch(self, batch: AgentBatch) -> None: + """Load data from AgentBatch (internal method)""" + # Need to extract data based on batch structure + # Specific implementation depends on your trajdata data format + raise NotImplementedError( + "from_agent_batch needs to be implemented based on your trajdata data format. " + "It is recommended to use from_trajdata_scene method instead." + ) + + @property + def scene_id(self) -> str: + """Scene ID""" + if self._scene_id: + return self._scene_id + if self._scene is not None: + return self._scene.name + raise ValueError("scene_id is not set and cannot be obtained from scene") + + @scene_id.setter + def scene_id(self, value: str) -> None: + self._scene_id = value + + @property + def asset_path(self) -> str | None: + """ + Resolve asset folder path for this scene. + + The asset path is constructed by appending the scene name to _asset_base_path. + The _asset_base_path should already contain any dataset-specific subdirectories + (e.g., it might be /data/WE_processed/navtest/assets for MTGS). + + Returns: + Resolved asset folder path, or None if _asset_base_path is not set + """ + if self._asset_base_path is None: + return None + + # Extract asset folder name from scene metadata + scene_name = self._extract_asset_folder_name() + + # Simple join: asset_base_path already contains dataset-specific subdirs + return os.path.join(self._asset_base_path, scene_name) + + def _extract_asset_folder_name(self) -> str: + """ + Extract the asset folder name from scene metadata. + + This method attempts to determine the appropriate asset folder name + based on scene metadata. Override this in subclasses if needed. + + Resolution order: + 1. USDZ: Use usdz_stem from data_access_info + 2. Other datasets: Use log_id or asset_folder from data_access_info + 3. Fallback: Use scene_id with common suffixes removed + + Returns: + Asset folder name (defaults to scene_id if no specific name found) + """ + # Try to get from scene data_access_info + if self._scene is not None and hasattr(self._scene, "data_access_info"): + data_access_info = self._scene.data_access_info or {} + + # USDZ: Use usdz_stem (filename without .usdz extension) + if "usdz_stem" in data_access_info: + return data_access_info["usdz_stem"] + + # Look for asset_folder or similar keys + if "asset_folder" in data_access_info: + return data_access_info["asset_folder"] + + # NuPlan and other datasets: use log_id + if "log_id" in data_access_info: + return data_access_info["log_id"] + + # Default: use scene_id (potentially with suffix removed) + scene_name = self.scene_id + # Remove common suffixes like "-001" + if len(scene_name) > 4 and scene_name[-4] == "-" and scene_name[-3:].isdigit(): + scene_name = scene_name[:-4] + return scene_name + + def set_asset_base_path(self, path: str | None) -> None: + """Set the base path for rendering assets.""" + self._asset_base_path = path + + def _get_scene_cache(self) -> EnvCache: + """Get or create scene_cache""" + if self._scene_cache is not None: + return self._scene_cache + + if self._scene is None: + raise ValueError("Cannot create scene_cache: scene is not set") + + # Try to create from dataset if available + if self._dataset is not None: + logger.debug(f"Creating scene_cache for scene: {self._scene.name}") + try: + self._scene_cache = self._dataset.cache_class( + self._dataset.cache_path, self._scene, self._dataset.augmentations + ) + self._scene_cache.set_obs_format(self._dataset.obs_format) + logger.debug("Scene cache created successfully") + return self._scene_cache + except Exception as e: + logger.error(f"Failed to create scene_cache: {e}") + raise + + # If dataset is not set, scene_cache must be provided externally + raise ValueError( + "Cannot create scene_cache: dataset is not set and scene_cache was not provided. " + "Either pass 'dataset' parameter or pre-create 'scene_cache' when creating TrajdataDataSource. " + "Example: TrajdataDataSource.from_trajdata_scene(scene, dataset=your_dataset) " + "or TrajdataDataSource.from_trajdata_scene(scene, scene_cache=your_cache)" + ) + + def _extract_agent_trajectory( + self, + agent: AgentMetadata, + ) -> tuple[Optional[Trajectory], Optional[VehicleConfig]]: + """Extract complete trajectory for agent (refer to trajdata_artifact_converter.py implementation)""" + if self._scene is None: + return None, None + + scene_cache = self._get_scene_cache() + dt = self._scene.dt + base_timestamp_us = getattr(self, "_base_timestamp_us", None) + + try: + timestamps_us = [] + poses_vec3 = [] + poses_quat = [] + + # Iterate through all timesteps + for ts in range(agent.first_timestep, agent.last_timestep + 1): + try: + state = scene_cache.get_raw_state(agent.name, ts) + + # Get position and orientation + x = state.get_attr("x") if hasattr(state, "get_attr") else state.x + y = state.get_attr("y") if hasattr(state, "get_attr") else state.y + z = ( + state.get_attr("z") + if hasattr(state, "get_attr") + else (state.z if hasattr(state, "z") else 0.0) + ) + heading = ( + state.get_attr("h") if hasattr(state, "get_attr") else state.h + ) + + # Convert to numpy array (handle scalar case) + if isinstance(x, (int, float)): + x = np.array([x]) + if isinstance(y, (int, float)): + y = np.array([y]) + if isinstance(z, (int, float)): + z = np.array([z]) + if isinstance(heading, (int, float)): + heading = np.array([heading]) + + # Take first element (if array) + x_val = float(x[0] if x.ndim > 0 else x) + y_val = float(y[0] if y.ndim > 0 else y) + z_val = float(z[0] if z.ndim > 0 else z) + heading_val = float(heading[0] if heading.ndim > 0 else heading) + + # Calculate timestamp + if base_timestamp_us is None: + timestamp_us = int(ts * dt * 1e6) + else: + timestamp_us = int(base_timestamp_us + ts * dt * 1e6) + + timestamps_us.append(timestamp_us) + poses_vec3.append([x_val, y_val, z_val]) + + # Convert heading to quaternion + quat = R.from_euler("z", heading_val).as_quat() # [x, y, z, w] + poses_quat.append(quat) + + except Exception as e: + logger.debug( + f"Failed to get state for agent {agent.name} at ts {ts}: {e}" + ) + continue + + if len(timestamps_us) == 0: + return None, None + + # Create Trajectory + trajectory = Trajectory( + timestamps=np.array(timestamps_us, dtype=np.uint64), + positions=np.array(poses_vec3, dtype=np.float32), + quaternions=np.array(poses_quat, dtype=np.float32), + ) + + # Create VehicleConfig (extract from extent) + vehicle_config = None + if hasattr(agent.extent, "length"): + vehicle_config = VehicleConfig( + aabb_x_m=agent.extent.length, + aabb_y_m=agent.extent.width, + aabb_z_m=agent.extent.height, + aabb_x_offset_m=-agent.extent.length / 2, + aabb_y_offset_m=0.0, + aabb_z_offset_m=-agent.extent.height / 2, + ) + + return trajectory, vehicle_config + + except Exception as e: + logger.error(f"Failed to extract trajectory for agent {agent.name}: {e}") + return None, None + + @property + def rig(self) -> Rig: + """Load and return Rig object for ego vehicle""" + if self._rig is not None: + return self._rig + + if self._scene is None: + raise ValueError("Cannot load rig: scene is not set") + + # Get all agents + all_agents = self._scene.agents if self._scene.agents else [] + + # Identify ego agent + ego_agent = next((a for a in all_agents if a.name == "ego"), None) + if ego_agent is None and len(all_agents) > 0: + # If no ego, use first agent + ego_agent = all_agents[0] + logger.warning(f"No ego agent found, using first agent: {ego_agent.name}") + + if ego_agent is None: + raise ValueError("No ego agent found in scene") + + # Extract ego trajectory + ego_trajectory, ego_vehicle_config = self._extract_agent_trajectory(ego_agent) + + if ego_trajectory is None: + logger.error( + f"Failed to extract ego trajectory for agent {ego_agent.name}. " + f"Check if scene_cache is properly initialized and agent data is available." + ) + raise ValueError("Cannot extract ego trajectory") + + # Calculate world_to_nre transformation matrix (use first trajectory point as origin) + world_to_nre = np.eye(4) + if len(ego_trajectory) > 0: + first_pose_position = ego_trajectory.positions[0] + world_to_nre[:3, 3] = -first_pose_position + logger.info( + f"Setting world_to_nre origin at first pose: {first_pose_position}, " + f"translation: {world_to_nre[:3, 3]}" + ) + + # Convert ego trajectory to local coordinates (NRE) + if len(ego_trajectory) > 0: + translation = world_to_nre[:3, 3] + local_positions = ego_trajectory.positions + translation + + # Validate transform + first_pose_local = local_positions[0] + if np.linalg.norm(first_pose_local[:2]) > 1.0: + logger.warning( + f"First pose after transformation is not at origin: {first_pose_local}. " + f"Expected [0, 0, ~z], got {first_pose_local}" + ) + + local_quat = ego_trajectory.quaternions.copy() + ego_trajectory = Trajectory( + timestamps=ego_trajectory.timestamps_us.copy(), + positions=local_positions, + quaternions=local_quat, + ) + + logger.debug( + f"Transformed ego trajectory to local coordinates. " + f"First pose: {ego_trajectory.first_pose}, " + f"Range: X[{local_positions[:, 0].min():.2f}, {local_positions[:, 0].max():.2f}], " + f"Y[{local_positions[:, 1].min():.2f}, {local_positions[:, 1].max():.2f}], " + f"Z[{local_positions[:, 2].min():.2f}, {local_positions[:, 2].max():.2f}]" + ) + + # Extract camera information (refer to trajdata_artifact_converter.py) + camera_ids, _ = self._extract_camera_info_from_scene() + + self._rig = Rig( + sequence_id=self.scene_id, + trajectory=ego_trajectory, + camera_ids=camera_ids, + world_to_nre=world_to_nre, + vehicle_config=ego_vehicle_config, + ) + + return self._rig + + def _extract_camera_info_from_scene(self) -> tuple[list[CameraId], dict]: + """Extract camera information from scene (refer to trajdata_artifact_converter.py)""" + camera_ids = [] + camera_calibrations = {} + + if self._scene is None: + return camera_ids, camera_calibrations + + # Check if sensor_calibration information exists + if ( + not hasattr(self._scene, "data_access_info") + or not self._scene.data_access_info + ): + logger.warning( + "scene.data_access_info does not exist, skipping camera information extraction" + ) + return camera_ids, camera_calibrations + + sensor_calibration = self._scene.data_access_info.get("sensor_calibration") + if not sensor_calibration or not isinstance(sensor_calibration, dict): + logger.warning( + "sensor_calibration does not exist or has incorrect format, skipping camera information extraction" + ) + return camera_ids, camera_calibrations + + unique_sensor_idx = 0 + for camera_name, calibration_info in sensor_calibration.get( + "cameras", {} + ).items(): + try: + unique_camera_id = f"{camera_name}@{self.scene_id}" + + position = calibration_info.get( + "sensor2ego_translation", [0.0, 0.0, 0.0] + ) + rotation = calibration_info.get( + "sensor2ego_rotation", [0.0, 0.0, 0.0, 1.0] + ) + + if isinstance(position, (int, float)): + position = [float(position), 0.0, 0.0] + elif len(position) < 3: + position = list(position) + [0.0] * (3 - len(position)) + + if isinstance(rotation, (int, float)): + rotation = [0.0, 0.0, 0.0, 1.0] + elif len(rotation) < 4: + if len(rotation) == 3: + r = R.from_euler("xyz", rotation) + rotation = r.as_quat() + else: + rotation = [0.0, 0.0, 0.0, 1.0] + + camera_id = CameraId( + logical_name=camera_name, + trajectory_idx=0, + sequence_id=self.scene_id, + unique_id=unique_camera_id, + ) + camera_ids.append(camera_id) + unique_sensor_idx += 1 + + except Exception as e: + logger.warning( + f"Error extracting camera {camera_name} information: {e}" + ) + continue + + if len(camera_ids) == 0: + # If no camera information, create a default one + logger.warning( + f"Scene {self.scene_id} has no camera information, using default camera" + ) + camera_ids.append( + CameraId( + logical_name="camera_front", + trajectory_idx=0, + sequence_id=self.scene_id, + unique_id="0@camera_front", + ) + ) + + return camera_ids, camera_calibrations + + def _is_static_object( + self, trajectory: Trajectory, velocity_threshold: float = 0.1 + ) -> bool: + """Determine if object is static (based on velocity)""" + if len(trajectory) < 2: + return True + + positions = trajectory.positions + timestamps = trajectory.timestamps_us.astype(np.float64) / 1e6 + + velocities = [] + for i in range(1, len(positions)): + dt_sec = timestamps[i] - timestamps[i - 1] + if dt_sec > 0: + displacement = np.linalg.norm(positions[i] - positions[i - 1]) + velocity = displacement / dt_sec + velocities.append(velocity) + + if len(velocities) == 0: + return True + + avg_velocity = np.mean(velocities) + return avg_velocity < velocity_threshold + + @property + def traffic_objects(self) -> TrafficObjects: + """Load and return traffic objects""" + if self._traffic_objects is not None: + return self._traffic_objects + + if self._scene is None: + raise ValueError("Cannot load traffic_objects: scene is not set") + + # Get all agents + all_agents = self._scene.agents if self._scene.agents else [] + + # Identify ego agent + ego_agent = next((a for a in all_agents if a.name == "ego"), None) + if ego_agent is None and len(all_agents) > 0: + ego_agent = all_agents[0] + + traffic_dict = {} + for agent in all_agents: + # Skip ego agent + if agent.name == "ego" or agent == ego_agent: + continue + + # Extract trajectory + trajectory, _ = self._extract_agent_trajectory(agent) + + # Filter out empty trajectories or trajectories with only 1 data point + if trajectory is None or len(trajectory) < 2: + continue + + # Convert trajectory to local coordinates (NRE) - use rig's world_to_nre + if self._rig is None: + # If rig is not loaded yet, load it first + _ = self.rig + + world_to_nre = self._rig.world_to_nre + translation = world_to_nre[:3, 3] + local_positions = trajectory.positions + translation + local_quat = trajectory.quaternions.copy() + trajectory = Trajectory( + timestamps=trajectory.timestamps_us.copy(), + positions=local_positions, + quaternions=local_quat, + ) + + # Smooth if needed + if self._smooth_trajectories: + try: + import csaps + + css = csaps.CubicSmoothingSpline( + trajectory.timestamps_us / 1e6, + trajectory.positions.T, + normalizedsmooth=True, + ) + filtered_positions = css(trajectory.timestamps_us / 1e6).T + max_error = np.max( + np.abs(filtered_positions - trajectory.positions) + ) + if max_error > 1.0: + logger.warning( + f"Max error in cubic spline approximation: {max_error:.6f} m for {agent.name=}" + ) + # Create new trajectory with smoothed positions + trajectory = Trajectory( + timestamps=trajectory.timestamps_us.copy(), + positions=filtered_positions.astype(np.float32), + quaternions=trajectory.quaternions.copy(), + ) + except ImportError: + logger.warning("csaps not installed, skipping trajectory smoothing") + + # Get AABB + if hasattr(agent.extent, "length"): + aabb = AABB( + x=agent.extent.length, y=agent.extent.width, z=agent.extent.height + ) + else: + # Default AABB + aabb = AABB(x=4.5, y=1.8, z=1.5) + + # Determine if static object + is_static = self._is_static_object(trajectory) + + # Get category label + label_class = agent.type.name if hasattr(agent.type, "name") else "UNKNOWN" + + traffic_dict[agent.name] = TrafficObject( + track_id=agent.name, + aabb=aabb, + trajectory=trajectory, + is_static=is_static, + label_class=label_class, + ) + + self._traffic_objects = TrafficObjects(**traffic_dict) + return self._traffic_objects + + @property + def map(self) -> Optional[VectorMap]: + """Load and return VectorMap (obtained from dataset._map_api or scene.map_data)""" + if self._map is not None: + return self._map + + if VectorMap is None: + logger.warning("trajdata is not installed, cannot load map") + return None + + if self._scene is None: + logger.warning("Cannot load map: scene is not set") + return None + + # First, try to get map from scene.map_data (for USDZ and other datasets that attach map directly) + if hasattr(self._scene, "map_data") and self._scene.map_data is not None: + logger.info(f"Loading map from scene.map_data for {self.scene_id}") + # Make a deep copy to avoid modifying shared map object + self._map = copy.deepcopy(self._scene.map_data) + + # Apply coordinate transformation if needed + if self._rig is None: + # If rig is not loaded yet, load it first (this will set world_to_nre) + _ = self.rig + + world_to_nre = self._rig.world_to_nre + + # Check if transformation is needed (if world_to_nre is not identity) + if world_to_nre is not None and not np.allclose(world_to_nre, np.eye(4)): + translation = world_to_nre[:3, 3] + logger.info( + f"Transforming map to local coordinates with translation: {translation}" + ) + + # Transform map coordinates + self._map.translate(translation[0], translation[1], translation[2]) + + logger.info( + f"Successfully loaded map from scene.map_data for {self.scene_id}" + ) + return self._map + + # Otherwise, try to get map from dataset._map_api (for datasets with map_api) + if self._dataset is None: + logger.warning( + "Cannot load map: dataset is not set and scene.map_data is not available" + ) + return None + + # Get map from dataset._map_api (refer to trajdata_artifact_converter.py) + try: + # Check if dataset includes map support + if ( + not hasattr(self._dataset, "incl_vector_map") + or not self._dataset.incl_vector_map + or not hasattr(self._dataset, "_map_api") + or self._dataset._map_api is None + ): + logger.warning( + "Dataset does not have map support enabled or map_api is unavailable" + ) + return None + + # Build map name: "{env_name}:{location}" + if not hasattr(self._scene, "location") or not self._scene.location: + logger.warning( + f"Scene {self.scene_id} has no location information, cannot load map" + ) + return None + + map_name = f"{self._scene.env_name}:{self._scene.location}" + + # Get vector_map_params (if exists) + vector_map_params = {} + if hasattr(self._dataset, "vector_map_params"): + vector_map_params = self._dataset.vector_map_params + + # Get map from map_api + vec_map = self._dataset._map_api.get_map(map_name, **vector_map_params) + + if vec_map is None: + logger.debug( + f"Scene {self.scene_id} (map_name: {map_name}) has no map data" + ) + return None + + # Create an independent copy of VectorMap for current scene to avoid modifying + # map objects in shared cache. This allows continued use of MapAPI's disk cache + # and index loading capabilities while preventing coordinate system pollution + # between multiple scenes through shared VectorMap instances. + self._map = copy.deepcopy(vec_map) + + # Important: Transform map to local coordinate system (NRE) + # Since trajectories are already converted to local coordinates, map also needs + # to be converted to match. This is consistent with USDZ format handling: + # both map and trajectories need to be converted to the same coordinate system + if self._rig is None: + # If rig is not loaded yet, load it first (this will set world_to_nre) + _ = self.rig + + world_to_nre = self._rig.world_to_nre + + # Check if world_to_nre contains rotation + rotation_matrix = world_to_nre[:3, :3] + translation = world_to_nre[:3, 3] + has_rotation = not np.allclose(rotation_matrix, np.eye(3)) + + if has_rotation: + logger.warning( + "world_to_nre contains rotation. Map transformation may need rotation handling. " + "Currently only applying translation." + ) + + # Transform all points in map (center.points, left_boundary.points, right_boundary.points, etc.) + # Note: Must transform before finalize, as finalize may rebuild certain data structures + # Important: Only transform X, Y coordinates, align Z coordinate to trajectory Z baseline + # Because map Z usually represents height relative to ground (usually 0), while trajectory Z + # represents altitude. After transformation, map Z should align with trajectory Z + # (both relative to first trajectory point's Z, usually 0 after transformation) + translation_xy = translation[:2] # Only use X, Y translation + + # Get Z coordinate of first trajectory point (transformed baseline, usually 0) + first_traj_z = ( + self.rig.trajectory.positions[0][2] + if len(self.rig.trajectory) > 0 + else 0.0 + ) + + logger.info( + f"Map coordinate transformation: " + f"translation_xy={translation_xy}, " + f"first_traj_z={first_traj_z:.2f}m, " + f"map Z will be aligned to trajectory Z baseline" + ) + + def transform_map_points(points: np.ndarray) -> np.ndarray: + """Transform map points: only transform X, Y, align Z to trajectory Z baseline""" + if ( + points is None + or len(points) == 0 + or points.ndim != 2 + or points.shape[1] < 3 + ): + return points + + points_copy = points.copy() + + # Transform X, Y coordinates + if has_rotation: + # If rotation exists, need to rotate X, Y + xy_rotated = ( + points_copy[:, :2] @ rotation_matrix[:2, :2].T + ) + translation_xy + points_copy[:, 0] = xy_rotated[:, 0] + points_copy[:, 1] = xy_rotated[:, 1] + else: + # Only translate X, Y + points_copy[:, 0] = points_copy[:, 0] + translation_xy[0] + points_copy[:, 1] = points_copy[:, 1] + translation_xy[1] + + # Z coordinate alignment: align map Z to trajectory Z baseline + # Map Z is usually height relative to ground (usually 0), + # after transformation should align with trajectory Z (both relative to first + # trajectory point's Z, usually 0 after transformation) + # So: new_z = original_z + first_traj_z + # If map Z=0 (ground), after transformation it becomes first_traj_z (usually 0) + points_copy[:, 2] = points_copy[:, 2] + first_traj_z + + return points_copy + + if hasattr(self._map, "lanes"): + for lane_idx, lane in enumerate(self._map.lanes): + # Transform center.points + if hasattr(lane, "center") and hasattr(lane.center, "points"): + points = lane.center.points + if points is not None and len(points) > 0: + try: + transformed_points = transform_map_points(points) + lane.center.points = transformed_points + except Exception as e: + logger.warning( + f"Failed to transform lane {lane_idx} center.points: {e}" + ) + + # Transform left_boundary.points + if hasattr(lane, "left_boundary") and hasattr( + lane.left_boundary, "points" + ): + points = lane.left_boundary.points + if points is not None and len(points) > 0: + try: + transformed_points = transform_map_points(points) + lane.left_boundary.points = transformed_points + except Exception as e: + logger.warning( + f"Failed to transform lane {lane_idx} left_boundary.points: {e}" + ) + + # Transform right_boundary.points + if hasattr(lane, "right_boundary") and hasattr( + lane.right_boundary, "points" + ): + points = lane.right_boundary.points + if points is not None and len(points) > 0: + try: + transformed_points = transform_map_points(points) + lane.right_boundary.points = transformed_points + except Exception as e: + logger.warning( + f"Failed to transform lane {lane_idx} right_boundary.points: {e}" + ) + + # Transform other possible point attributes (if any) + for attr_name in ["intersections", "crosswalks", "stop_lines"]: + if hasattr(lane, attr_name): + attr_value = getattr(lane, attr_name) + if isinstance(attr_value, list): + for item in attr_value: + if hasattr(item, "points"): + points = item.points + if points is not None and len(points) > 0: + try: + transformed_points = ( + transform_map_points(points) + ) + item.points = transformed_points + except Exception as e: + logger.debug( + f"Failed to transform {attr_name} points: {e}" + ) + + # If map needs finalize, call it (after transformation) + # Note: finalize may rebuild certain indices but won't change point coordinates + # But for safety, we verify transformation again after finalize + if hasattr(self._map, "__post_init__"): + self._map.__post_init__() + if hasattr(self._map, "compute_search_indices"): + self._map.compute_search_indices() + + # Verify again: check if first point's coordinates are still correct after finalize + if hasattr(self._map, "lanes") and len(self._map.lanes) > 0: + first_lane = self._map.lanes[0] + if hasattr(first_lane, "center") and hasattr( + first_lane.center, "points" + ): + first_map_point_after_finalize = ( + first_lane.center.points[0, :3] + if len(first_lane.center.points) > 0 + else None + ) + if first_map_point_after_finalize is not None: + # Check if Z coordinate aligns with trajectory (should both be first_traj_z, usually 0) + actual_z = first_map_point_after_finalize[2] + if abs(actual_z - first_traj_z) > 1.0: # Allow 1 meter error + logger.warning( + f"Map Z coordinate may have been reset after finalize. " + f"Expected Z≈{first_traj_z:.2f}m (aligned to trajectory Z baseline), " + f"got Z={actual_z:.2f}m. " + f"This may cause coordinate misalignment." + ) + # If alignment is correct, no need to log (reduce noise) + + # Fix data types (if needed) + if hasattr(self._map, "lanes"): + for lane in self._map.lanes: + if hasattr(lane, "next_lanes") and isinstance( + lane.next_lanes, list + ): + lane.next_lanes = set(lane.next_lanes) + if hasattr(lane, "prev_lanes") and isinstance( + lane.prev_lanes, list + ): + lane.prev_lanes = set(lane.prev_lanes) + if hasattr(lane, "adj_lanes_right") and isinstance( + lane.adj_lanes_right, list + ): + lane.adj_lanes_right = set(lane.adj_lanes_right) + if hasattr(lane, "adj_lanes_left") and isinstance( + lane.adj_lanes_left, list + ): + lane.adj_lanes_left = set(lane.adj_lanes_left) + + # Verify map transformation: check if first lane's first point is within reasonable range + if hasattr(self._map, "lanes") and len(self._map.lanes) > 0: + first_lane = self._map.lanes[0] + if hasattr(first_lane, "center") and hasattr( + first_lane.center, "points" + ): + first_map_point = ( + first_lane.center.points[0, :3] + if len(first_lane.center.points) > 0 + else None + ) + if first_map_point is not None: + # Map point should be near trajectory (within hundreds of meters) + distance_from_origin_xy = np.linalg.norm( + first_map_point[:2] + ) # Only check X, Y + distance_from_origin_xyz = np.linalg.norm( + first_map_point + ) # Check X, Y, Z + + # Get first trajectory point for comparison + first_traj_point = self.rig.trajectory.positions[0] + + logger.info( + f"Map transformation verification: " + f"first lane center point: {first_map_point}, " + f"first trajectory point: {first_traj_point}, " + f"distance (X,Y): {distance_from_origin_xy:.2f}m, " + f"distance (X,Y,Z): {distance_from_origin_xyz:.2f}m, " + f"Z difference: {abs(first_map_point[2] - first_traj_point[2]):.2f}m" + ) + + # Warn if Z coordinate is too far from trajectory Z + z_diff = abs(first_map_point[2] - first_traj_point[2]) + if z_diff > 10.0: # Z coordinate difference exceeds 10 meters + logger.warning( + f"Map Z coordinate may not be correctly aligned with trajectory. " + f"Map Z={first_map_point[2]:.2f}m, Trajectory Z={first_traj_point[2]:.2f}m, " + f"difference={z_diff:.2f}m. This may cause route generation to fail." + ) + + logger.info( + f"Successfully loaded map: {map_name} (transformed to local coordinate system)" + ) + return self._map + except Exception as e: + logger.error(f"Error loading map: {e}", exc_info=True) + return None + + @property + def metadata(self) -> Metadata: + """Create and return Metadata object""" + if self._metadata is not None: + return self._metadata + + # Extract metadata from scene + scene_id = self.scene_id + + # Ensure rig is loaded + rig = self.rig + + # Extract camera ID list from rig + camera_id_names = [] + if rig and rig.camera_ids: + camera_id_names = [camera_id.logical_name for camera_id in rig.camera_ids] + + # Calculate time range + if self._scene is not None: + dt = self._scene.dt + length_timesteps = self._scene.length_timesteps + base_timestamp_us = getattr(self, "_base_timestamp_us", 0.0) + time_range_start = float(base_timestamp_us) / 1e6 + time_range_end = ( + float(base_timestamp_us + length_timesteps * dt * 1e6) / 1e6 + ) + else: + time_range_start = float(rig.trajectory.time_range_us.start) / 1e6 + time_range_end = float(rig.trajectory.time_range_us.stop) / 1e6 + + # Create metadata + import uuid + from datetime import datetime + + self._metadata = Metadata( + scene_id=scene_id, + version_string="trajdata_direct", + training_date=datetime.now().strftime("%Y-%m-%d"), + dataset_hash=str(uuid.uuid4()), + uuid=str(uuid.uuid4()), + is_resumable=False, + sensors=Metadata.Sensors( + camera_ids=camera_id_names, + lidar_ids=[], + ), + logger=Metadata.Logger(), + time_range=Metadata.TimeRange( + start=time_range_start, + end=time_range_end, + ), + ) + + return self._metadata From 1a4965498fbe7a7179840ec35293d34c439443a1 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:28:25 +0000 Subject: [PATCH 03/17] refactor!: migrate RuntimeContext to use scene_id_to_idx mapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace heavy scene_id_to_data_source dict with lightweight scene_id_to_idx mapping in RuntimeContext. This enables efficient serialization and reduces memory overhead when passing context to worker processes. Changes: - RuntimeContext.scene_id_to_data_source → scene_id_to_idx - Build scene_id to trajdata index mapping once at startup - Workers can reconstruct data sources on-demand using the mapping Benefits: - Lightweight RuntimeContext (dict[str, int] vs dict[str, DataSource]) - Fully serializable with pickle for multiprocessing - Avoids duplicating heavy data objects across workers - Aligns with trajdata's index-based API BREAKING CHANGE: RuntimeContext.scene_id_to_data_source field replaced with scene_id_to_idx. Code accessing the old field must be updated to use the new scene_id_to_idx mapping and load data sources on-demand. --- .../alpasim_runtime/runtime_context.py | 72 +++++++++++++++---- 1 file changed, 58 insertions(+), 14 deletions(-) diff --git a/src/runtime/alpasim_runtime/runtime_context.py b/src/runtime/alpasim_runtime/runtime_context.py index eb1f4c12..bf7fada0 100644 --- a/src/runtime/alpasim_runtime/runtime_context.py +++ b/src/runtime/alpasim_runtime/runtime_context.py @@ -3,6 +3,7 @@ from __future__ import annotations +import logging import copy import math from dataclasses import dataclass @@ -19,10 +20,20 @@ gather_versions_from_addresses, validate_scenarios, ) -from alpasim_utils.artifact import Artifact from eval.schema import EvalConfig +logger = logging.getLogger(__name__) + +# Optional trajdata support +try: + from trajdata.dataset import UnifiedDataset + + TRAJDATA_AVAILABLE = True +except ImportError: + TRAJDATA_AVAILABLE = False + UnifiedDataset = None + ALL_SKIP_PER_WORKER_CONCURRENCY = 16 @@ -121,13 +132,13 @@ class RuntimeContext: """Immutable snapshot of all runtime state needed to dispatch simulation jobs. Built once during startup by ``build_runtime_context`` after config parsing, - service version probing, scenario validation, and address pool creation. + service version probing, scenario validation, and scene mapping creation. """ config: SimulatorConfig eval_config: EvalConfig version_ids: RolloutMetadata.VersionIds - scene_id_to_artifact_path: dict[str, str] + scene_id_to_idx: dict[str, int] pools: dict[str, AddressPool] max_in_flight: int @@ -147,7 +158,6 @@ async def build_runtime_context( user_config_path: str, network_config_path: str, eval_config_path: str, - usdz_glob: str, validate_config_scenes: bool = True, ) -> RuntimeContext: """Build the RuntimeContext by parsing configs, probing services, and validating scenarios. @@ -156,20 +166,26 @@ async def build_runtime_context( 1. Parse user and network configs. 2. Probe all service addresses for version IDs. 3. Validate scenario compatibility (unless *validate_config_scenes* is False). - 4. Discover scene artifacts from *usdz_glob*. + 4. Create UnifiedDataset from data_source config and build scene data sources. 5. Create address pools and compute max in-flight concurrency. Args: user_config_path: Path to user YAML config. network_config_path: Path to network YAML config. eval_config_path: Path to evaluation YAML config. - usdz_glob: Glob pattern for USDZ artifact discovery. validate_config_scenes: If False, skip scene compatibility checks (useful for daemon mode where scenes come from requests). """ config = parse_simulator_config(user_config_path, network_config_path) eval_config = typed_parse_config(eval_config_path, EvalConfig) + # Require data_source in config (unified data flow) + if config.user.data_source is None: + raise ValueError( + "No data source specified in user config. " + "Please set 'data_source' in your YAML config file." + ) + version_ids = await gather_versions_from_addresses( config.network, config.user.endpoints, @@ -185,13 +201,41 @@ async def build_runtime_context( ) await validate_scenarios(config_for_validation) - scene_id_to_artifact_path = { - scene_id: artifact.source - for scene_id, artifact in Artifact.discover_from_glob( - usdz_glob, - smooth_trajectories=config.user.smooth_trajectories, - ).items() - } + # Create UnifiedDataset and build scene_id to data source mapping + logger.info("Creating UnifiedDataset from config") + if not TRAJDATA_AVAILABLE: + raise ImportError( + "trajdata is required for data source loading. " "Please install trajdata." + ) + + data_source_config = config.user.data_source + dataset = UnifiedDataset( + desired_data=data_source_config.desired_data, + data_dirs=data_source_config.data_dirs, + cache_location=data_source_config.cache_location, + incl_vector_map=data_source_config.incl_vector_map, + rebuild_cache=data_source_config.rebuild_cache, + rebuild_maps=data_source_config.rebuild_maps, + desired_dt=data_source_config.desired_dt, + num_workers=data_source_config.num_workers, + ) + logger.info( + f"Created UnifiedDataset with {dataset.num_scenes()} scenes, " + f"desired_data={data_source_config.desired_data}" + ) + + # Build scene_id to index mapping (once, in main process) + scene_id_to_idx = {} + num_scenes = dataset.num_scenes() + for idx in range(num_scenes): + try: + scene = dataset.get_scene(idx) + scene_id_to_idx[scene.name] = idx + except Exception as e: + logger.warning(f"Failed to get scene at index {idx}: {e}") + continue + logger.info(f"Built scene_id mapping for {len(scene_id_to_idx)} scenes") + pools = create_address_pools(config) max_in_flight = compute_max_in_flight(pools, config) @@ -199,7 +243,7 @@ async def build_runtime_context( config=config, eval_config=eval_config, version_ids=version_ids, - scene_id_to_artifact_path=scene_id_to_artifact_path, + scene_id_to_idx=scene_id_to_idx, pools=pools, max_in_flight=max_in_flight, ) From bcfaaa35506f7d5b9584939b0c0f95f8446f2e98 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:28:42 +0000 Subject: [PATCH 04/17] feat: update DaemonEngine to support lazy scene loading Adapt DaemonEngine to work with new RuntimeContext and support on-demand scene data source loading from trajdata. Changes: - Store scene_id_to_idx mapping from RuntimeContext - Create UnifiedDataset at engine startup for scene loading - Add _get_data_source() method for lazy loading with caching - Update build_pending_jobs_from_request to use callback pattern - Pre-create scene_cache to avoid pickle errors This enables daemon mode to efficiently handle multiple scenes without loading all data upfront. --- src/runtime/alpasim_runtime/daemon/engine.py | 126 +++++++++++++++--- .../alpasim_runtime/daemon/scheduler.py | 2 +- 2 files changed, 111 insertions(+), 17 deletions(-) diff --git a/src/runtime/alpasim_runtime/daemon/engine.py b/src/runtime/alpasim_runtime/daemon/engine.py index 7985dc6d..d41ac1d4 100644 --- a/src/runtime/alpasim_runtime/daemon/engine.py +++ b/src/runtime/alpasim_runtime/daemon/engine.py @@ -5,8 +5,17 @@ import logging from collections import defaultdict +from typing import Callable from uuid import uuid4 +try: + from trajdata.dataset import UnifiedDataset + + TRAJDATA_AVAILABLE = True +except ImportError: + TRAJDATA_AVAILABLE = False + UnifiedDataset = None + from alpasim_grpc.v0 import logging_pb2, runtime_pb2 from alpasim_runtime.address_pool import AddressPool from alpasim_runtime.daemon.scheduler import DaemonScheduler, DaemonUnavailableError @@ -16,6 +25,8 @@ ) from alpasim_runtime.worker.ipc import JobResult, PendingRolloutJob from alpasim_runtime.worker.runtime import WorkerRuntime, start_worker_runtime +from alpasim_utils.scene_data_source import SceneDataSource +from alpasim_utils.trajdata_data_source import TrajdataDataSource from eval.data import AggregationType @@ -104,30 +115,34 @@ class InvalidRequestError(ValueError): class UnknownSceneError(InvalidRequestError): - """Raised when a simulation request references a scene_id with no known artifact.""" + """Raised when a simulation request references a scene_id with no known data source.""" def __init__(self, scene_id: str): - super().__init__(f"No artifact found for scene_id: {scene_id}") + super().__init__(f"No data source found for scene_id: {scene_id}") self.scene_id = scene_id def build_pending_jobs_from_request( request: runtime_pb2.SimulationRequest, - scene_id_to_artifact_path: dict[str, str], + get_data_source: Callable[[str], SceneDataSource], ) -> list[PendingRolloutJob]: """Expand a SimulationRequest into individual PendingRolloutJob entries. Each RolloutSpec is expanded by its ``nr_rollouts`` count. Specs with ``nr_rollouts=0`` are silently dropped with a warning. + Args: + request: The simulation request to expand. + get_data_source: Callable that returns a SceneDataSource for a given scene_id. + Should raise UnknownSceneError if the scene_id is not found. + Raises: - UnknownSceneError: If a spec references a scene_id not present in - *scene_id_to_artifact_path*. + UnknownSceneError: If a spec references an unknown scene_id. """ jobs: list[PendingRolloutJob] = [] for spec_index, spec in enumerate(request.rollout_specs): - if spec.scenario_id not in scene_id_to_artifact_path: - raise UnknownSceneError(spec.scenario_id) + # This will raise UnknownSceneError if scene_id is not found + data_source = get_data_source(spec.scenario_id) if spec.nr_rollouts == 0: logger.warning( @@ -142,7 +157,7 @@ def build_pending_jobs_from_request( job_id=uuid4().hex, scene_id=spec.scenario_id, rollout_spec_index=spec_index, - artifact_path=scene_id_to_artifact_path[spec.scenario_id], + data_source=data_source, ) ) return jobs @@ -166,19 +181,20 @@ def __init__( user_config: str, network_config: str, eval_config: str, - usdz_glob: str, log_dir: str, validate_config_scenes: bool = True, ) -> None: self._user_config_path = user_config self._network_config_path = network_config self._eval_config_path = eval_config - self._usdz_glob = usdz_glob self._log_dir = log_dir self._validate_config_scenes = validate_config_scenes self._version_ids: logging_pb2.RolloutMetadata.VersionIds | None = None - self._scene_id_to_artifact_path: dict[str, str] = {} + self._config = None # Will be set during startup + self._dataset: UnifiedDataset | None = None + self._scene_id_to_idx: dict[str, int] = {} + self._scene_id_to_data_source: dict[str, SceneDataSource] = {} self._scheduler: DaemonScheduler | None = None self._worker_runtime: WorkerRuntime | None = None self._started = False @@ -189,12 +205,63 @@ def version_ids(self) -> logging_pb2.RolloutMetadata.VersionIds: raise RuntimeError("daemon is not started") return self._version_ids + def _get_data_source(self, scene_id: str) -> SceneDataSource: + """Get or create a data source for the given scene_id.""" + # Check cache first + if scene_id in self._scene_id_to_data_source: + return self._scene_id_to_data_source[scene_id] + + # Lazy load from dataset + if self._dataset is None or not TRAJDATA_AVAILABLE: + raise RuntimeError(f"Dataset not initialized, cannot load scene {scene_id}") + + if self._config is None: + raise RuntimeError("Config not initialized") + + scene_idx = self._scene_id_to_idx.get(scene_id) + if scene_idx is None: + raise UnknownSceneError(scene_id) + + try: + scene = self._dataset.get_scene(scene_idx) + if scene is None: + raise UnknownSceneError(scene_id) + + # Get asset_base_path from config + asset_base_path = None + if self._config.user.data_source is not None: + asset_base_path = self._config.user.data_source.asset_base_path + + # Create scene_cache (pre-create to avoid pickle errors) + scene_cache = self._dataset.cache_class( + self._dataset.cache_path, scene, self._dataset.augmentations + ) + scene_cache.set_obs_format(self._dataset.obs_format) + + # Create TrajdataDataSource + data_source = TrajdataDataSource.from_trajdata_scene( + scene=scene, + dataset=None, # Don't pass dataset to avoid pickle errors + scene_cache=scene_cache, + smooth_trajectories=self._config.user.smooth_trajectories, + asset_base_path=asset_base_path, + ) + + # Cache for future use + self._scene_id_to_data_source[scene_id] = data_source + logger.debug(f"Loaded data source for scene {scene_id}") + return data_source + + except Exception as e: + logger.error(f"Failed to load scene {scene_id}: {e}") + raise UnknownSceneError(scene_id) + async def startup(self) -> None: """Initialize the runtime context, start workers, and begin scheduling. Builds the RuntimeContext (parses configs, probes service versions, - validates scenarios, discovers scene artifacts), then creates the - worker runtime and daemon scheduler. Idempotent: subsequent calls + validates scenarios, creates scene mapping from trajdata), then creates + the worker runtime and daemon scheduler. Idempotent: subsequent calls after the first are no-ops. """ if self._started: @@ -204,10 +271,36 @@ async def startup(self) -> None: user_config_path=self._user_config_path, network_config_path=self._network_config_path, eval_config_path=self._eval_config_path, - usdz_glob=self._usdz_glob, validate_config_scenes=self._validate_config_scenes, ) + # Create UnifiedDataset for on-demand scene loading + if not TRAJDATA_AVAILABLE: + raise ImportError( + "trajdata is required for data source loading. " + "Please install trajdata." + ) + + data_source_config = runtime_context.config.user.data_source + if data_source_config is None: + raise ValueError("data_source is required in user config") + + self._dataset = UnifiedDataset( + desired_data=data_source_config.desired_data, + data_dirs=data_source_config.data_dirs, + cache_location=data_source_config.cache_location, + incl_vector_map=data_source_config.incl_vector_map, + rebuild_cache=data_source_config.rebuild_cache, + rebuild_maps=data_source_config.rebuild_maps, + desired_dt=data_source_config.desired_dt, + num_workers=data_source_config.num_workers, + ) + logger.info(f"Created UnifiedDataset with {self._dataset.num_scenes()} scenes") + + # Store scene_id to index mapping from RuntimeContext + self._scene_id_to_idx = runtime_context.scene_id_to_idx + self._config = runtime_context.config + num_consumers_per_worker = compute_num_consumers_per_worker( max_in_flight=runtime_context.max_in_flight, nr_workers=runtime_context.config.user.nr_workers, @@ -228,7 +321,6 @@ async def startup(self) -> None: ) self._version_ids = runtime_context.version_ids - self._scene_id_to_artifact_path = runtime_context.scene_id_to_artifact_path self._worker_runtime = worker_runtime self._scheduler = scheduler self._started = True @@ -252,7 +344,9 @@ async def simulate( assert self._scheduler is not None request_id = uuid4().hex - jobs = build_pending_jobs_from_request(request, self._scene_id_to_artifact_path) + + # Use instance method for getting data sources + jobs = build_pending_jobs_from_request(request, self._get_data_source) driver_pool: AddressPool | None = None if request.available_drivers: diff --git a/src/runtime/alpasim_runtime/daemon/scheduler.py b/src/runtime/alpasim_runtime/daemon/scheduler.py index 35099758..7ef1b96f 100644 --- a/src/runtime/alpasim_runtime/daemon/scheduler.py +++ b/src/runtime/alpasim_runtime/daemon/scheduler.py @@ -146,7 +146,7 @@ async def dispatch_once(self) -> None: job_id=pending_job.job_id, scene_id=pending_job.scene_id, rollout_spec_index=pending_job.rollout_spec_index, - artifact_path=pending_job.artifact_path, + data_source=pending_job.data_source, endpoints=ServiceEndpoints( driver=acquired["driver"], sensorsim=acquired["sensorsim"], From 1cb807aeb0b74170a1bcabe9c2fee79ff7d2d987 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:29:15 +0000 Subject: [PATCH 05/17] feat: add prepare_data CLI for trajdata cache preprocessing Add command-line tool for preprocessing scene data and building trajdata cache before running simulations. Features: - Basic mode: preprocess all scenes in a dataset - YAML config mode: batch process specific scenes from YAML files - Central token mode: process scenes around specific tokens (NuPlan) - Support for smooth_trajectories parameter - Configurable cache rebuilding and vector map inclusion Usage: # Basic preprocessing python -m alpasim_runtime.prepare_data --user-config user.yaml # With explicit parameters python -m alpasim_runtime.prepare_data \ --desired-data nuplan_test \ --data-dir /path/to/data \ --cache-location /tmp/cache This preprocessing step improves simulation startup time by pre-building the trajdata cache. --- .../alpasim_runtime/prepare_data/__init__.py | 56 ++ .../alpasim_runtime/prepare_data/__main__.py | 720 ++++++++++++++++++ 2 files changed, 776 insertions(+) create mode 100644 src/runtime/alpasim_runtime/prepare_data/__init__.py create mode 100644 src/runtime/alpasim_runtime/prepare_data/__main__.py diff --git a/src/runtime/alpasim_runtime/prepare_data/__init__.py b/src/runtime/alpasim_runtime/prepare_data/__init__.py new file mode 100644 index 00000000..cd4c763a --- /dev/null +++ b/src/runtime/alpasim_runtime/prepare_data/__init__.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025-2026 NVIDIA Corporation + +""" +Data preprocessing module for building trajdata cache. + +This module provides tools for preparing scene data before running simulations. +It supports: + +1. Basic preprocessing - Build trajdata cache for all scenes in a dataset +2. YAML config preprocessing - Batch process specific scenes based on YAML configs +3. Central token mode - Process scenes around specific central tokens (NuPlan) + +Main functions: + - preprocess_basic: Basic preprocessing for all scenes + - preprocess_from_yaml_configs: Batch preprocessing from YAML configs + - load_yaml_configs: Load YAML configuration files + +Example usage: + + from alpasim_runtime.prepare_data import preprocess_basic, preprocess_from_yaml_configs + + # Basic preprocessing + preprocess_basic( + desired_data=["nuplan_test"], + data_dirs={"nuplan_test": "/path/to/nuplan"}, + cache_location="/path/to/cache", + ) + + # YAML config batch preprocessing + preprocess_from_yaml_configs( + config_dir=Path("/path/to/configs"), + cache_location="/path/to/cache", + data_dirs={"nuplan_test": "/path/to/nuplan"}, + env_name="nuplan_test", + num_timesteps_before=30, + num_timesteps_after=80, + ) + +CLI usage: + python -m alpasim_runtime.prepare_data --help +""" + +from alpasim_runtime.prepare_data.__main__ import ( + load_yaml_configs, + main, + preprocess_basic, + preprocess_from_yaml_configs, +) + +__all__ = [ + "preprocess_basic", + "preprocess_from_yaml_configs", + "load_yaml_configs", + "main", +] diff --git a/src/runtime/alpasim_runtime/prepare_data/__main__.py b/src/runtime/alpasim_runtime/prepare_data/__main__.py new file mode 100644 index 00000000..5200349f --- /dev/null +++ b/src/runtime/alpasim_runtime/prepare_data/__main__.py @@ -0,0 +1,720 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025-2026 NVIDIA Corporation + +""" +Data preprocessing CLI for building trajdata cache. + +This module provides a command-line interface for preparing scene data before +running simulations. It supports: + +1. Basic preprocessing - Build trajdata cache for all scenes in a dataset +2. YAML config preprocessing - Batch process specific scenes based on YAML configs +3. Central token mode - Process scenes around specific central tokens (NuPlan) + +Usage Examples: + + # Basic preprocessing using user-config + python -m alpasim_runtime.prepare_data --user-config user.yaml + + # Basic preprocessing with explicit parameters + python -m alpasim_runtime.prepare_data \\ + --desired-data nuplan_test \\ + --data-dir /path/to/nuplan \\ + --cache-location /path/to/cache + + # Rebuild cache even if it exists + python -m alpasim_runtime.prepare_data --user-config user.yaml --rebuild-cache +""" + +from __future__ import annotations + +import argparse +import logging +import sys +import time +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, List, Optional + +import yaml + +logger = logging.getLogger(__name__) + +# Optional trajdata import +try: + from trajdata.dataset import UnifiedDataset + + TRAJDATA_AVAILABLE = True + +except ImportError: + TRAJDATA_AVAILABLE = False + UnifiedDataset = None + env_utils = None + + +def load_yaml_configs(config_dir: Path) -> Dict[str, List[Dict[str, str]]]: + """ + Load all yaml configuration files and group them by central_log. + + Supports both simple YAML files and NuPlan-generated files with Python object tags. + Uses a custom loader to handle Python objects without requiring module imports. + + Args: + config_dir: Directory containing yaml configuration files. + + Returns: + Dict where key is central_log and value is the list of central_tokens configs for that log. + """ + configs_by_log = defaultdict(list) + + yaml_files = list(config_dir.glob("*.yaml")) + logger.info(f"Found {len(yaml_files)} yaml configuration files.") + + # Custom YAML loader that converts unknown Python objects to dicts + class SafeLoaderWithObjects(yaml.SafeLoader): + """Custom YAML loader that treats Python objects as plain dicts.""" + + pass + + def python_object_constructor(loader, tag_suffix, node): + """Convert Python object tags to plain dicts. + + Args: + loader: YAML loader instance + tag_suffix: Tag suffix (for multi_constructor, ignored for single constructor) + node: YAML node to construct + """ + return loader.construct_mapping(node, deep=True) + + def python_tuple_constructor(loader, tag_suffix, node): + """Convert Python tuple tags to lists. + + Args: + loader: YAML loader instance + tag_suffix: Tag suffix (ignored) + node: YAML node to construct + """ + return loader.construct_sequence(node, deep=True) + + # Register constructors for Python objects and tuples + # Note: add_multi_constructor passes 3 args (loader, tag_suffix, node) + SafeLoaderWithObjects.add_multi_constructor( + "tag:yaml.org,2002:python/object", python_object_constructor + ) + SafeLoaderWithObjects.add_multi_constructor( + "tag:yaml.org,2002:python/tuple", python_tuple_constructor + ) + + for yaml_file in yaml_files: + try: + # Load YAML with custom loader that handles Python objects + config = yaml.load(yaml_file.read_text(), Loader=SafeLoaderWithObjects) + + # Support both attribute-style (config.central_log) and dict-style access + if hasattr(config, "central_log"): + central_log = config.central_log + central_tokens = config.central_tokens + else: + central_log = config.get("central_log", "") + central_tokens = config.get("central_tokens", []) + + if not central_log or not central_tokens: + logger.warning( + f"{yaml_file.name} is missing central_log or central_tokens, skipping." + ) + continue + + # Extract first central_token + configs_by_log[central_log].append( + { + "central_token": central_tokens[0], + "logfile": central_log, + "yaml_file": str(yaml_file), + } + ) + # Extract every central_token + # for token in central_tokens: + # configs_by_log[central_log].append({ + # 'central_token': token, + # 'logfile': central_log, + # 'yaml_file': str(yaml_file), + # }) + + except Exception as e: + logger.error(f"Failed to load {yaml_file.name}: {e}") + continue + + logger.info( + f"\nAfter grouping by central_log, there are {len(configs_by_log)} different log files." + ) + for log, configs in configs_by_log.items(): + logger.info(f" {log}: {len(configs)} central tokens") + + return dict(configs_by_log) + + +def preprocess_from_yaml_configs( + config_dir: Path, + cache_location: str, + data_dirs: Dict[str, str], + env_name: str = "nuplan_test", + rebuild_cache: bool = True, + rebuild_maps: bool = False, + num_workers: int = 1, + desired_dt: float = 0.5, + num_timesteps_before: int = 30, + num_timesteps_after: int = 80, + verbose: bool = True, +) -> bool: + """ + Batch preprocess data based on YAML configuration files. + + This function reads YAML config files containing central_log and central_tokens, + and processes only those specific scenes. This is useful for processing + specific scenarios without loading the entire dataset. + + Args: + config_dir: Directory containing YAML configuration files. + cache_location: Path to cache directory. + data_dirs: Dictionary of dataset name to data directory paths. + env_name: Environment name (e.g., "nuplan_test"). + rebuild_cache: Whether to rebuild cache. + rebuild_maps: Whether to rebuild maps. + num_workers: Number of worker processes. + desired_dt: Desired timestep duration in seconds. + num_timesteps_before: Number of timesteps before the central token. + num_timesteps_after: Number of timesteps after the central token. + verbose: Whether to show verbose logs. + + Returns: + True if successful, False otherwise. + """ + if not TRAJDATA_AVAILABLE: + logger.error("trajdata is not installed. Please install it first.") + return False + + # Load all YAML configs + configs_by_log = load_yaml_configs(config_dir) + + if not configs_by_log: + logger.error("No valid YAML configuration files found.") + return False + + # Merge all configs into a single central_tokens_config list + all_central_tokens_config: List[Dict[str, Any]] = [] + for _, configs in configs_by_log.items(): + for cfg in configs: + all_central_tokens_config.append( + { + "central_token": cfg["central_token"], + "logfile": cfg["logfile"], + "num_timesteps_before": num_timesteps_before, + "num_timesteps_after": num_timesteps_after, + } + ) + + logger.info(f"Total {len(all_central_tokens_config)} central tokens to process") + + try: + # Create cache directory + Path(cache_location).mkdir(parents=True, exist_ok=True) + + # Create UnifiedDataset (this triggers cache building) + logger.info("Creating UnifiedDataset with YAML configs...") + start_time = time.perf_counter() + + dataset = UnifiedDataset( + dataset_kwargs={ + "central_tokens_config": all_central_tokens_config, + "num_timesteps_before": num_timesteps_before, + "num_timesteps_after": num_timesteps_after, + }, + desired_data=[env_name], + cache_location=cache_location, + rebuild_cache=rebuild_cache, + rebuild_maps=rebuild_maps, + require_map_cache=False, + num_workers=num_workers, + desired_dt=desired_dt, + verbose=verbose, + data_dirs=data_dirs, + ) + + elapsed = time.perf_counter() - start_time + + # Get both scene index count and total dataset length + num_scenes = dataset.num_scenes() + logger.info("=" * 80) + logger.info("Preprocessing completed!") + logger.info(f" Num Scenes: {num_scenes}") + logger.info(f" Time elapsed: {elapsed:.2f} seconds") + logger.info("=" * 80) + + return True + + except Exception as e: + logger.error(f"Error during preprocessing: {e}") + if verbose: + import traceback + + traceback.print_exc() + return False + + +def preprocess_basic( + desired_data: List[str], + data_dirs: Dict[str, str], + cache_location: str, + rebuild_cache: bool = False, + rebuild_maps: bool = False, + incl_vector_map: bool = True, + desired_dt: float = 0.1, + num_workers: int = 1, + verbose: bool = True, + list_scenes: bool = False, + smooth_trajectories: bool = True, +) -> bool: + """ + Basic preprocessing - build trajdata cache for all scenes. + + Args: + desired_data: List of dataset names to load. + data_dirs: Dict mapping dataset names to their data directories. + cache_location: Path to trajdata cache directory. + rebuild_cache: Whether to force rebuild cache. + rebuild_maps: Whether to rebuild map cache. + incl_vector_map: Whether to include vector maps. + desired_dt: Desired time delta between frames in seconds. + num_workers: Number of workers for data loading. + verbose: Whether to show verbose output. + list_scenes: Whether to list all available scenes after preparation. + + Returns: + True if successful, False otherwise. + """ + if not TRAJDATA_AVAILABLE: + logger.error("trajdata is not installed. Please install it first.") + return False + + logger.info("Data source configuration:") + logger.info(f" desired_data: {desired_data}") + logger.info(f" data_dirs: {data_dirs}") + logger.info(f" cache_location: {cache_location}") + logger.info(f" rebuild_cache: {rebuild_cache}") + logger.info(f" desired_dt: {desired_dt}") + logger.info(f" smooth_trajectories: {smooth_trajectories}") + + # Create cache directory + cache_path = Path(cache_location) + cache_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Cache directory: {cache_path}") + + # Build UnifiedDataset (this triggers cache building) + logger.info("Creating UnifiedDataset (this may take a while)...") + start_time = time.perf_counter() + + try: + dataset = UnifiedDataset( + desired_data=desired_data, + data_dirs=data_dirs, + cache_location=cache_location, + incl_vector_map=incl_vector_map, + rebuild_cache=rebuild_cache, + rebuild_maps=rebuild_maps, + # require_map_cache=False, + desired_dt=desired_dt, + num_workers=num_workers, + verbose=verbose, + dataset_kwargs={"smooth_trajectories": smooth_trajectories}, + ) + except Exception as e: + logger.error(f"Failed to create UnifiedDataset: {e}") + import traceback + + traceback.print_exc() + return False + + elapsed = time.perf_counter() - start_time + logger.info(f"UnifiedDataset created in {elapsed:.2f} seconds") + + # Get both scene index count and total dataset length + num_scenes = dataset.num_scenes() + logger.info(f"Scene files (logs): {num_scenes}") + + # List scenes if requested + if list_scenes and num_scenes > 0: + logger.info("Available scenes:") + max_display = min(num_scenes, 100) + for i in range(max_display): + try: + scene = dataset.get_scene(i) + logger.info(f" [{i}] {scene.name}") + except Exception as e: + logger.warning(f" [{i}] (failed to load: {e})") + + if num_scenes > 100: + logger.info(f" ... and {num_scenes - 100} more scenes") + + logger.info("Data preparation complete!") + return True + + +def create_arg_parser() -> argparse.ArgumentParser: + """Create argument parser for prepare_data CLI.""" + parser = argparse.ArgumentParser( + description="Prepare scene data and build trajdata cache for alpasim simulations.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + + # Mode selection + mode_group = parser.add_argument_group("Mode Selection") + mode_group.add_argument( + "--user-config", + type=str, + help="Path to user config YAML file containing data_source configuration", + ) + + # Data source parameters + data_group = parser.add_argument_group("Data Source") + data_group.add_argument( + "--desired-data", + type=str, + nargs="+", + help="List of dataset names to prepare (e.g., nuplan_test, waymo_val, usdz)", + ) + data_group.add_argument( + "--data-dir", + type=str, + action="append", + dest="data_dirs", + help="Data directory (format: dataset_name=/path/to/data or just /path/to/data)", + ) + data_group.add_argument( + "--cache-location", + type=str, + help="Path to trajdata cache directory", + ) + + # Preprocessing options + preprocess_group = parser.add_argument_group("Preprocessing Options") + preprocess_group.add_argument( + "--rebuild-cache", + action="store_true", + help="Force rebuild cache even if it already exists", + ) + preprocess_group.add_argument( + "--rebuild-maps", + action="store_true", + help="Force rebuild map cache", + ) + preprocess_group.add_argument( + "--desired-dt", + type=float, + default=0.1, + help="Desired timestep duration in seconds (default: 0.1)", + ) + preprocess_group.add_argument( + "--num-workers", + type=int, + default=1, + help="Number of worker processes (default: 1)", + ) + preprocess_group.add_argument( + "--no-vector-map", + action="store_true", + help="Exclude vector maps (default: include)", + ) + preprocess_group.add_argument( + "--smooth-trajectories", + type=str, + choices=["true", "false", "True", "False"], + default=None, + help="Enable/disable cubic spline smoothing for trajectories (default: True)", + ) + + # YAML config mode options (NuPlan specific) + yaml_group = parser.add_argument_group("YAML Config Mode (NuPlan)") + yaml_group.add_argument( + "--num-timesteps-before", + type=int, + default=30, + help="Number of timesteps before central token (default: 30)", + ) + yaml_group.add_argument( + "--num-timesteps-after", + type=int, + default=80, + help="Number of timesteps after central token (default: 80)", + ) + + # Output options + output_group = parser.add_argument_group("Output Options") + output_group.add_argument( + "--validate-only", + action="store_true", + help="Only validate configuration without building cache", + ) + output_group.add_argument( + "--list-scenes", + action="store_true", + help="List all available scenes after preparation", + ) + output_group.add_argument( + "--log-level", + type=str, + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Logging level (default: INFO)", + ) + output_group.add_argument( + "--verbose", + action="store_true", + default=True, + help="Show verbose output (default: True)", + ) + output_group.add_argument( + "--quiet", + action="store_true", + help="Suppress verbose output", + ) + + return parser + + +def load_config_from_file(config_path: str) -> Dict[str, Any]: + """Load data source configuration from user config file.""" + from alpasim_runtime.config import UserSimulatorConfig, typed_parse_config + + user_config = typed_parse_config(config_path, UserSimulatorConfig) + + if user_config.data_source is None: + raise ValueError( + f"No data_source configuration found in {config_path}. " + "Please add a 'data_source' section to your user config." + ) + + ds = user_config.data_source + return { + "desired_data": ds.desired_data, + "data_dirs": ds.data_dirs, + "cache_location": ds.cache_location, + "config_dir": ds.config_dir, + "incl_vector_map": ds.incl_vector_map, + "rebuild_cache": ds.rebuild_cache, + "rebuild_maps": ds.rebuild_maps, + "desired_dt": ds.desired_dt, + "num_workers": ds.num_workers, + "num_timesteps_before": ds.num_timesteps_before, + "num_timesteps_after": ds.num_timesteps_after, + "smooth_trajectories": user_config.smooth_trajectories, + } + + +def parse_data_dirs( + data_dirs_args: Optional[List[str]], desired_data: List[str] +) -> Dict[str, str]: + """ + Parse data directory arguments into a dict. + + Supports two formats: + - "dataset_name=/path/to/data" - explicit mapping + - "/path/to/data" - auto-map to desired_data entries in order + """ + if not data_dirs_args: + return {} + + result: Dict[str, str] = {} + + for i, arg in enumerate(data_dirs_args): + if "=" in arg: + # Explicit mapping: dataset_name=/path/to/data + parts = arg.split("=", 1) + result[parts[0]] = parts[1] + else: + # Implicit mapping: use desired_data order + if i < len(desired_data): + result[desired_data[i]] = arg + else: + # Use as default for remaining datasets + for ds in desired_data[len(result) :]: + if ds not in result: + result[ds] = arg + + return result + + +def run_yaml_batch_preprocessing( + config_dir: str, + desired_data: List[str], + data_dirs: Dict[str, str], + cache_location: str, + rebuild_cache: bool = False, + rebuild_maps: bool = False, + num_workers: int = 1, + desired_dt: float = 0.5, + num_timesteps_before: int = 30, + num_timesteps_after: int = 80, + verbose: bool = True, +) -> int: + """ + Execute YAML batch preprocessing mode. + + This is a unified entry point for YAML-based batch preprocessing, + whether triggered from user config file or CLI arguments. + + Args: + config_dir: Directory containing YAML scene config files. + desired_data: List of dataset names (first one will be used as env_name). + data_dirs: Dictionary mapping dataset names to data directories. + cache_location: Path to cache directory. + rebuild_cache: Whether to rebuild cache. + rebuild_maps: Whether to rebuild maps. + num_workers: Number of worker processes. + desired_dt: Desired timestep duration in seconds. + num_timesteps_before: Number of timesteps before central token. + num_timesteps_after: Number of timesteps after central token. + verbose: Whether to show verbose output. + + Returns: + Exit code (0 for success, 1 for failure). + """ + logger.info("Using YAML config batch preprocessing mode") + + success = preprocess_from_yaml_configs( + config_dir=Path(config_dir), + cache_location=cache_location, + data_dirs=data_dirs, + env_name=desired_data[0], # Use first dataset as environment name + rebuild_cache=rebuild_cache, + rebuild_maps=rebuild_maps, + num_workers=num_workers, + desired_dt=desired_dt, + num_timesteps_before=num_timesteps_before, + num_timesteps_after=num_timesteps_after, + verbose=verbose, + ) + + return 0 if success else 1 + + +def main(arg_list: Optional[List[str]] = None) -> int: + """Main entry point for prepare_data CLI.""" + parser = create_arg_parser() + args = parser.parse_args(arg_list) + + # Configure logging + log_level = getattr(logging, args.log_level.upper(), logging.INFO) + logging.basicConfig( + level=log_level, + format="%(asctime)s.%(msecs)03d %(levelname)s:\t%(message)s", + datefmt="%H:%M:%S", + ) + + verbose = args.verbose and not args.quiet + + logger.info("=" * 60) + logger.info("Alpasim Data Preparation Tool") + logger.info("=" * 60) + + if not TRAJDATA_AVAILABLE: + logger.error("trajdata is not installed. Please install it first:") + logger.error("pip install trajdata") + return 1 + + # Determine mode and load configuration + try: + if args.user_config: + logger.info(f"Loading configuration from: {args.user_config}") + config = load_config_from_file(args.user_config) + + # Override with command line args if provided + if args.rebuild_cache: + config["rebuild_cache"] = True + if args.rebuild_maps: + config["rebuild_maps"] = True + if args.smooth_trajectories is not None: + smooth_value = args.smooth_trajectories.lower() in ["true", "1", "yes"] + config["smooth_trajectories"] = smooth_value + logger.info( + f"Command line overrides config: smooth_trajectories={smooth_value}" + ) + + # Check if user config contains config_dir -> automatically use Mode 2 + if config.get("config_dir") is not None: + logger.info("Detected 'config_dir' in user config") + return run_yaml_batch_preprocessing( + config_dir=config["config_dir"], + desired_data=config["desired_data"], + data_dirs=config["data_dirs"], + cache_location=config["cache_location"], + rebuild_cache=config.get("rebuild_cache", False), + rebuild_maps=config.get("rebuild_maps", False), + num_workers=config.get("num_workers", 8), + desired_dt=config.get("desired_dt", 0.5), + num_timesteps_before=config.get("num_timesteps_before", 30), + num_timesteps_after=config.get("num_timesteps_after", 80), + verbose=verbose, + ) + # Otherwise continue with Mode 1 (basic mode) + + else: + if not args.desired_data: + logger.error("Either --user-config, or --desired-data must be provided") + return 1 + if not args.data_dirs: + logger.error("--data-dir is required when not using --user-config") + return 1 + if not args.cache_location: + logger.error( + "--cache-location is required when not using --user-config" + ) + return 1 + + data_dirs = parse_data_dirs(args.data_dirs, args.desired_data) + incl_vector_map = not args.no_vector_map + + if args.smooth_trajectories is not None: + smooth_value = args.smooth_trajectories.lower() in ["true", "1", "yes"] + else: + smooth_value = True + + config = { + "desired_data": args.desired_data, + "data_dirs": data_dirs, + "cache_location": args.cache_location, + "incl_vector_map": incl_vector_map, + "rebuild_cache": args.rebuild_cache, + "rebuild_maps": args.rebuild_maps, + "desired_dt": args.desired_dt, + "num_workers": args.num_workers, + "smooth_trajectories": smooth_value, + } + + except Exception as e: + logger.error(f"Configuration error: {e}") + import traceback + + traceback.print_exc() + return 1 + + # Run basic preprocessing + success = preprocess_basic( + desired_data=config["desired_data"], + data_dirs=config["data_dirs"], + cache_location=config["cache_location"], + rebuild_cache=config.get("rebuild_cache", False), + rebuild_maps=config.get("rebuild_maps", False), + incl_vector_map=config.get("incl_vector_map", True), + desired_dt=config.get("desired_dt", 0.1), + num_workers=config.get("num_workers", 8), + verbose=verbose, + list_scenes=args.list_scenes, + smooth_trajectories=config.get("smooth_trajectories", True), + ) + + return 0 if success else 1 + + +if __name__ == "__main__": + sys.exit(main()) From 1ded645f7e5b788568b18097788b1f314a3e0243 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:29:32 +0000 Subject: [PATCH 06/17] feat: integrate SceneDataSource into worker and simulation flow Update worker processes and simulation entry points to use the new SceneDataSource abstraction instead of direct Artifact access. Changes: - Worker IPC: PendingRolloutJob and AssignedRolloutJob use SceneDataSource - Worker main: Pass data_source from job to rollout execution - UnboundRollout: Accept SceneDataSource parameter - Simulate CLI: Use RuntimeContext.scene_id_to_idx for scene lookup This completes the migration from Artifact-based to SceneDataSource-based data loading, enabling support for multiple data source backends. --- .../alpasim_runtime/simulate/__main__.py | 10 ++++-- .../alpasim_runtime/unbound_rollout.py | 36 ++++++++++--------- src/runtime/alpasim_runtime/worker/ipc.py | 9 ++--- src/runtime/alpasim_runtime/worker/main.py | 25 +++---------- 4 files changed, 36 insertions(+), 44 deletions(-) diff --git a/src/runtime/alpasim_runtime/simulate/__main__.py b/src/runtime/alpasim_runtime/simulate/__main__.py index ac8649b5..1fc24e98 100644 --- a/src/runtime/alpasim_runtime/simulate/__main__.py +++ b/src/runtime/alpasim_runtime/simulate/__main__.py @@ -54,7 +54,7 @@ def create_arg_parser() -> argparse.ArgumentParser: # we split user and network config files because the latter is commonly auto-generated by kubernetes parser.add_argument("--user-config", type=str, required=True) parser.add_argument("--network-config", type=str, required=True) - parser.add_argument("--usdz-glob", type=str, required=True) + parser.add_argument( "--log-dir", type=str, @@ -95,7 +95,6 @@ async def _serve(args: argparse.Namespace) -> None: user_config=args.user_config, network_config=args.network_config, eval_config=args.eval_config, - usdz_glob=args.usdz_glob, log_dir=args.log_dir, validate_config_scenes=False, ) @@ -184,6 +183,12 @@ async def run_simulation(args: argparse.Namespace) -> bool: config = parse_simulator_config(args.user_config, args.network_config) eval_config = typed_parse_config(args.eval_config, EvalConfig) + # Require data_source in config (unified data flow) + if config.user.data_source is None: + raise ValueError( + "No data source specified. Please set 'data_source' in user config YAML. " + ) + # Derive output directories from log_dir rollouts_dir = os.path.join(args.log_dir, "rollouts") telemetry_dir = os.path.join(args.log_dir, "telemetry") @@ -265,7 +270,6 @@ async def _run_one_shot_request( user_config=args.user_config, network_config=args.network_config, eval_config=args.eval_config, - usdz_glob=args.usdz_glob, log_dir=args.log_dir, ) diff --git a/src/runtime/alpasim_runtime/unbound_rollout.py b/src/runtime/alpasim_runtime/unbound_rollout.py index 8166f7a9..cecf1e5a 100644 --- a/src/runtime/alpasim_runtime/unbound_rollout.py +++ b/src/runtime/alpasim_runtime/unbound_rollout.py @@ -21,9 +21,9 @@ VehicleConfig, ) from alpasim_runtime.services.sensorsim_service import ImageFormat -from alpasim_utils.artifact import Artifact from alpasim_utils.geometry import Pose, Trajectory from alpasim_utils.scenario import AABB, TrafficObjects +from alpasim_utils.scene_data_source import SceneDataSource from trajdata.maps import VectorMap logger = logging.getLogger(__name__) @@ -108,15 +108,19 @@ def create( simulation_config: SimulationConfig, scene_id: str, version_ids: RolloutMetadata.VersionIds, - available_artifacts: dict[str, Artifact], + data_source: SceneDataSource, rollouts_dir: str, ) -> UnboundRollout: - artifact = available_artifacts[scene_id] + """Create UnboundRollout from SceneDataSource.""" camera_configs = list(simulation_config.cameras) + # Get time range from data source rig + rig_time_range_start = data_source.rig.trajectory.time_range_us.start + rig_time_range_stop = data_source.rig.trajectory.time_range_us.stop + control_timestamps_us_arr: np.ndarray = ( - artifact.rig.trajectory.time_range_us.start + rig_time_range_start + simulation_config.time_start_offset_us + np.arange( simulation_config.n_sim_steps + 2 @@ -125,19 +129,17 @@ def create( ) control_timestamps_us = [ - int(min(t, artifact.rig.trajectory.time_range_us.stop - 1)) + int(min(t, rig_time_range_stop - 1)) for t in control_timestamps_us_arr - if t - < artifact.rig.trajectory.time_range_us.stop - + ORIGINAL_TRAJECTORY_DURATION_EXTENSION_US + if t < rig_time_range_stop + ORIGINAL_TRAJECTORY_DURATION_EXTENSION_US ] start_us = control_timestamps_us[0] end_us = control_timestamps_us[-1] - gt_ego_trajectory = artifact.rig.trajectory + gt_ego_trajectory = data_source.rig.trajectory # Filter out objects that are not in the time window - all_objs_in_window = artifact.traffic_objects.clip_trajectories( + all_objs_in_window = data_source.traffic_objects.clip_trajectories( start_us, end_us + 1, exclude_empty=True ) @@ -191,8 +193,8 @@ def create( if simulation_config.vehicle is not None: vehicle = simulation_config.vehicle - elif artifact.rig.vehicle_config is not None: - vehicle = artifact.rig.vehicle_config + elif data_source.rig.vehicle_config is not None: + vehicle = data_source.rig.vehicle_config else: raise ValueError("No vehicle config provided/found.") @@ -208,7 +210,7 @@ def create( gt_ego_trajectory=gt_ego_trajectory, traffic_objs=traffic_objects, n_sim_steps=simulation_config.n_sim_steps, - start_timestamp_us=artifact.rig.trajectory.time_range_us.start, + start_timestamp_us=data_source.rig.trajectory.time_range_us.start, force_gt_duration_us=simulation_config.force_gt_duration_us, control_timestep_us=simulation_config.control_timestep_us, follow_log=None, @@ -228,15 +230,15 @@ def create( vehicle ), ego_aabb=ego_aabb, - nre_runid=str(artifact.metadata.logger.run_id), - nre_version=artifact.metadata.version_string, - nre_uuid=str(artifact.metadata.uuid), + nre_runid=str(data_source.metadata.logger.run_id), + nre_version=data_source.metadata.version_string, + nre_uuid=str(data_source.metadata.uuid), planner_delay_us=simulation_config.planner_delay_us, pose_reporting_interval_us=simulation_config.pose_reporting_interval_us, route_generator_type=simulation_config.route_generator_type, send_recording_ground_truth=simulation_config.send_recording_ground_truth, vehicle_config=vehicle, - vector_map=artifact.map, + vector_map=data_source.map, hidden_traffic_objs=hidden_traffic_objs, group_render_requests=simulation_config.group_render_requests, ) diff --git a/src/runtime/alpasim_runtime/worker/ipc.py b/src/runtime/alpasim_runtime/worker/ipc.py index 9568f937..0e937fe2 100644 --- a/src/runtime/alpasim_runtime/worker/ipc.py +++ b/src/runtime/alpasim_runtime/worker/ipc.py @@ -14,6 +14,7 @@ from alpasim_grpc.v0.logging_pb2 import RolloutMetadata from alpasim_runtime.address_pool import ServiceAddress from alpasim_runtime.telemetry.rpc_wrapper import SharedRpcTracking +from alpasim_utils.scene_data_source import SceneDataSource from eval.scenario_evaluator import ScenarioEvalResult from eval.schema import EvalConfig @@ -42,8 +43,8 @@ class PendingRolloutJob: scene_id: str # Index of rollout spec in SimulationRequest.rollout_specs rollout_spec_index: int - # Artifact source path for this job's scene. - artifact_path: str + # SceneDataSource for this job's scene + data_source: SceneDataSource @dataclass @@ -58,8 +59,8 @@ class AssignedRolloutJob: scene_id: str # Index of rollout spec in SimulationRequest.rollout_specs rollout_spec_index: int - # Artifact source path for this job's scene. - artifact_path: str + # SceneDataSource for this job's scene + data_source: SceneDataSource # Concrete service addresses assigned by the parent dispatch loop. endpoints: ServiceEndpoints diff --git a/src/runtime/alpasim_runtime/worker/main.py b/src/runtime/alpasim_runtime/worker/main.py index cd1b4b02..f527d69a 100644 --- a/src/runtime/alpasim_runtime/worker/main.py +++ b/src/runtime/alpasim_runtime/worker/main.py @@ -39,14 +39,13 @@ from alpasim_runtime.telemetry.rpc_wrapper import set_shared_rpc_tracking from alpasim_runtime.telemetry.telemetry_context import TelemetryContext from alpasim_runtime.unbound_rollout import UnboundRollout -from alpasim_runtime.worker.artifact_cache import make_artifact_loader from alpasim_runtime.worker.ipc import ( AssignedRolloutJob, JobResult, WorkerArgs, _ShutdownSentinel, ) -from alpasim_utils.artifact import Artifact +from alpasim_utils.scene_data_source import SceneDataSource from eval.schema import EvalConfig @@ -61,7 +60,7 @@ def _is_orphaned(parent_pid: int) -> bool: async def run_single_rollout( job: AssignedRolloutJob, user_config: UserSimulatorConfig, - artifacts: dict[str, Artifact], + data_source: SceneDataSource, camera_catalog: CameraCatalog, version_ids: RolloutMetadata.VersionIds, rollouts_dir: str, @@ -105,7 +104,7 @@ async def run_single_rollout( simulation_config=user_config.simulation_config, scene_id=job.scene_id, version_ids=version_ids, - available_artifacts=artifacts, + data_source=data_source, rollouts_dir=rollouts_dir, ), ) @@ -161,8 +160,6 @@ async def run_worker_loop( result_queue: Queue, num_consumers: int, user_config: UserSimulatorConfig, - smooth_trajectories: bool, - artifact_cache_size: int | None, camera_catalog: CameraCatalog, version_ids: RolloutMetadata.VersionIds, rollouts_dir: str, @@ -178,9 +175,6 @@ async def run_worker_loop( result_queue: Queue to push JobResult to. num_consumers: Number of concurrent consumer tasks. user_config: User simulator configuration. - smooth_trajectories: Whether to smooth trajectories when loading artifacts. - artifact_cache_size: Max worker-local artifact cache size. - None = unlimited cache, 0 = disable cache. camera_catalog: Camera catalog for sensorsim. version_ids: Canonical version IDs from the parent process. rollouts_dir: Directory for rollout outputs. @@ -205,11 +199,6 @@ async def run_worker_loop( # Install event loop idle profiler install_event_loop_idle_profiler(loop) - load_artifact = make_artifact_loader( - smooth_trajectories=smooth_trajectories, - max_cache_size=artifact_cache_size, - ) - # Create a process pool for offloading CPU-bound eval computation. # One slot per consumer so no consumer blocks waiting for a pool slot. eval_executor = ProcessPoolExecutor(max_workers=num_consumers) @@ -251,13 +240,11 @@ def _poll_job() -> AssignedRolloutJob | _ShutdownSentinel | None: shutdown_event.set() break - artifact = load_artifact(job.scene_id, job.artifact_path) - - # Process the job + # Process the job using data_source from the job result = await run_single_rollout( job=job, user_config=user_config, - artifacts={job.scene_id: artifact}, + data_source=job.data_source, camera_catalog=camera_catalog, version_ids=version_ids, rollouts_dir=rollouts_dir, @@ -352,8 +339,6 @@ async def worker_async_main(args: WorkerArgs) -> None: result_queue=args.result_queue, num_consumers=args.num_consumers, user_config=user_config, - smooth_trajectories=user_config.smooth_trajectories, - artifact_cache_size=user_config.artifact_cache_size, camera_catalog=camera_catalog, version_ids=args.version_ids, rollouts_dir=rollouts_dir, From aade78bd7393b676a5d98aa3b437754ae39821e5 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 13:54:07 +0000 Subject: [PATCH 07/17] refactor: improve coordinate frame variable naming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow CONTRIBUTING.md naming conventions for coordinate frames. Changes in trajdata_data_source.py: - poses_vec3 → positions_agent_world - poses_quat → quaternions_agent_world - first_pose_position → position_ego_first_world - first_pose_local → position_ego_first_local These changes improve code readability by making coordinate frames explicit in variable names, following the position_{what}_{frame} naming pattern required by CONTRIBUTING.md. --- .../alpasim_utils/trajdata_data_source.py | 56 ++++++++++--------- 1 file changed, 29 insertions(+), 27 deletions(-) diff --git a/src/utils/alpasim_utils/trajdata_data_source.py b/src/utils/alpasim_utils/trajdata_data_source.py index 4cd4146f..0fa88f2a 100644 --- a/src/utils/alpasim_utils/trajdata_data_source.py +++ b/src/utils/alpasim_utils/trajdata_data_source.py @@ -293,8 +293,8 @@ def _extract_agent_trajectory( try: timestamps_us = [] - poses_vec3 = [] - poses_quat = [] + positions_agent_world = [] + quaternions_agent_world = [] # Iterate through all timesteps for ts in range(agent.first_timestep, agent.last_timestep + 1): @@ -336,11 +336,11 @@ def _extract_agent_trajectory( timestamp_us = int(base_timestamp_us + ts * dt * 1e6) timestamps_us.append(timestamp_us) - poses_vec3.append([x_val, y_val, z_val]) + positions_agent_world.append([x_val, y_val, z_val]) # Convert heading to quaternion quat = R.from_euler("z", heading_val).as_quat() # [x, y, z, w] - poses_quat.append(quat) + quaternions_agent_world.append(quat) except Exception as e: logger.debug( @@ -354,8 +354,8 @@ def _extract_agent_trajectory( # Create Trajectory trajectory = Trajectory( timestamps=np.array(timestamps_us, dtype=np.uint64), - positions=np.array(poses_vec3, dtype=np.float32), - quaternions=np.array(poses_quat, dtype=np.float32), + positions=np.array(positions_agent_world, dtype=np.float32), + quaternions=np.array(quaternions_agent_world, dtype=np.float32), ) # Create VehicleConfig (extract from extent) @@ -411,10 +411,10 @@ def rig(self) -> Rig: # Calculate world_to_nre transformation matrix (use first trajectory point as origin) world_to_nre = np.eye(4) if len(ego_trajectory) > 0: - first_pose_position = ego_trajectory.positions[0] - world_to_nre[:3, 3] = -first_pose_position + position_ego_first_world = ego_trajectory.positions[0] + world_to_nre[:3, 3] = -position_ego_first_world logger.info( - f"Setting world_to_nre origin at first pose: {first_pose_position}, " + f"Setting world_to_nre origin at first pose: {position_ego_first_world}, " f"translation: {world_to_nre[:3, 3]}" ) @@ -424,11 +424,11 @@ def rig(self) -> Rig: local_positions = ego_trajectory.positions + translation # Validate transform - first_pose_local = local_positions[0] - if np.linalg.norm(first_pose_local[:2]) > 1.0: + position_ego_first_local = local_positions[0] + if np.linalg.norm(position_ego_first_local[:2]) > 1.0: logger.warning( - f"First pose after transformation is not at origin: {first_pose_local}. " - f"Expected [0, 0, ~z], got {first_pose_local}" + f"First pose after transformation is not at origin: {position_ego_first_local}. " + f"Expected [0, 0, ~z], got {position_ego_first_local}" ) local_quat = ego_trajectory.quaternions.copy() @@ -491,26 +491,28 @@ def _extract_camera_info_from_scene(self) -> tuple[list[CameraId], dict]: try: unique_camera_id = f"{camera_name}@{self.scene_id}" - position = calibration_info.get( + position_sensor_to_ego = calibration_info.get( "sensor2ego_translation", [0.0, 0.0, 0.0] ) - rotation = calibration_info.get( + rotation_sensor_to_ego = calibration_info.get( "sensor2ego_rotation", [0.0, 0.0, 0.0, 1.0] ) - if isinstance(position, (int, float)): - position = [float(position), 0.0, 0.0] - elif len(position) < 3: - position = list(position) + [0.0] * (3 - len(position)) - - if isinstance(rotation, (int, float)): - rotation = [0.0, 0.0, 0.0, 1.0] - elif len(rotation) < 4: - if len(rotation) == 3: - r = R.from_euler("xyz", rotation) - rotation = r.as_quat() + if isinstance(position_sensor_to_ego, (int, float)): + position_sensor_to_ego = [float(position_sensor_to_ego), 0.0, 0.0] + elif len(position_sensor_to_ego) < 3: + position_sensor_to_ego = list(position_sensor_to_ego) + [0.0] * ( + 3 - len(position_sensor_to_ego) + ) + + if isinstance(rotation_sensor_to_ego, (int, float)): + rotation_sensor_to_ego = [0.0, 0.0, 0.0, 1.0] + elif len(rotation_sensor_to_ego) < 4: + if len(rotation_sensor_to_ego) == 3: + r = R.from_euler("xyz", rotation_sensor_to_ego) + rotation_sensor_to_ego = r.as_quat() else: - rotation = [0.0, 0.0, 0.0, 1.0] + rotation_sensor_to_ego = [0.0, 0.0, 0.0, 1.0] camera_id = CameraId( logical_name=camera_name, From e487b147efdb373dc66ae53c59619f8f426062a1 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 14:14:30 +0000 Subject: [PATCH 08/17] test: update tests for trajdata data source migration Update all runtime tests to work with the new trajdata-based data flow: - test_daemon_engine: Replace scene_id_to_artifact_path with scene_id_to_idx and mock scene_id_to_data_source cache in engine tests - test_daemon_main: Remove usdz_glob parameter from DaemonEngine construction - test_daemon_request_plumbing: Update build_pending_jobs_from_request tests to use get_data_source callable instead of scene_id_to_artifact_path dict - test_config: Add unit tests for new DataSourceConfig class - test_trajdata_integration: New integration tests for trajdata data source functionality, including caching, error handling, and job creation - test_runtime_integration_replay: Add TODO note for future update (this manual integration test needs significant rework for new data flow) All tests now reflect the shift from USDZ artifacts to lazy-loaded SceneDataSource instances via trajdata UnifiedDataset. Co-Authored-By: Claude Sonnet 4.5 --- src/runtime/tests/test_config.py | 38 ++++ src/runtime/tests/test_daemon_engine.py | 18 +- src/runtime/tests/test_daemon_main.py | 4 - .../tests/test_daemon_request_plumbing.py | 80 ++++--- .../tests/test_runtime_integration_replay.py | 5 + .../tests/test_trajdata_integration.py | 208 ++++++++++++++++++ 6 files changed, 308 insertions(+), 45 deletions(-) create mode 100644 src/runtime/tests/test_trajdata_integration.py diff --git a/src/runtime/tests/test_config.py b/src/runtime/tests/test_config.py index 9b0d321c..33f27c54 100644 --- a/src/runtime/tests/test_config.py +++ b/src/runtime/tests/test_config.py @@ -50,3 +50,41 @@ def test_typed_parse_config_invalid_yaml(tmp_path): # TODO(mwatson, mtyszkiewicz): What should happen when the config is empty? Currently, # no error is raised, and we return an empty config object. Is this the desired behavior? + + +def test_data_source_config_defaults(): + """Test that DataSourceConfig has sensible defaults.""" + cfg = config.DataSourceConfig( + desired_data=["nuplan_mini"], + data_dirs={"nuplan_mini": "/data/nuplan"}, + cache_location="/tmp/cache", + ) + assert cfg.desired_data == ["nuplan_mini"] + assert cfg.data_dirs == {"nuplan_mini": "/data/nuplan"} + assert cfg.cache_location == "/tmp/cache" + assert cfg.incl_vector_map is True + assert cfg.rebuild_cache is False + assert cfg.rebuild_maps is False + assert cfg.desired_dt == 0.1 + assert cfg.num_workers == 1 + assert cfg.num_timesteps_before == 30 + assert cfg.num_timesteps_after == 80 + assert cfg.config_dir is None + assert cfg.asset_base_path is None + + +def test_data_source_config_optional_fields(): + """Test that DataSourceConfig optional fields work correctly.""" + cfg = config.DataSourceConfig( + desired_data=["usdz"], + data_dirs={"usdz": "/data/usdz"}, + cache_location="/tmp/cache", + config_dir="/configs", + asset_base_path="/assets", + rebuild_cache=True, + desired_dt=0.05, + ) + assert cfg.config_dir == "/configs" + assert cfg.asset_base_path == "/assets" + assert cfg.rebuild_cache is True + assert cfg.desired_dt == 0.05 diff --git a/src/runtime/tests/test_daemon_engine.py b/src/runtime/tests/test_daemon_engine.py index 13c1522a..7ffb047d 100644 --- a/src/runtime/tests/test_daemon_engine.py +++ b/src/runtime/tests/test_daemon_engine.py @@ -65,7 +65,7 @@ async def _fake_build_runtime_context(*args, **kwargs): config=config, eval_config=eval_config, version_ids=version_ids, - scene_id_to_artifact_path={"clipgt-a": "/tmp/scene-a.usdz"}, + scene_id_to_idx={"clipgt-a": 0}, pools={"driver": MagicMock()}, max_in_flight=1, ) @@ -87,7 +87,6 @@ async def _fake_build_runtime_context(*args, **kwargs): user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", ) @@ -130,7 +129,7 @@ async def _fake_build_runtime_context(*args, **kwargs): config=config, eval_config=eval_config, version_ids=version_ids, - scene_id_to_artifact_path={"clipgt-a": "/tmp/scene-a.usdz"}, + scene_id_to_idx={"clipgt-a": 0}, pools={"driver": MagicMock()}, max_in_flight=1, ) @@ -152,7 +151,6 @@ async def _fake_build_runtime_context(*args, **kwargs): user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", validate_config_scenes=False, ) @@ -184,11 +182,11 @@ async def wait_request(self, request_id: str): user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_artifact_path = {"clipgt-a": "/tmp/scene-a.usdz"} + engine._scene_id_to_idx = {"clipgt-a": 0} + engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), @@ -232,11 +230,11 @@ async def wait_request(self, request_id): user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_artifact_path = {"clipgt-a": "/tmp/scene-a.usdz"} + engine._scene_id_to_idx = {"clipgt-a": 0} + engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), @@ -285,11 +283,11 @@ async def wait_request(self, request_id): user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_artifact_path = {"clipgt-a": "/tmp/scene-a.usdz"} + engine._scene_id_to_idx = {"clipgt-a": 0} + engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), diff --git a/src/runtime/tests/test_daemon_main.py b/src/runtime/tests/test_daemon_main.py index e31827bd..8923b463 100644 --- a/src/runtime/tests/test_daemon_main.py +++ b/src/runtime/tests/test_daemon_main.py @@ -143,7 +143,6 @@ async def run(self) -> None: user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", listen_address="[::]:50051", ) @@ -194,7 +193,6 @@ async def run(self) -> None: user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", listen_address="[::]:50051", ) @@ -418,7 +416,6 @@ def _make_one_shot_args() -> Namespace: user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", array_job_dir=None, ) @@ -532,7 +529,6 @@ async def test_run_simulation_one_shot_uses_daemon_engine( user_config="u.yaml", network_config="n.yaml", eval_config="e.yaml", - usdz_glob="/tmp/*.usdz", log_dir="/tmp/log", ) fake_engine.startup.assert_awaited_once() diff --git a/src/runtime/tests/test_daemon_request_plumbing.py b/src/runtime/tests/test_daemon_request_plumbing.py index 88600196..78ecba01 100644 --- a/src/runtime/tests/test_daemon_request_plumbing.py +++ b/src/runtime/tests/test_daemon_request_plumbing.py @@ -3,6 +3,8 @@ from __future__ import annotations +from unittest.mock import MagicMock + import pytest from alpasim_grpc.v0 import runtime_pb2 from alpasim_runtime.daemon.engine import ( @@ -16,13 +18,17 @@ def test_adapter_expands_nr_rollouts() -> None: rollout_specs=[runtime_pb2.RolloutSpec(scenario_id="clipgt-a", nr_rollouts=3)] ) - jobs = build_pending_jobs_from_request( - req, - scene_id_to_artifact_path={"clipgt-a": "/tmp/clipgt-a.usdz"}, - ) + mock_data_source = MagicMock() + + def fake_get_data_source(scene_id: str): + if scene_id == "clipgt-a": + return mock_data_source + raise UnknownSceneError(scene_id) + + jobs = build_pending_jobs_from_request(req, fake_get_data_source) assert [job.scene_id for job in jobs] == ["clipgt-a", "clipgt-a", "clipgt-a"] assert [job.rollout_spec_index for job in jobs] == [0, 0, 0] - assert all(job.artifact_path == "/tmp/clipgt-a.usdz" for job in jobs) + assert all(job.data_source is mock_data_source for job in jobs) def test_adapter_drops_zero_nr_rollouts_with_warning( @@ -33,10 +39,12 @@ def test_adapter_drops_zero_nr_rollouts_with_warning( rollout_specs=[runtime_pb2.RolloutSpec(scenario_id="clipgt-a")] ) - jobs = build_pending_jobs_from_request( - req, - scene_id_to_artifact_path={"clipgt-a": "/tmp/clipgt-a.usdz"}, - ) + mock_data_source = MagicMock() + + def fake_get_data_source(scene_id: str): + return mock_data_source + + jobs = build_pending_jobs_from_request(req, fake_get_data_source) assert jobs == [] assert "Dropping rollout spec with nr_rollouts=0" in caplog.text @@ -48,11 +56,13 @@ def test_adapter_rejects_scene_without_artifact() -> None: ] ) + def fake_get_data_source(scene_id: str): + if scene_id == "clipgt-missing": + raise UnknownSceneError(scene_id) + return MagicMock() + with pytest.raises(UnknownSceneError): - build_pending_jobs_from_request( - req, - scene_id_to_artifact_path={"clipgt-a": "/tmp/clipgt-a.usdz"}, - ) + build_pending_jobs_from_request(req, fake_get_data_source) def test_adapter_assigns_rollout_spec_indexes_in_request_order() -> None: @@ -63,20 +73,24 @@ def test_adapter_assigns_rollout_spec_indexes_in_request_order() -> None: ] ) - jobs = build_pending_jobs_from_request( - req, - scene_id_to_artifact_path={ - "clipgt-a": "/tmp/clipgt-a.usdz", - "clipgt-b": "/tmp/clipgt-b.usdz", - }, - ) + mock_data_source_a = MagicMock() + mock_data_source_b = MagicMock() + + def fake_get_data_source(scene_id: str): + if scene_id == "clipgt-a": + return mock_data_source_a + elif scene_id == "clipgt-b": + return mock_data_source_b + raise UnknownSceneError(scene_id) + + jobs = build_pending_jobs_from_request(req, fake_get_data_source) assert len(jobs) == 3 assert [job.scene_id for job in jobs] == ["clipgt-a", "clipgt-b", "clipgt-b"] assert [job.rollout_spec_index for job in jobs] == [0, 1, 1] - assert [job.artifact_path for job in jobs] == [ - "/tmp/clipgt-a.usdz", - "/tmp/clipgt-b.usdz", - "/tmp/clipgt-b.usdz", + assert [job.data_source for job in jobs] == [ + mock_data_source_a, + mock_data_source_b, + mock_data_source_b, ] @@ -88,13 +102,17 @@ def test_adapter_ignores_zero_rollout_specs_when_indexing() -> None: ] ) - jobs = build_pending_jobs_from_request( - req, - scene_id_to_artifact_path={ - "clipgt-a": "/tmp/clipgt-a.usdz", - "clipgt-b": "/tmp/clipgt-b.usdz", - }, - ) + mock_data_source_a = MagicMock() + mock_data_source_b = MagicMock() + + def fake_get_data_source(scene_id: str): + if scene_id == "clipgt-a": + return mock_data_source_a + elif scene_id == "clipgt-b": + return mock_data_source_b + raise UnknownSceneError(scene_id) + + jobs = build_pending_jobs_from_request(req, fake_get_data_source) assert len(jobs) == 2 assert [job.scene_id for job in jobs] == ["clipgt-b", "clipgt-b"] assert [job.rollout_spec_index for job in jobs] == [1, 1] diff --git a/src/runtime/tests/test_runtime_integration_replay.py b/src/runtime/tests/test_runtime_integration_replay.py index c55b50c8..ad84f9e8 100644 --- a/src/runtime/tests/test_runtime_integration_replay.py +++ b/src/runtime/tests/test_runtime_integration_replay.py @@ -3,6 +3,11 @@ """Manual integration replay test for runtime request determinism. +TODO: This test needs to be updated for the new trajdata-based data flow. +The test currently references usdz_glob which has been removed in favor of +DataSourceConfig. This test is marked as @pytest.mark.manual and is not +run in regular CI, so it can be updated as part of a follow-up task. + This test boots one replay gRPC server per service and re-runs runtime against recorded artifacts. It asserts that runtime emits requests matching the ASL recording exactly (ignoring expected dynamic fields). diff --git a/src/runtime/tests/test_trajdata_integration.py b/src/runtime/tests/test_trajdata_integration.py new file mode 100644 index 00000000..428a2b44 --- /dev/null +++ b/src/runtime/tests/test_trajdata_integration.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA Corporation + +"""Integration tests for trajdata data source functionality.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from alpasim_grpc.v0 import runtime_pb2 +from alpasim_runtime.daemon.engine import ( + DaemonEngine, + UnknownSceneError, + build_pending_jobs_from_request, +) + + +@pytest.fixture +def mock_trajdata_scene(): + """Create a mock trajdata Scene object.""" + scene = MagicMock() + scene.name = "test_scene_001" + scene.length_timesteps = 100 + scene.dt = 0.1 + return scene + + +@pytest.fixture +def mock_trajdata_dataset(mock_trajdata_scene): + """Create a mock UnifiedDataset.""" + dataset = MagicMock() + dataset.num_scenes.return_value = 1 + dataset.get_scene.return_value = mock_trajdata_scene + dataset.cache_class = MagicMock + dataset.cache_path = "/tmp/cache" + dataset.augmentations = None + dataset.obs_format = MagicMock() + return dataset + + +def test_build_pending_jobs_with_valid_scene(): + """Test building pending jobs from a simulation request with valid scene IDs.""" + request = runtime_pb2.SimulationRequest( + rollout_specs=[ + runtime_pb2.RolloutSpec(scenario_id="scene_a", nr_rollouts=2), + runtime_pb2.RolloutSpec(scenario_id="scene_b", nr_rollouts=1), + ] + ) + + mock_data_source_a = MagicMock() + mock_data_source_b = MagicMock() + + def fake_get_data_source(scene_id: str): + if scene_id == "scene_a": + return mock_data_source_a + elif scene_id == "scene_b": + return mock_data_source_b + raise UnknownSceneError(scene_id) + + jobs = build_pending_jobs_from_request(request, fake_get_data_source) + + # Should create 3 jobs total: 2 for scene_a, 1 for scene_b + assert len(jobs) == 3 + + # Check first two jobs are for scene_a + assert jobs[0].scene_id == "scene_a" + assert jobs[0].rollout_spec_index == 0 + assert jobs[0].data_source is mock_data_source_a + assert jobs[1].scene_id == "scene_a" + assert jobs[1].rollout_spec_index == 0 + assert jobs[1].data_source is mock_data_source_a + + # Check third job is for scene_b + assert jobs[2].scene_id == "scene_b" + assert jobs[2].rollout_spec_index == 1 + assert jobs[2].data_source is mock_data_source_b + + +def test_build_pending_jobs_with_unknown_scene(): + """Test that UnknownSceneError is raised for unknown scene IDs.""" + request = runtime_pb2.SimulationRequest( + rollout_specs=[ + runtime_pb2.RolloutSpec(scenario_id="unknown_scene", nr_rollouts=1), + ] + ) + + def fake_get_data_source(scene_id: str): + raise UnknownSceneError(scene_id) + + with pytest.raises(UnknownSceneError) as exc_info: + build_pending_jobs_from_request(request, fake_get_data_source) + + assert exc_info.value.scene_id == "unknown_scene" + + +def test_build_pending_jobs_drops_zero_rollouts(): + """Test that specs with nr_rollouts=0 are dropped with a warning.""" + request = runtime_pb2.SimulationRequest( + rollout_specs=[ + runtime_pb2.RolloutSpec(scenario_id="scene_a", nr_rollouts=1), + runtime_pb2.RolloutSpec(scenario_id="scene_b", nr_rollouts=0), + ] + ) + + mock_data_source = MagicMock() + + def fake_get_data_source(scene_id: str): + return mock_data_source + + jobs = build_pending_jobs_from_request(request, fake_get_data_source) + + # Should only create 1 job (scene_b with 0 rollouts is dropped) + assert len(jobs) == 1 + assert jobs[0].scene_id == "scene_a" + + +def test_daemon_engine_get_data_source_caching( + mock_trajdata_dataset, mock_trajdata_scene +): + """Test that DaemonEngine caches data sources and doesn't reload them.""" + with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + # Set up engine state + engine._started = True + engine._dataset = mock_trajdata_dataset + engine._scene_id_to_idx = {"test_scene_001": 0} + engine._scene_id_to_data_source = {} + engine._config = SimpleNamespace( + user=SimpleNamespace( + smooth_trajectories=True, + data_source=SimpleNamespace(asset_base_path="/tmp/assets"), + ) + ) + + # Mock TrajdataDataSource.from_trajdata_scene + with patch( + "alpasim_runtime.daemon.engine.TrajdataDataSource.from_trajdata_scene" + ) as mock_from_scene: + mock_data_source = MagicMock() + mock_from_scene.return_value = mock_data_source + + # First call should create the data source + data_source_1 = engine._get_data_source("test_scene_001") + assert data_source_1 is mock_data_source + assert mock_from_scene.call_count == 1 + + # Second call should return cached data source + data_source_2 = engine._get_data_source("test_scene_001") + assert data_source_2 is mock_data_source + assert mock_from_scene.call_count == 1 # Not called again + + # Verify cache contains the data source + assert "test_scene_001" in engine._scene_id_to_data_source + assert engine._scene_id_to_data_source["test_scene_001"] is mock_data_source + + +def test_daemon_engine_get_data_source_unknown_scene(): + """Test that _get_data_source raises UnknownSceneError for unknown scenes.""" + with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + engine._started = True + engine._dataset = MagicMock() + engine._scene_id_to_idx = {"known_scene": 0} + engine._scene_id_to_data_source = {} + engine._config = SimpleNamespace( + user=SimpleNamespace( + smooth_trajectories=True, + data_source=SimpleNamespace(asset_base_path=None), + ) + ) + + with pytest.raises(UnknownSceneError) as exc_info: + engine._get_data_source("unknown_scene") + + assert exc_info.value.scene_id == "unknown_scene" + + +def test_daemon_engine_get_data_source_without_dataset(): + """Test that _get_data_source raises RuntimeError if dataset is not initialized.""" + with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + engine._started = True + engine._dataset = None + engine._scene_id_to_idx = {"test_scene": 0} + engine._scene_id_to_data_source = {} + + with pytest.raises(RuntimeError, match="Dataset not initialized"): + engine._get_data_source("test_scene") From 486d87ddfef59030db07bef07da2ec12db6b8fa5 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 19 Mar 2026 14:52:56 +0000 Subject: [PATCH 09/17] docs(wizard): add trajdata documentation and wizard config Update documentation and wizard configuration to support the new trajdata-based data flow: Documentation updates: - README: Add Quick Start section explaining data preparation - ONBOARDING: Document trajdata dependency (alpasim branch) - TUTORIAL: Add Data Preparation section with prepare_data examples - TUTORIAL: Fix deprecated --usdz-glob references in debug examples Wizard configuration: - Add defines.trajdata_cache for unified cache location - Add runtime.data_source with USDZ defaults - Configure recursive scanning of all-usdzs directory - Set sensible defaults: 10Hz, vector maps, 4 workers This ensures users understand the data preparation workflow and wizard-generated configs work without manual edits. Co-Authored-By: Claude Sonnet 4.5 --- README.md | 34 ++++++++- docs/ONBOARDING.md | 23 +++++++ docs/TUTORIAL.md | 103 +++++++++++++++++++++++++++- src/wizard/configs/base_config.yaml | 53 ++++++++++++++ 4 files changed, 208 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 8e78eac1..9ea51b7c 100644 --- a/README.md +++ b/README.md @@ -57,8 +57,38 @@ appreciated. ## Getting Started -To run simulations locally (Docker Compose, single machine), see the [Tutorial](docs/TUTORIAL.md). -For cluster or SLURM deployment, see `src/tools/run-on-slurm`. +### Quick Start + +1. **Setup Environment** + Follow the [Onboarding Guide](docs/ONBOARDING.md) to install dependencies and configure your + environment. + +2. **Prepare Scene Data** + AlpaSim uses [trajdata](https://github.com/NVlabs/trajdata/tree/alpasim) (custom `alpasim` + branch) for unified data loading. This dependency is automatically installed via `uv`. The wizard + automatically prepares scene caches, but you can also do it manually: + + ```bash + # For USDZ scenes (after downloading from Hugging Face) + uv run python -m alpasim_runtime.prepare_data \ + --desired-data=usdz \ + --data-dir=./data/nre-artifacts/all-usdzs \ + --cache-location=./cache/trajdata_usdz + ``` + + See the [Tutorial](docs/TUTORIAL.md#data-preparation) for more details on data preparation + options. + +3. **Run Your First Simulation** + ```bash + source setup_local_env.sh + uv run alpasim_wizard +deploy=local wizard.log_dir=$PWD/my_first_run + ``` + + Results will be in `my_first_run/` including videos and metrics. + +For detailed instructions, see the [Tutorial](docs/TUTORIAL.md). For cluster or SLURM deployment, see +`src/tools/run-on-slurm`. ## Documentation & Resources diff --git a/docs/ONBOARDING.md b/docs/ONBOARDING.md index cd29a047..36a5e45c 100644 --- a/docs/ONBOARDING.md +++ b/docs/ONBOARDING.md @@ -62,6 +62,29 @@ uv run --project src/runtime python -c "..." # run in a sub-project context All members share one dependency resolution; there is no per-member version isolation. See [Plugin System](PLUGIN_SYSTEM.md) for how plugins integrate with the workspace. +## Key Dependencies + +### Trajdata (Unified Data Loading) + +AlpaSim uses [trajdata](https://github.com/NVlabs/trajdata) for unified scene data loading across +different autonomous driving datasets (USDZ, NuPlan, Waymo, etc.). The dependency is automatically +installed via `uv` from the `alpasim` branch: + +```toml +# src/runtime/pyproject.toml +[tool.uv.sources.trajdata-alpasim] +git = "https://github.com/NVlabs/trajdata" +branch = "alpasim" # Custom branch with AlpaSim-specific extensions +``` + +**You don't need to manually install trajdata** - it's handled automatically when you run +`setup_local_env.sh` or `uv sync --extra runtime`. + +The `alpasim` branch includes: +- USDZ data format support +- Custom trajectory smoothing +- Performance optimizations for simulation workloads + ## Troubleshooting **`CUDA_ERROR_UNSUPPORTED_PTX_VERSION` during NRE warmup** diff --git a/docs/TUTORIAL.md b/docs/TUTORIAL.md index e5efebbd..5718a1de 100644 --- a/docs/TUTORIAL.md +++ b/docs/TUTORIAL.md @@ -336,6 +336,90 @@ the predictions of a policy, you can set `runtime.simulation_config.physics_update_mode: NONE` and `runtime.simulation_config.force_gt_duration_us` to a very high value (20s+). +## Data Preparation + +AlpaSim uses [trajdata](https://github.com/NVlabs/trajdata/tree/alpasim) (custom `alpasim` branch) +for unified data loading across different autonomous driving datasets. The trajdata library is +automatically installed via `uv` when you run `setup_local_env.sh`. + +Before running simulations, you need to prepare a trajdata cache from your scene data. + +### Preparing USDZ Scene Cache + +The wizard automatically handles data preparation for downloaded USDZ scenes. However, if you need to +manually prepare or rebuild the cache, use the `prepare_data` tool: + +```bash +# Prepare cache from USDZ scenes +uv run python -m alpasim_runtime.prepare_data \ + --desired-data=usdz \ + --data-dir=./data/nre-artifacts/all-usdzs \ + --cache-location=./cache/trajdata_usdz \ + --smooth_trajectories=true \ + --log-level=INFO +``` + +This command: +- Scans USDZ files in the specified directory +- Extracts trajectory and map data +- Creates a trajdata cache for fast scene loading +- The cache is reused across runs to improve startup time + +### Using Configuration Files + +For more complex setups or repeated use, create a configuration file (e.g., +`user_config/config_prepare_usdz.yaml`): + +```yaml +data_source: + desired_data: ["usdz"] + data_dirs: + usdz: "./data/nre-artifacts/all-usdzs" + cache_location: "./cache/trajdata_usdz" + incl_vector_map: true + rebuild_cache: false # Set to true to force rebuild + rebuild_maps: false # Set to true to force rebuild + desired_dt: 0.1 # 10 Hz sampling rate + num_workers: 4 # Parallel workers for cache creation + +smooth_trajectories: true +``` + +Then run: + +```bash +uv run python -m alpasim_runtime.prepare_data \ + --user-config=user_config/config_prepare_usdz.yaml +``` + +### Cache Location + +The trajdata cache contains: +- Preprocessed scene metadata and indices +- Trajectory data in a unified format +- Vector map data (when enabled) +- Scene lookup tables for fast access + +You can share the cache directory across machines to avoid redundant preprocessing. + +### Rebuilding the Cache + +If you encounter data inconsistencies or add new scenes, rebuild the cache: + +```bash +# Command line +uv run python -m alpasim_runtime.prepare_data \ + --desired-data=usdz \ + --data-dir=./data/nre-artifacts/all-usdzs \ + --cache-location=./cache/trajdata_usdz \ + --rebuild-cache + +# Or in config file, set: rebuild_cache: true +``` + +> :green_book: The wizard uses configuration files from `user_config/` that include data source +> settings. These configs are automatically used during simulation runs. + ## Scenes The scene in AlpaSim is a NuRec reconstruction of a real-world driving log. @@ -510,7 +594,17 @@ from shutting down the docker containers after each simulation by setting 1. (Terminal 2) `cd` into the the runtime src directory (`/src/runtime/`) and prepare to start the runtime. The exact command paths will vary, but, to use the configuration generated from the earlier steps, an example command would be: - `bash cd /src/runtime/ # Following command is based on the docker-compose.yaml generated by the wizard uv run python -m alpasim_runtime.simulate \ --usdz-glob=../../data/nre-artifacts/all-usdzs/**/*.usdz \ --user-config=../../tutorial_dbg_runtime/generated-user-config-0.yaml \ --network-config=../../tutorial_dbg_runtime/generated-network-config.yaml \ --log-dir=../../tutorial_dbg_runtime \ --log-level=INFO ` + ```bash + cd /src/runtime/ + # Following command is based on the docker-compose.yaml generated by the wizard + # Ensure the user config contains the data_source configuration + uv run python -m alpasim_runtime.simulate \ + --user-config=../../tutorial_dbg_runtime/generated-user-config-0.yaml \ + --network-config=../../tutorial_dbg_runtime/generated-network-config.yaml \ + --log-dir=../../tutorial_dbg_runtime \ + --eval-config=../../tutorial_dbg_runtime/eval-config.yaml \ + --log-level=INFO + ``` ### Using VSCode Debugger (Optional) @@ -528,13 +622,16 @@ built-in debugger: "justMyCode": false, "cwd": "${workspaceFolder}/src/runtime", "args": [ - "--usdz-glob=../../data/nre-artifacts/all-usdzs/**/*.usdz", "--user-config=../../tutorial_dbg_runtime/generated-user-config-0.yaml", "--network-config=../../tutorial_dbg_runtime/generated-network-config.yaml", + "--eval-config=../../tutorial_dbg_runtime/eval-config.yaml", "--log-dir=../../tutorial_dbg_runtime", "--log-level=INFO" ], - "console": "integratedTerminal" + "console": "integratedTerminal", + "env": { + "PYTHONPATH": "${workspaceFolder}/src/grpc:${workspaceFolder}/src/eval/src:${workspaceFolder}/src/utils:${workspaceFolder}/src/runtime:${env:PYTHONPATH}" + } } ``` diff --git a/src/wizard/configs/base_config.yaml b/src/wizard/configs/base_config.yaml index cdde7a13..4b5c13ed 100644 --- a/src/wizard/configs/base_config.yaml +++ b/src/wizard/configs/base_config.yaml @@ -30,6 +30,7 @@ defines: drivers: "${defines.filesystem}/drivers" sensordata: "${defines.filesystem}/nre-artifacts" trafficsim_map_cache: "${defines.filesystem}/trafficsim/unified_data_cache" + trajdata_cache: "${defines.filesystem}/trajdata_cache" # Trajdata unified cache location sensorsim_entrypoint: "/app/internal/scripts/pycena/runtime/pycena_run" helper: scripts @@ -224,7 +225,59 @@ services: runtime: # nr_workers and endpoints.*.n_concurrent_rollouts are set by topology configs. + # Unified data source configuration using trajdata + # This replaces the old USDZ artifact glob pattern approach + data_source: + desired_data: + - usdz # Load USDZ scenes (NuRec artifacts) + + data_dirs: + # Point to all-usdzs directory to scan all scene batches recursively + # This covers paths like: all-usdzs/sample_set/25.07_release/Batch000X/ + usdz: "${scenes.scene_cache}/all-usdzs/sample_set/25.07_release/Batch0001" + + cache_location: "${defines.trajdata_cache}" # Trajdata unified cache + + # Map and cache settings + incl_vector_map: true # Include vector map data (roads, lanes, etc.) + rebuild_cache: false # Set to true to force rebuild cache + rebuild_maps: false # Set to true to force rebuild maps + + # Trajectory processing + desired_dt: 0.1 # 10 Hz sampling rate (matches control_timestep_us) + num_workers: 4 # Parallel workers for cache creation + + # Optional: Base path for MTGS rendering assets + asset_base_path: null + + # YAML config mode parameters (for NuPlan central token processing) + config_dir: null # Set to enable YAML batch preprocessing mode + num_timesteps_before: 30 # Timesteps before central token + num_timesteps_after: 80 # Timesteps after central token + + # Enable cubic spline smoothing for trajectories + smooth_trajectories: true + endpoints: + sensorsim: + # how many rollouts can run on a single worker at once + n_concurrent_rollouts: 2 + + driver: + n_concurrent_rollouts: 14 + + physics: + n_concurrent_rollouts: 14 + skip: false # physics supports skipping + + controller: + n_concurrent_rollouts: 14 + skip: false # controller supports skipping + + trafficsim: + n_concurrent_rollouts: 14 + skip: true # trafficsim is not included yet + # shut down the system after simulation is finished. without this flag the microservice servers # will remain on forever requring a manual interrupt (useful for debugging) do_shutdown: true From 848d89e8be9fec8e52969387504d53e5d399d3da Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 20 Mar 2026 04:46:58 +0000 Subject: [PATCH 10/17] fix: update copyright year to 2026 for new files Address reviewer feedback: new files added in this PR should use copyright year 2026 only, not 2025-2026. Changed files: - src/runtime/alpasim_runtime/prepare_data/__init__.py - src/runtime/alpasim_runtime/prepare_data/__main__.py - src/utils/alpasim_utils/scene_data_source.py - src/utils/alpasim_utils/trajdata_data_source.py --- src/runtime/alpasim_runtime/prepare_data/__init__.py | 2 +- src/runtime/alpasim_runtime/prepare_data/__main__.py | 2 +- src/utils/alpasim_utils/scene_data_source.py | 2 +- src/utils/alpasim_utils/trajdata_data_source.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/runtime/alpasim_runtime/prepare_data/__init__.py b/src/runtime/alpasim_runtime/prepare_data/__init__.py index cd4c763a..8b2ada83 100644 --- a/src/runtime/alpasim_runtime/prepare_data/__init__.py +++ b/src/runtime/alpasim_runtime/prepare_data/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025-2026 NVIDIA Corporation +# Copyright (c) 2026 NVIDIA Corporation """ Data preprocessing module for building trajdata cache. diff --git a/src/runtime/alpasim_runtime/prepare_data/__main__.py b/src/runtime/alpasim_runtime/prepare_data/__main__.py index 5200349f..b695efc3 100644 --- a/src/runtime/alpasim_runtime/prepare_data/__main__.py +++ b/src/runtime/alpasim_runtime/prepare_data/__main__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025-2026 NVIDIA Corporation +# Copyright (c) 2026 NVIDIA Corporation """ Data preprocessing CLI for building trajdata cache. diff --git a/src/utils/alpasim_utils/scene_data_source.py b/src/utils/alpasim_utils/scene_data_source.py index 0e4ca073..865754d6 100644 --- a/src/utils/alpasim_utils/scene_data_source.py +++ b/src/utils/alpasim_utils/scene_data_source.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025-2026 NVIDIA Corporation +# Copyright (c) 2026 NVIDIA Corporation """ SceneDataSource Protocol for abstracting scene data loading. diff --git a/src/utils/alpasim_utils/trajdata_data_source.py b/src/utils/alpasim_utils/trajdata_data_source.py index 0fa88f2a..5cc31216 100644 --- a/src/utils/alpasim_utils/trajdata_data_source.py +++ b/src/utils/alpasim_utils/trajdata_data_source.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025-2026 NVIDIA Corporation +# Copyright (c) 2026 NVIDIA Corporation """ TrajdataDataSource: Implementation for loading scene data directly from trajdata From eff45136e90149050e934ad6748793108a7a5aff Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 20 Mar 2026 05:06:16 +0000 Subject: [PATCH 11/17] refactor: improve configuration structure and validation Address reviewer comments 3, 5, 6, 7, 9 on PR #56. Changes: - Restructure DataSourceConfig to hierarchical format (Comment 3) - Separate common settings from source-specific configuration - Add USDZSourceConfig and NuPlanSourceConfig dataclasses - Add to_trajdata_params() conversion method - Makes it clear which parameters affect which data sources - Make data_source required via Hydra MISSING (Comment 7) - Change from Optional[...] = None to MISSING - Leverage Hydra's built-in validation instead of runtime checks - Provides better error messages and fails fast - Use dynamic sceneset_path instead of hardcoded path (Comment 9) - Preserves wizard's ability to create temporary sceneset directories - Add sceneset_path to scenes config with sensible default - Remove unused configuration (Comments 5, 6) - Remove --usdz-glob command line argument - Remove artifact_cache_size config field Co-Authored-By: Claude Sonnet 4.5 --- src/runtime/alpasim_runtime/config.py | 136 ++++++++++++++---- src/runtime/alpasim_runtime/daemon/engine.py | 38 +---- .../alpasim_runtime/runtime_context.py | 36 +---- .../alpasim_runtime/simulate/__main__.py | 6 - src/runtime/tests/test_config.py | 79 +++++++--- src/wizard/configs/base_config.yaml | 54 +++---- 6 files changed, 203 insertions(+), 146 deletions(-) diff --git a/src/runtime/alpasim_runtime/config.py b/src/runtime/alpasim_runtime/config.py index 47896f97..30740d94 100644 --- a/src/runtime/alpasim_runtime/config.py +++ b/src/runtime/alpasim_runtime/config.py @@ -18,39 +18,120 @@ @dataclass -class DataSourceConfig: - """Configuration for trajdata's UnifiedDataset. - - This config unifies all scene data loading through trajdata's UnifiedDataset, - supporting both USDZ files and standard trajdata datasets like NuPlan, Waymo, etc. +class USDZSourceConfig: + """Configuration for USDZ data source. Attributes: - desired_data: List of dataset names to load (e.g., ["nuplan_test", "usdz"]) - data_dirs: Dict mapping dataset names to their data directories - cache_location: Path to trajdata cache directory - config_dir: Optional directory containing YAML scene config files for batch mode + enabled: Whether this data source is enabled + data_dir: Path to directory containing USDZ files + desired_dt: Desired time delta between trajectory frames in seconds + incl_vector_map: Whether to load vector maps (roads, lanes, etc.) asset_base_path: Base path for rendering assets (e.g., MTGS assets) - incl_vector_map: Whether to load vector maps - rebuild_cache: Whether to force rebuild the trajdata cache - rebuild_maps: Whether to force rebuild maps - desired_dt: Desired time delta between frames in seconds - num_workers: Number of workers for data loading + """ + + enabled: bool = True + data_dir: str = MISSING + desired_dt: float = 0.1 # 10 Hz sampling + incl_vector_map: bool = True + asset_base_path: Optional[str] = None + + +@dataclass +class NuPlanSourceConfig: + """Configuration for NuPlan data source. + + Attributes: + enabled: Whether this data source is enabled + data_dir: Path to NuPlan dataset directory + config_dir: Directory containing YAML scene config files for batch mode num_timesteps_before: Number of timesteps before central token (batch mode) num_timesteps_after: Number of timesteps after central token (batch mode) + desired_dt: Desired time delta between frames in seconds + incl_vector_map: Whether to load vector maps """ - desired_data: list[str] = MISSING - data_dirs: dict[str, str] = MISSING - cache_location: str = MISSING - config_dir: Optional[str] = None # For YAML batch mode - asset_base_path: Optional[str] = None + enabled: bool = False + data_dir: Optional[str] = None + config_dir: Optional[str] = None + num_timesteps_before: int = 30 + num_timesteps_after: int = 80 + desired_dt: float = 0.1 incl_vector_map: bool = True + + +@dataclass +class DataSourceConfig: + """Configuration for unified data loading through trajdata. + + This provides a hierarchical structure where common configuration is separated + from data source-specific settings, making it easier to understand which + parameters affect which data sources. + + Attributes: + cache_location: Path to shared trajdata cache directory + rebuild_cache: Whether to force rebuild the cache for all sources + rebuild_maps: Whether to force rebuild maps for all sources + num_workers: Number of parallel workers for cache creation + usdz: USDZ-specific configuration + nuplan: NuPlan-specific configuration + """ + + # Common configuration (applies to all data sources) + cache_location: str = MISSING rebuild_cache: bool = False rebuild_maps: bool = False - desired_dt: float = 0.1 # 10 Hz default - num_workers: int = 1 - num_timesteps_before: int = 30 # For batch mode - num_timesteps_after: int = 80 # For batch mode + num_workers: int = 4 + + # Source-specific configurations + usdz: Optional[USDZSourceConfig] = None + nuplan: Optional[NuPlanSourceConfig] = None + + def to_trajdata_params(self) -> dict: + """Convert hierarchical config to flat parameters for trajdata's UnifiedDataset. + + Returns: + Dictionary with keys expected by UnifiedDataset constructor + + Raises: + ValueError: If no data sources are enabled + """ + desired_data = [] + data_dirs = {} + + # Collect enabled sources + if self.usdz is not None and self.usdz.enabled: + desired_data.append("usdz") + data_dirs["usdz"] = self.usdz.data_dir + + if self.nuplan is not None and self.nuplan.enabled: + desired_data.append("nuplan") + data_dirs["nuplan"] = self.nuplan.data_dir + + if not desired_data: + raise ValueError("No data sources enabled in configuration") + + # Use first enabled source for common parameters + # (desired_dt, incl_vector_map are typically consistent across sources) + primary_source = self.usdz if (self.usdz and self.usdz.enabled) else self.nuplan + + params = { + "desired_data": desired_data, + "data_dirs": data_dirs, + "cache_location": self.cache_location, + "rebuild_cache": self.rebuild_cache, + "rebuild_maps": self.rebuild_maps, + "num_workers": self.num_workers, + "desired_dt": primary_source.desired_dt, + "incl_vector_map": primary_source.incl_vector_map, + } + + # Add source-specific parameters + if self.nuplan and self.nuplan.enabled and self.nuplan.config_dir: + params["config_dir"] = self.nuplan.config_dir + params["num_timesteps_before"] = self.nuplan.num_timesteps_before + params["num_timesteps_after"] = self.nuplan.num_timesteps_after + + return params def typed_parse_config(path: str | Path, config_type: Type[C]) -> C: @@ -279,9 +360,6 @@ class UserSimulatorConfig: endpoints: UserEndpointConfig = MISSING smooth_trajectories: bool = True # whether to smooth trajectories with cubic spline - # Max worker-local artifact cache size. - # None = unlimited, 0 = disable cache and always reload artifacts. - artifact_cache_size: Optional[int] = None extra_cameras: list[CameraDefinitionConfig] = field(default_factory=list) # Number of worker processes for parallel rollout execution. @@ -289,9 +367,9 @@ class UserSimulatorConfig: # >1 = multi-worker mode with subprocess-based parallelism nr_workers: int = MISSING - # Unified data source configuration (optional for backward compatibility) - # When provided, data loading goes through trajdata's UnifiedDataset - data_source: Optional[DataSourceConfig] = None + # Unified data source configuration (required) + # Data loading goes through trajdata's UnifiedDataset + data_source: DataSourceConfig = MISSING @dataclass diff --git a/src/runtime/alpasim_runtime/daemon/engine.py b/src/runtime/alpasim_runtime/daemon/engine.py index d41ac1d4..362534e5 100644 --- a/src/runtime/alpasim_runtime/daemon/engine.py +++ b/src/runtime/alpasim_runtime/daemon/engine.py @@ -8,14 +8,6 @@ from typing import Callable from uuid import uuid4 -try: - from trajdata.dataset import UnifiedDataset - - TRAJDATA_AVAILABLE = True -except ImportError: - TRAJDATA_AVAILABLE = False - UnifiedDataset = None - from alpasim_grpc.v0 import logging_pb2, runtime_pb2 from alpasim_runtime.address_pool import AddressPool from alpasim_runtime.daemon.scheduler import DaemonScheduler, DaemonUnavailableError @@ -27,6 +19,7 @@ from alpasim_runtime.worker.runtime import WorkerRuntime, start_worker_runtime from alpasim_utils.scene_data_source import SceneDataSource from alpasim_utils.trajdata_data_source import TrajdataDataSource +from trajdata.dataset import UnifiedDataset from eval.data import AggregationType @@ -212,7 +205,7 @@ def _get_data_source(self, scene_id: str) -> SceneDataSource: return self._scene_id_to_data_source[scene_id] # Lazy load from dataset - if self._dataset is None or not TRAJDATA_AVAILABLE: + if self._dataset is None: raise RuntimeError(f"Dataset not initialized, cannot load scene {scene_id}") if self._config is None: @@ -227,10 +220,10 @@ def _get_data_source(self, scene_id: str) -> SceneDataSource: if scene is None: raise UnknownSceneError(scene_id) - # Get asset_base_path from config + # Get asset_base_path from USDZ config asset_base_path = None - if self._config.user.data_source is not None: - asset_base_path = self._config.user.data_source.asset_base_path + if self._config.user.data_source.usdz is not None: + asset_base_path = self._config.user.data_source.usdz.asset_base_path # Create scene_cache (pre-create to avoid pickle errors) scene_cache = self._dataset.cache_class( @@ -275,26 +268,9 @@ async def startup(self) -> None: ) # Create UnifiedDataset for on-demand scene loading - if not TRAJDATA_AVAILABLE: - raise ImportError( - "trajdata is required for data source loading. " - "Please install trajdata." - ) - data_source_config = runtime_context.config.user.data_source - if data_source_config is None: - raise ValueError("data_source is required in user config") - - self._dataset = UnifiedDataset( - desired_data=data_source_config.desired_data, - data_dirs=data_source_config.data_dirs, - cache_location=data_source_config.cache_location, - incl_vector_map=data_source_config.incl_vector_map, - rebuild_cache=data_source_config.rebuild_cache, - rebuild_maps=data_source_config.rebuild_maps, - desired_dt=data_source_config.desired_dt, - num_workers=data_source_config.num_workers, - ) + trajdata_params = data_source_config.to_trajdata_params() + self._dataset = UnifiedDataset(**trajdata_params) logger.info(f"Created UnifiedDataset with {self._dataset.num_scenes()} scenes") # Store scene_id to index mapping from RuntimeContext diff --git a/src/runtime/alpasim_runtime/runtime_context.py b/src/runtime/alpasim_runtime/runtime_context.py index bf7fada0..932ffd38 100644 --- a/src/runtime/alpasim_runtime/runtime_context.py +++ b/src/runtime/alpasim_runtime/runtime_context.py @@ -20,20 +20,12 @@ gather_versions_from_addresses, validate_scenarios, ) +from trajdata.dataset import UnifiedDataset from eval.schema import EvalConfig logger = logging.getLogger(__name__) -# Optional trajdata support -try: - from trajdata.dataset import UnifiedDataset - - TRAJDATA_AVAILABLE = True -except ImportError: - TRAJDATA_AVAILABLE = False - UnifiedDataset = None - ALL_SKIP_PER_WORKER_CONCURRENCY = 16 @@ -179,13 +171,7 @@ async def build_runtime_context( config = parse_simulator_config(user_config_path, network_config_path) eval_config = typed_parse_config(eval_config_path, EvalConfig) - # Require data_source in config (unified data flow) - if config.user.data_source is None: - raise ValueError( - "No data source specified in user config. " - "Please set 'data_source' in your YAML config file." - ) - + # Validate configuration version_ids = await gather_versions_from_addresses( config.network, config.user.endpoints, @@ -203,25 +189,13 @@ async def build_runtime_context( # Create UnifiedDataset and build scene_id to data source mapping logger.info("Creating UnifiedDataset from config") - if not TRAJDATA_AVAILABLE: - raise ImportError( - "trajdata is required for data source loading. " "Please install trajdata." - ) data_source_config = config.user.data_source - dataset = UnifiedDataset( - desired_data=data_source_config.desired_data, - data_dirs=data_source_config.data_dirs, - cache_location=data_source_config.cache_location, - incl_vector_map=data_source_config.incl_vector_map, - rebuild_cache=data_source_config.rebuild_cache, - rebuild_maps=data_source_config.rebuild_maps, - desired_dt=data_source_config.desired_dt, - num_workers=data_source_config.num_workers, - ) + trajdata_params = data_source_config.to_trajdata_params() + dataset = UnifiedDataset(**trajdata_params) logger.info( f"Created UnifiedDataset with {dataset.num_scenes()} scenes, " - f"desired_data={data_source_config.desired_data}" + f"desired_data={trajdata_params['desired_data']}" ) # Build scene_id to index mapping (once, in main process) diff --git a/src/runtime/alpasim_runtime/simulate/__main__.py b/src/runtime/alpasim_runtime/simulate/__main__.py index 1fc24e98..3cb802fc 100644 --- a/src/runtime/alpasim_runtime/simulate/__main__.py +++ b/src/runtime/alpasim_runtime/simulate/__main__.py @@ -183,12 +183,6 @@ async def run_simulation(args: argparse.Namespace) -> bool: config = parse_simulator_config(args.user_config, args.network_config) eval_config = typed_parse_config(args.eval_config, EvalConfig) - # Require data_source in config (unified data flow) - if config.user.data_source is None: - raise ValueError( - "No data source specified. Please set 'data_source' in user config YAML. " - ) - # Derive output directories from log_dir rollouts_dir = os.path.join(args.log_dir, "rollouts") telemetry_dir = os.path.join(args.log_dir, "telemetry") diff --git a/src/runtime/tests/test_config.py b/src/runtime/tests/test_config.py index 33f27c54..37fdc87e 100644 --- a/src/runtime/tests/test_config.py +++ b/src/runtime/tests/test_config.py @@ -53,38 +53,71 @@ def test_typed_parse_config_invalid_yaml(tmp_path): def test_data_source_config_defaults(): - """Test that DataSourceConfig has sensible defaults.""" + """Test that DataSourceConfig has sensible defaults for hierarchical structure.""" cfg = config.DataSourceConfig( - desired_data=["nuplan_mini"], - data_dirs={"nuplan_mini": "/data/nuplan"}, cache_location="/tmp/cache", + usdz=config.USDZSourceConfig(data_dir="/data/usdz"), ) - assert cfg.desired_data == ["nuplan_mini"] - assert cfg.data_dirs == {"nuplan_mini": "/data/nuplan"} assert cfg.cache_location == "/tmp/cache" - assert cfg.incl_vector_map is True assert cfg.rebuild_cache is False assert cfg.rebuild_maps is False - assert cfg.desired_dt == 0.1 - assert cfg.num_workers == 1 - assert cfg.num_timesteps_before == 30 - assert cfg.num_timesteps_after == 80 - assert cfg.config_dir is None - assert cfg.asset_base_path is None + assert cfg.num_workers == 4 + assert cfg.usdz is not None + assert cfg.usdz.enabled is True + assert cfg.usdz.data_dir == "/data/usdz" + assert cfg.usdz.desired_dt == 0.1 + assert cfg.usdz.incl_vector_map is True -def test_data_source_config_optional_fields(): - """Test that DataSourceConfig optional fields work correctly.""" +def test_data_source_config_to_trajdata_params(): + """Test conversion from hierarchical config to flat trajdata parameters.""" cfg = config.DataSourceConfig( - desired_data=["usdz"], - data_dirs={"usdz": "/data/usdz"}, cache_location="/tmp/cache", - config_dir="/configs", - asset_base_path="/assets", rebuild_cache=True, - desired_dt=0.05, + num_workers=8, + usdz=config.USDZSourceConfig( + data_dir="/data/usdz", + desired_dt=0.05, + asset_base_path="/assets", + ), ) - assert cfg.config_dir == "/configs" - assert cfg.asset_base_path == "/assets" - assert cfg.rebuild_cache is True - assert cfg.desired_dt == 0.05 + params = cfg.to_trajdata_params() + + assert params["desired_data"] == ["usdz"] + assert params["data_dirs"] == {"usdz": "/data/usdz"} + assert params["cache_location"] == "/tmp/cache" + assert params["rebuild_cache"] is True + assert params["num_workers"] == 8 + assert params["desired_dt"] == 0.05 + assert params["incl_vector_map"] is True + + +def test_data_source_config_multiple_sources(): + """Test configuration with multiple data sources enabled.""" + cfg = config.DataSourceConfig( + cache_location="/tmp/cache", + usdz=config.USDZSourceConfig(data_dir="/data/usdz"), + nuplan=config.NuPlanSourceConfig( + enabled=True, + data_dir="/data/nuplan", + config_dir="/configs", + ), + ) + params = cfg.to_trajdata_params() + + assert set(params["desired_data"]) == {"usdz", "nuplan"} + assert params["data_dirs"]["usdz"] == "/data/usdz" + assert params["data_dirs"]["nuplan"] == "/data/nuplan" + assert params["config_dir"] == "/configs" + assert params["num_timesteps_before"] == 30 + + +def test_data_source_config_no_sources_enabled(): + """Test that error is raised when no data sources are enabled.""" + cfg = config.DataSourceConfig( + cache_location="/tmp/cache", + nuplan=config.NuPlanSourceConfig(enabled=False), + ) + + with pytest.raises(ValueError, match="No data sources enabled"): + cfg.to_trajdata_params() diff --git a/src/wizard/configs/base_config.yaml b/src/wizard/configs/base_config.yaml index 4b5c13ed..ad6f2af4 100644 --- a/src/wizard/configs/base_config.yaml +++ b/src/wizard/configs/base_config.yaml @@ -111,6 +111,11 @@ scenes: suites_csv: - "${repo-relative:'data/scenes/sim_suites.csv'}" + # Relative path within scene_cache to the sceneset directory + # Set automatically by wizard; override for manual runtime usage + # Example: "scenesets/abc123" or "all-usdzs/sample_set/25.07_release/Batch0001" + sceneset_path: "all-usdzs/sample_set/25.07_release/Batch0001" + # \/ services.* defines the individual components of the simulation. Each of them is deployed from an image so the layout of # each item is similar. For the services that pull in local code (e.g., controller, runtime, eval) we mount the repo-relative # `src/` directory into `/mnt/src` in the container and run from there. The virtual environment is reused from the container. @@ -213,7 +218,6 @@ services: gpus: null # uses no GPUs command: - "uv run python -m alpasim_runtime.simulate" - - "--usdz-glob=/mnt/nre-data/{sceneset}/**/*.usdz" - "--user-config=/mnt/log_dir/{runtime_config_name}" - "--network-config=/mnt/log_dir/generated-network-config.yaml" - "--log-dir=/mnt/log_dir" @@ -226,34 +230,34 @@ runtime: # nr_workers and endpoints.*.n_concurrent_rollouts are set by topology configs. # Unified data source configuration using trajdata - # This replaces the old USDZ artifact glob pattern approach + # Hierarchical structure separates common settings from source-specific config data_source: - desired_data: - - usdz # Load USDZ scenes (NuRec artifacts) - - data_dirs: - # Point to all-usdzs directory to scan all scene batches recursively - # This covers paths like: all-usdzs/sample_set/25.07_release/Batch000X/ - usdz: "${scenes.scene_cache}/all-usdzs/sample_set/25.07_release/Batch0001" - - cache_location: "${defines.trajdata_cache}" # Trajdata unified cache - - # Map and cache settings - incl_vector_map: true # Include vector map data (roads, lanes, etc.) + # Common configuration (applies to all data sources) + cache_location: "${defines.trajdata_cache}" # Shared trajdata cache rebuild_cache: false # Set to true to force rebuild cache rebuild_maps: false # Set to true to force rebuild maps - - # Trajectory processing - desired_dt: 0.1 # 10 Hz sampling rate (matches control_timestep_us) num_workers: 4 # Parallel workers for cache creation - # Optional: Base path for MTGS rendering assets - asset_base_path: null - - # YAML config mode parameters (for NuPlan central token processing) - config_dir: null # Set to enable YAML batch preprocessing mode - num_timesteps_before: 30 # Timesteps before central token - num_timesteps_after: 80 # Timesteps after central token + # USDZ data source configuration (NuRec artifacts) + usdz: + enabled: true + # Use wizard's dynamic sceneset path. The wizard creates a sceneset directory + # based on selected scenes and sets sceneset_path at runtime. + # For manual runtime usage without wizard, set sceneset_path in scenes config. + data_dir: "${scenes.scene_cache}/${scenes.sceneset_path}" + desired_dt: 0.1 # 10 Hz sampling rate for trajectories + incl_vector_map: true # Include vector map data (roads, lanes, etc.) + asset_base_path: null # Optional: Base path for MTGS rendering assets + + # NuPlan data source configuration (disabled by default) + nuplan: + enabled: false + data_dir: null + config_dir: null # Set to enable YAML batch preprocessing mode + num_timesteps_before: 30 # Timesteps before central token + num_timesteps_after: 80 # Timesteps after central token + desired_dt: 0.1 + incl_vector_map: true # Enable cubic spline smoothing for trajectories smooth_trajectories: true @@ -285,8 +289,6 @@ runtime: sensorsim_cache_size: ${defines.nre_cache_size} enable_autoresume: false - # How many scenes (in particular maps) to cache in the worker local artifact cache. - artifact_cache_size: 10 simulation_config: n_sim_steps: 200 # how many steps to simulate in a rollout From b773aff174ca388eb77e9773b043f9917ac6a90d Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 20 Mar 2026 04:48:29 +0000 Subject: [PATCH 12/17] refactor: simplify and clean up codebase Address reviewer comments 2, 4, 10, 12, 13 on PR #56. Changes: - Remove unnecessary temporary variables (Comment 4) - Use data_source.rig.trajectory.time_range_us directly - Reduces cognitive load by eliminating intermediate variables - Remove trajdata availability checks (Comment 10) - trajdata is a required dependency in pyproject.toml - Trust dependency management instead of runtime try-except guards - Removes TRAJDATA_AVAILABLE checks from prepare_data and runtime_context - Improve documentation clarity (Comments 2, 12, 13) - Remove 'Key Dependencies' section from ONBOARDING.md (Comment 12) - Clarify data preparation is optional and automatic in TUTORIAL.md (Comment 13) - Remove misleading comment about desired_dt matching control_timestep_us (Comment 2) Co-Authored-By: Claude Sonnet 4.5 --- docs/ONBOARDING.md | 23 ----------------- docs/TUTORIAL.md | 3 ++- .../alpasim_runtime/prepare_data/__main__.py | 25 ++----------------- .../alpasim_runtime/unbound_rollout.py | 12 ++++----- 4 files changed, 9 insertions(+), 54 deletions(-) diff --git a/docs/ONBOARDING.md b/docs/ONBOARDING.md index 36a5e45c..cd29a047 100644 --- a/docs/ONBOARDING.md +++ b/docs/ONBOARDING.md @@ -62,29 +62,6 @@ uv run --project src/runtime python -c "..." # run in a sub-project context All members share one dependency resolution; there is no per-member version isolation. See [Plugin System](PLUGIN_SYSTEM.md) for how plugins integrate with the workspace. -## Key Dependencies - -### Trajdata (Unified Data Loading) - -AlpaSim uses [trajdata](https://github.com/NVlabs/trajdata) for unified scene data loading across -different autonomous driving datasets (USDZ, NuPlan, Waymo, etc.). The dependency is automatically -installed via `uv` from the `alpasim` branch: - -```toml -# src/runtime/pyproject.toml -[tool.uv.sources.trajdata-alpasim] -git = "https://github.com/NVlabs/trajdata" -branch = "alpasim" # Custom branch with AlpaSim-specific extensions -``` - -**You don't need to manually install trajdata** - it's handled automatically when you run -`setup_local_env.sh` or `uv sync --extra runtime`. - -The `alpasim` branch includes: -- USDZ data format support -- Custom trajectory smoothing -- Performance optimizations for simulation workloads - ## Troubleshooting **`CUDA_ERROR_UNSUPPORTED_PTX_VERSION` during NRE warmup** diff --git a/docs/TUTORIAL.md b/docs/TUTORIAL.md index 5718a1de..8d7d909a 100644 --- a/docs/TUTORIAL.md +++ b/docs/TUTORIAL.md @@ -342,7 +342,8 @@ AlpaSim uses [trajdata](https://github.com/NVlabs/trajdata/tree/alpasim) (custom for unified data loading across different autonomous driving datasets. The trajdata library is automatically installed via `uv` when you run `setup_local_env.sh`. -Before running simulations, you need to prepare a trajdata cache from your scene data. +The wizard automatically prepares the trajdata cache when needed. However, you can manually +prepare or rebuild the cache for optimization or debugging purposes. ### Preparing USDZ Scene Cache diff --git a/src/runtime/alpasim_runtime/prepare_data/__main__.py b/src/runtime/alpasim_runtime/prepare_data/__main__.py index b695efc3..6ea73f0f 100644 --- a/src/runtime/alpasim_runtime/prepare_data/__main__.py +++ b/src/runtime/alpasim_runtime/prepare_data/__main__.py @@ -37,20 +37,10 @@ from typing import Any, Dict, List, Optional import yaml +from trajdata.dataset import UnifiedDataset logger = logging.getLogger(__name__) -# Optional trajdata import -try: - from trajdata.dataset import UnifiedDataset - - TRAJDATA_AVAILABLE = True - -except ImportError: - TRAJDATA_AVAILABLE = False - UnifiedDataset = None - env_utils = None - def load_yaml_configs(config_dir: Path) -> Dict[str, List[Dict[str, str]]]: """ @@ -189,10 +179,6 @@ def preprocess_from_yaml_configs( Returns: True if successful, False otherwise. """ - if not TRAJDATA_AVAILABLE: - logger.error("trajdata is not installed. Please install it first.") - return False - # Load all YAML configs configs_by_log = load_yaml_configs(config_dir) @@ -292,9 +278,7 @@ def preprocess_basic( Returns: True if successful, False otherwise. """ - if not TRAJDATA_AVAILABLE: - logger.error("trajdata is not installed. Please install it first.") - return False + # Convert config to flat params for UnifiedDataset logger.info("Data source configuration:") logger.info(f" desired_data: {desired_data}") @@ -617,11 +601,6 @@ def main(arg_list: Optional[List[str]] = None) -> int: logger.info("Alpasim Data Preparation Tool") logger.info("=" * 60) - if not TRAJDATA_AVAILABLE: - logger.error("trajdata is not installed. Please install it first:") - logger.error("pip install trajdata") - return 1 - # Determine mode and load configuration try: if args.user_config: diff --git a/src/runtime/alpasim_runtime/unbound_rollout.py b/src/runtime/alpasim_runtime/unbound_rollout.py index cecf1e5a..929c0d6b 100644 --- a/src/runtime/alpasim_runtime/unbound_rollout.py +++ b/src/runtime/alpasim_runtime/unbound_rollout.py @@ -115,12 +115,8 @@ def create( camera_configs = list(simulation_config.cameras) - # Get time range from data source rig - rig_time_range_start = data_source.rig.trajectory.time_range_us.start - rig_time_range_stop = data_source.rig.trajectory.time_range_us.stop - control_timestamps_us_arr: np.ndarray = ( - rig_time_range_start + data_source.rig.trajectory.time_range_us.start + simulation_config.time_start_offset_us + np.arange( simulation_config.n_sim_steps + 2 @@ -129,9 +125,11 @@ def create( ) control_timestamps_us = [ - int(min(t, rig_time_range_stop - 1)) + int(min(t, data_source.rig.trajectory.time_range_us.stop - 1)) for t in control_timestamps_us_arr - if t < rig_time_range_stop + ORIGINAL_TRAJECTORY_DURATION_EXTENSION_US + if t + < data_source.rig.trajectory.time_range_us.stop + + ORIGINAL_TRAJECTORY_DURATION_EXTENSION_US ] start_us = control_timestamps_us[0] From 8cd49584d365483143ef1dc199d3b92efd17e4e9 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 20 Mar 2026 06:01:19 +0000 Subject: [PATCH 13/17] refactor: add UnifiedDataset to RuntimeContext to avoid duplication Address reviewer feedback (Comment 11): UnifiedDataset was being created twice - once in build_runtime_context() and once in DaemonEngine.startup(). This was wasteful as dataset creation can be slow (directory scanning, cache loading). Changes: - Add dataset field to RuntimeContext - DaemonEngine now uses runtime_context.dataset directly - Remove duplicate UnifiedDataset creation from engine startup The dataset is only used in the main process (not serialized to workers), so adding it to RuntimeContext is safe and eliminates the duplication overhead. --- src/runtime/alpasim_runtime/daemon/engine.py | 9 ++------- src/runtime/alpasim_runtime/runtime_context.py | 2 ++ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/runtime/alpasim_runtime/daemon/engine.py b/src/runtime/alpasim_runtime/daemon/engine.py index 362534e5..7424d133 100644 --- a/src/runtime/alpasim_runtime/daemon/engine.py +++ b/src/runtime/alpasim_runtime/daemon/engine.py @@ -267,13 +267,8 @@ async def startup(self) -> None: validate_config_scenes=self._validate_config_scenes, ) - # Create UnifiedDataset for on-demand scene loading - data_source_config = runtime_context.config.user.data_source - trajdata_params = data_source_config.to_trajdata_params() - self._dataset = UnifiedDataset(**trajdata_params) - logger.info(f"Created UnifiedDataset with {self._dataset.num_scenes()} scenes") - - # Store scene_id to index mapping from RuntimeContext + # Use UnifiedDataset from RuntimeContext (created once in build_runtime_context) + self._dataset = runtime_context.dataset self._scene_id_to_idx = runtime_context.scene_id_to_idx self._config = runtime_context.config diff --git a/src/runtime/alpasim_runtime/runtime_context.py b/src/runtime/alpasim_runtime/runtime_context.py index 932ffd38..8ecddd6e 100644 --- a/src/runtime/alpasim_runtime/runtime_context.py +++ b/src/runtime/alpasim_runtime/runtime_context.py @@ -131,6 +131,7 @@ class RuntimeContext: eval_config: EvalConfig version_ids: RolloutMetadata.VersionIds scene_id_to_idx: dict[str, int] + dataset: UnifiedDataset pools: dict[str, AddressPool] max_in_flight: int @@ -218,6 +219,7 @@ async def build_runtime_context( eval_config=eval_config, version_ids=version_ids, scene_id_to_idx=scene_id_to_idx, + dataset=dataset, pools=pools, max_in_flight=max_in_flight, ) From 6791e1b58edcd944908b024bf8659c62a2661981 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 20 Mar 2026 06:44:30 +0000 Subject: [PATCH 14/17] refactor: encapsulate scene loading into SceneLoader class Extract scene data loading logic from DaemonEngine into dedicated SceneLoader class to improve separation of concerns and reduce cognitive load. Changes: - Add SceneLoader class with lazy loading and caching - Replace _dataset, _scene_id_to_idx, _scene_id_to_data_source with single _scene_loader field in DaemonEngine - Update tests to mock SceneLoader instead of 3 separate fields Addresses reviewer comment 14 on PR #56. Co-Authored-By: Claude Sonnet 4.5 --- src/runtime/alpasim_runtime/daemon/engine.py | 77 +++------- src/runtime/alpasim_runtime/scene_loader.py | 120 +++++++++++++++ src/runtime/tests/test_daemon_engine.py | 18 ++- .../tests/test_trajdata_integration.py | 137 +++++++----------- 4 files changed, 209 insertions(+), 143 deletions(-) create mode 100644 src/runtime/alpasim_runtime/scene_loader.py diff --git a/src/runtime/alpasim_runtime/daemon/engine.py b/src/runtime/alpasim_runtime/daemon/engine.py index 7424d133..11e05b85 100644 --- a/src/runtime/alpasim_runtime/daemon/engine.py +++ b/src/runtime/alpasim_runtime/daemon/engine.py @@ -15,11 +15,10 @@ build_runtime_context, compute_num_consumers_per_worker, ) +from alpasim_runtime.scene_loader import SceneLoader from alpasim_runtime.worker.ipc import JobResult, PendingRolloutJob from alpasim_runtime.worker.runtime import WorkerRuntime, start_worker_runtime from alpasim_utils.scene_data_source import SceneDataSource -from alpasim_utils.trajdata_data_source import TrajdataDataSource -from trajdata.dataset import UnifiedDataset from eval.data import AggregationType @@ -185,9 +184,7 @@ def __init__( self._version_ids: logging_pb2.RolloutMetadata.VersionIds | None = None self._config = None # Will be set during startup - self._dataset: UnifiedDataset | None = None - self._scene_id_to_idx: dict[str, int] = {} - self._scene_id_to_data_source: dict[str, SceneDataSource] = {} + self._scene_loader: SceneLoader | None = None self._scheduler: DaemonScheduler | None = None self._worker_runtime: WorkerRuntime | None = None self._started = False @@ -199,55 +196,22 @@ def version_ids(self) -> logging_pb2.RolloutMetadata.VersionIds: return self._version_ids def _get_data_source(self, scene_id: str) -> SceneDataSource: - """Get or create a data source for the given scene_id.""" - # Check cache first - if scene_id in self._scene_id_to_data_source: - return self._scene_id_to_data_source[scene_id] - - # Lazy load from dataset - if self._dataset is None: - raise RuntimeError(f"Dataset not initialized, cannot load scene {scene_id}") - - if self._config is None: - raise RuntimeError("Config not initialized") - - scene_idx = self._scene_id_to_idx.get(scene_id) - if scene_idx is None: - raise UnknownSceneError(scene_id) - - try: - scene = self._dataset.get_scene(scene_idx) - if scene is None: - raise UnknownSceneError(scene_id) - - # Get asset_base_path from USDZ config - asset_base_path = None - if self._config.user.data_source.usdz is not None: - asset_base_path = self._config.user.data_source.usdz.asset_base_path - - # Create scene_cache (pre-create to avoid pickle errors) - scene_cache = self._dataset.cache_class( - self._dataset.cache_path, scene, self._dataset.augmentations - ) - scene_cache.set_obs_format(self._dataset.obs_format) - - # Create TrajdataDataSource - data_source = TrajdataDataSource.from_trajdata_scene( - scene=scene, - dataset=None, # Don't pass dataset to avoid pickle errors - scene_cache=scene_cache, - smooth_trajectories=self._config.user.smooth_trajectories, - asset_base_path=asset_base_path, - ) + """Get or create a data source for the given scene_id. + + Delegates to SceneLoader for lazy loading and caching. + + Args: + scene_id: Scene identifier to load - # Cache for future use - self._scene_id_to_data_source[scene_id] = data_source - logger.debug(f"Loaded data source for scene {scene_id}") - return data_source + Returns: + SceneDataSource for the requested scene - except Exception as e: - logger.error(f"Failed to load scene {scene_id}: {e}") - raise UnknownSceneError(scene_id) + Raises: + RuntimeError: If SceneLoader not initialized + """ + if self._scene_loader is None: + raise RuntimeError("SceneLoader not initialized") + return self._scene_loader.get_data_source(scene_id) async def startup(self) -> None: """Initialize the runtime context, start workers, and begin scheduling. @@ -267,9 +231,12 @@ async def startup(self) -> None: validate_config_scenes=self._validate_config_scenes, ) - # Use UnifiedDataset from RuntimeContext (created once in build_runtime_context) - self._dataset = runtime_context.dataset - self._scene_id_to_idx = runtime_context.scene_id_to_idx + # Create SceneLoader from RuntimeContext + self._scene_loader = SceneLoader( + dataset=runtime_context.dataset, + scene_id_to_idx=runtime_context.scene_id_to_idx, + config=runtime_context.config, + ) self._config = runtime_context.config num_consumers_per_worker = compute_num_consumers_per_worker( diff --git a/src/runtime/alpasim_runtime/scene_loader.py b/src/runtime/alpasim_runtime/scene_loader.py new file mode 100644 index 00000000..1893a6cc --- /dev/null +++ b/src/runtime/alpasim_runtime/scene_loader.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA Corporation + +"""SceneLoader: Manages scene data loading and caching.""" + +from __future__ import annotations + +import logging + +from alpasim_runtime.config import SimulatorConfig +from alpasim_runtime.daemon.scheduler import UnknownSceneError +from alpasim_utils.scene_data_source import SceneDataSource +from alpasim_utils.trajdata_data_source import TrajdataDataSource +from trajdata.dataset import UnifiedDataset + +logger = logging.getLogger(__name__) + + +class SceneLoader: + """Manages scene data loading and caching. + + Encapsulates UnifiedDataset, scene ID to index mapping, and lazy loading + of SceneDataSource objects. Provides a clean interface for on-demand scene + data access with automatic caching. + + Attributes: + _dataset: UnifiedDataset for accessing trajdata scenes + _scene_id_to_idx: Mapping from scene IDs to dataset indices + _config: Simulator configuration for scene parameters + _cache: Cache of loaded SceneDataSource objects + """ + + def __init__( + self, + dataset: UnifiedDataset, + scene_id_to_idx: dict[str, int], + config: SimulatorConfig, + ): + """Initialize SceneLoader with dataset and configuration. + + Args: + dataset: UnifiedDataset for scene access + scene_id_to_idx: Mapping from scene ID to dataset index + config: Simulator configuration + """ + self._dataset = dataset + self._scene_id_to_idx = scene_id_to_idx + self._config = config + self._cache: dict[str, SceneDataSource] = {} + + def get_data_source(self, scene_id: str) -> SceneDataSource: + """Get or create a data source for the given scene_id. + + Implements lazy loading with caching. On first access, creates a + TrajdataDataSource from the UnifiedDataset scene. Subsequent accesses + return the cached instance. + + Args: + scene_id: Scene identifier to load + + Returns: + SceneDataSource for the requested scene + + Raises: + UnknownSceneError: If scene_id is not found in the dataset + RuntimeError: If scene loading fails + """ + # Check cache first + if scene_id in self._cache: + return self._cache[scene_id] + + # Validate scene exists + scene_idx = self._scene_id_to_idx.get(scene_id) + if scene_idx is None: + raise UnknownSceneError(scene_id) + + try: + # Load scene from dataset + scene = self._dataset.get_scene(scene_idx) + if scene is None: + raise UnknownSceneError(scene_id) + + # Get asset_base_path from USDZ config if available + asset_base_path = None + if self._config.user.data_source.usdz is not None: + asset_base_path = self._config.user.data_source.usdz.asset_base_path + + # Create scene_cache (pre-create to avoid pickle errors) + scene_cache = self._dataset.cache_class( + self._dataset.cache_path, scene, self._dataset.augmentations + ) + scene_cache.set_obs_format(self._dataset.obs_format) + + # Create TrajdataDataSource + data_source = TrajdataDataSource.from_trajdata_scene( + scene=scene, + dataset=None, # Don't pass dataset to avoid pickle errors + scene_cache=scene_cache, + smooth_trajectories=self._config.user.smooth_trajectories, + asset_base_path=asset_base_path, + ) + + # Cache for future use + self._cache[scene_id] = data_source + logger.debug(f"Loaded data source for scene {scene_id}") + return data_source + + except Exception as e: + logger.error(f"Failed to load scene {scene_id}: {e}") + raise RuntimeError(f"Scene loading failed for {scene_id}") from e + + @property + def num_scenes(self) -> int: + """Return total number of scenes available in the dataset.""" + return self._dataset.num_scenes() + + @property + def num_cached(self) -> int: + """Return number of scenes currently cached.""" + return len(self._cache) diff --git a/src/runtime/tests/test_daemon_engine.py b/src/runtime/tests/test_daemon_engine.py index 7ffb047d..5f7929c8 100644 --- a/src/runtime/tests/test_daemon_engine.py +++ b/src/runtime/tests/test_daemon_engine.py @@ -185,8 +185,10 @@ async def wait_request(self, request_id: str): log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_idx = {"clipgt-a": 0} - engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} + # Mock SceneLoader with a fake data source + mock_scene_loader = MagicMock() + mock_scene_loader.get_data_source.return_value = MagicMock() + engine._scene_loader = mock_scene_loader engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), @@ -233,8 +235,10 @@ async def wait_request(self, request_id): log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_idx = {"clipgt-a": 0} - engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} + # Mock SceneLoader with a fake data source + mock_scene_loader = MagicMock() + mock_scene_loader.get_data_source.return_value = MagicMock() + engine._scene_loader = mock_scene_loader engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), @@ -286,8 +290,10 @@ async def wait_request(self, request_id): log_dir="/tmp/log", ) engine._started = True - engine._scene_id_to_idx = {"clipgt-a": 0} - engine._scene_id_to_data_source = {"clipgt-a": MagicMock()} + # Mock SceneLoader with a fake data source + mock_scene_loader = MagicMock() + mock_scene_loader.get_data_source.return_value = MagicMock() + engine._scene_loader = mock_scene_loader engine._version_ids = RolloutMetadata.VersionIds( runtime_version=VersionId(version_id="runtime", git_hash="a"), sensorsim_version=VersionId(version_id="sensorsim", git_hash="b"), diff --git a/src/runtime/tests/test_trajdata_integration.py b/src/runtime/tests/test_trajdata_integration.py index 428a2b44..77185f18 100644 --- a/src/runtime/tests/test_trajdata_integration.py +++ b/src/runtime/tests/test_trajdata_integration.py @@ -5,8 +5,7 @@ from __future__ import annotations -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest from alpasim_grpc.v0 import runtime_pb2 @@ -119,90 +118,64 @@ def fake_get_data_source(scene_id: str): def test_daemon_engine_get_data_source_caching( mock_trajdata_dataset, mock_trajdata_scene ): - """Test that DaemonEngine caches data sources and doesn't reload them.""" - with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): - engine = DaemonEngine( - user_config="u.yaml", - network_config="n.yaml", - eval_config="e.yaml", - log_dir="/tmp/log", - ) - - # Set up engine state - engine._started = True - engine._dataset = mock_trajdata_dataset - engine._scene_id_to_idx = {"test_scene_001": 0} - engine._scene_id_to_data_source = {} - engine._config = SimpleNamespace( - user=SimpleNamespace( - smooth_trajectories=True, - data_source=SimpleNamespace(asset_base_path="/tmp/assets"), - ) - ) - - # Mock TrajdataDataSource.from_trajdata_scene - with patch( - "alpasim_runtime.daemon.engine.TrajdataDataSource.from_trajdata_scene" - ) as mock_from_scene: - mock_data_source = MagicMock() - mock_from_scene.return_value = mock_data_source - - # First call should create the data source - data_source_1 = engine._get_data_source("test_scene_001") - assert data_source_1 is mock_data_source - assert mock_from_scene.call_count == 1 - - # Second call should return cached data source - data_source_2 = engine._get_data_source("test_scene_001") - assert data_source_2 is mock_data_source - assert mock_from_scene.call_count == 1 # Not called again - - # Verify cache contains the data source - assert "test_scene_001" in engine._scene_id_to_data_source - assert engine._scene_id_to_data_source["test_scene_001"] is mock_data_source + """Test that DaemonEngine uses SceneLoader for caching data sources.""" + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + # Set up engine with mock SceneLoader + engine._started = True + mock_scene_loader = MagicMock() + mock_data_source = MagicMock() + mock_scene_loader.get_data_source.return_value = mock_data_source + engine._scene_loader = mock_scene_loader + + # First call should delegate to SceneLoader + data_source_1 = engine._get_data_source("test_scene_001") + assert data_source_1 is mock_data_source + assert mock_scene_loader.get_data_source.call_count == 1 + + # Second call should also delegate (SceneLoader handles caching internally) + data_source_2 = engine._get_data_source("test_scene_001") + assert data_source_2 is mock_data_source + assert mock_scene_loader.get_data_source.call_count == 2 def test_daemon_engine_get_data_source_unknown_scene(): """Test that _get_data_source raises UnknownSceneError for unknown scenes.""" - with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): - engine = DaemonEngine( - user_config="u.yaml", - network_config="n.yaml", - eval_config="e.yaml", - log_dir="/tmp/log", - ) - - engine._started = True - engine._dataset = MagicMock() - engine._scene_id_to_idx = {"known_scene": 0} - engine._scene_id_to_data_source = {} - engine._config = SimpleNamespace( - user=SimpleNamespace( - smooth_trajectories=True, - data_source=SimpleNamespace(asset_base_path=None), - ) - ) - - with pytest.raises(UnknownSceneError) as exc_info: - engine._get_data_source("unknown_scene") - - assert exc_info.value.scene_id == "unknown_scene" + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + engine._started = True + # Mock SceneLoader that raises UnknownSceneError + mock_scene_loader = MagicMock() + mock_scene_loader.get_data_source.side_effect = UnknownSceneError("unknown_scene") + engine._scene_loader = mock_scene_loader + + with pytest.raises(UnknownSceneError) as exc_info: + engine._get_data_source("unknown_scene") + + assert exc_info.value.scene_id == "unknown_scene" def test_daemon_engine_get_data_source_without_dataset(): - """Test that _get_data_source raises RuntimeError if dataset is not initialized.""" - with patch("alpasim_runtime.daemon.engine.TRAJDATA_AVAILABLE", True): - engine = DaemonEngine( - user_config="u.yaml", - network_config="n.yaml", - eval_config="e.yaml", - log_dir="/tmp/log", - ) - - engine._started = True - engine._dataset = None - engine._scene_id_to_idx = {"test_scene": 0} - engine._scene_id_to_data_source = {} - - with pytest.raises(RuntimeError, match="Dataset not initialized"): - engine._get_data_source("test_scene") + """Test that _get_data_source raises RuntimeError if SceneLoader is not initialized.""" + engine = DaemonEngine( + user_config="u.yaml", + network_config="n.yaml", + eval_config="e.yaml", + log_dir="/tmp/log", + ) + + engine._started = True + engine._scene_loader = None + + with pytest.raises(RuntimeError, match="SceneLoader not initialized"): + engine._get_data_source("test_scene") From 74399c972e73e8865cfbd6a99dca77c6b1bac2cc Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Tue, 24 Mar 2026 15:01:56 +0000 Subject: [PATCH 15/17] refactor: improve prepare_data code quality and unify preprocessing Major improvements: - Unify preprocess_basic and preprocess_from_yaml_configs into single function - Fix to_trajdata_params() to handle multiple datasets via dataset_kwargs - Separate user-config and CLI paths with clear validation - Extract process_nuplan_yaml_configs() for better modularity - Simplify code structure and remove duplication Co-Authored-By: Claude Sonnet 4.5 (1M context) --- src/runtime/alpasim_runtime/config.py | 96 ++- .../alpasim_runtime/prepare_data/__init__.py | 70 ++- .../alpasim_runtime/prepare_data/__main__.py | 578 +++++++----------- 3 files changed, 293 insertions(+), 451 deletions(-) diff --git a/src/runtime/alpasim_runtime/config.py b/src/runtime/alpasim_runtime/config.py index 30740d94..6a1199da 100644 --- a/src/runtime/alpasim_runtime/config.py +++ b/src/runtime/alpasim_runtime/config.py @@ -8,7 +8,7 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Optional, Type, TypeVar, cast +from typing import Any, Dict, Optional, Type, TypeVar, cast from alpasim_utils.scenario import VehicleConfig from alpasim_utils.yaml_utils import load_yaml_dict @@ -18,73 +18,58 @@ @dataclass -class USDZSourceConfig: - """Configuration for USDZ data source. +class GenericSourceConfig: + """Generic configuration for any trajdata-supported dataset. + + This unified config supports all trajdata datasets (USDZ, NuPlan, Waymo, + nuScenes, Lyft, Argoverse, etc.) with a flexible extra_params field for + dataset-specific options. Attributes: enabled: Whether this data source is enabled - data_dir: Path to directory containing USDZ files + data_dir: Path to dataset directory desired_dt: Desired time delta between trajectory frames in seconds incl_vector_map: Whether to load vector maps (roads, lanes, etc.) - asset_base_path: Base path for rendering assets (e.g., MTGS assets) + extra_params: Dataset-specific parameters (e.g., NuPlan's config_dir, + USDZ's asset_base_path, etc.) + + Example extra_params: + - NuPlan: {"config_dir": "/path", "num_timesteps_before": 30, "num_timesteps_after": 80} + - USDZ: {"asset_base_path": "/assets"} + - Waymo: {} (no extra params needed) """ enabled: bool = True data_dir: str = MISSING - desired_dt: float = 0.1 # 10 Hz sampling - incl_vector_map: bool = True - asset_base_path: Optional[str] = None - - -@dataclass -class NuPlanSourceConfig: - """Configuration for NuPlan data source. - - Attributes: - enabled: Whether this data source is enabled - data_dir: Path to NuPlan dataset directory - config_dir: Directory containing YAML scene config files for batch mode - num_timesteps_before: Number of timesteps before central token (batch mode) - num_timesteps_after: Number of timesteps after central token (batch mode) - desired_dt: Desired time delta between frames in seconds - incl_vector_map: Whether to load vector maps - """ - - enabled: bool = False - data_dir: Optional[str] = None - config_dir: Optional[str] = None - num_timesteps_before: int = 30 - num_timesteps_after: int = 80 - desired_dt: float = 0.1 - incl_vector_map: bool = True + extra_params: Dict[str, Any] = field(default_factory=dict) @dataclass class DataSourceConfig: """Configuration for unified data loading through trajdata. - This provides a hierarchical structure where common configuration is separated - from data source-specific settings, making it easier to understand which - parameters affect which data sources. + Supports dynamic registration of any trajdata dataset (USDZ, NuPlan, Waymo, + nuScenes, Lyft, Argoverse, etc.) through the 'sources' dictionary. Attributes: cache_location: Path to shared trajdata cache directory rebuild_cache: Whether to force rebuild the cache for all sources rebuild_maps: Whether to force rebuild maps for all sources num_workers: Number of parallel workers for cache creation - usdz: USDZ-specific configuration - nuplan: NuPlan-specific configuration + sources: Dictionary mapping dataset names to their configurations + (e.g., {"usdz": GenericSourceConfig(...), "waymo": ...}) """ # Common configuration (applies to all data sources) cache_location: str = MISSING + desired_dt: float = 0.1 # 10 Hz sampling + incl_vector_map: bool = True rebuild_cache: bool = False rebuild_maps: bool = False - num_workers: int = 4 + num_workers: int = 1 # Conservative default for stability; increase for production - # Source-specific configurations - usdz: Optional[USDZSourceConfig] = None - nuplan: Optional[NuPlanSourceConfig] = None + # New extensible source configuration (preferred) + sources: Dict[str, GenericSourceConfig] = field(default_factory=dict) def to_trajdata_params(self) -> dict: """Convert hierarchical config to flat parameters for trajdata's UnifiedDataset. @@ -97,39 +82,32 @@ def to_trajdata_params(self) -> dict: """ desired_data = [] data_dirs = {} + dataset_kwargs = {} - # Collect enabled sources - if self.usdz is not None and self.usdz.enabled: - desired_data.append("usdz") - data_dirs["usdz"] = self.usdz.data_dir - - if self.nuplan is not None and self.nuplan.enabled: - desired_data.append("nuplan") - data_dirs["nuplan"] = self.nuplan.data_dir + for dataset_name, source in self.sources.items(): + if source.enabled: + desired_data.append(dataset_name) + data_dirs[dataset_name] = source.data_dir + if source.extra_params: + dataset_kwargs[dataset_name] = source.extra_params if not desired_data: raise ValueError("No data sources enabled in configuration") - # Use first enabled source for common parameters - # (desired_dt, incl_vector_map are typically consistent across sources) - primary_source = self.usdz if (self.usdz and self.usdz.enabled) else self.nuplan - params = { "desired_data": desired_data, "data_dirs": data_dirs, "cache_location": self.cache_location, + "incl_vector_map": self.incl_vector_map, "rebuild_cache": self.rebuild_cache, "rebuild_maps": self.rebuild_maps, "num_workers": self.num_workers, - "desired_dt": primary_source.desired_dt, - "incl_vector_map": primary_source.incl_vector_map, + "desired_dt": self.desired_dt, } - # Add source-specific parameters - if self.nuplan and self.nuplan.enabled and self.nuplan.config_dir: - params["config_dir"] = self.nuplan.config_dir - params["num_timesteps_before"] = self.nuplan.num_timesteps_before - params["num_timesteps_after"] = self.nuplan.num_timesteps_after + # Add dataset-specific kwargs if any source has extra_params + if dataset_kwargs: + params["dataset_kwargs"] = dataset_kwargs return params diff --git a/src/runtime/alpasim_runtime/prepare_data/__init__.py b/src/runtime/alpasim_runtime/prepare_data/__init__.py index 8b2ada83..0ccbb123 100644 --- a/src/runtime/alpasim_runtime/prepare_data/__init__.py +++ b/src/runtime/alpasim_runtime/prepare_data/__init__.py @@ -4,53 +4,69 @@ """ Data preprocessing module for building trajdata cache. -This module provides tools for preparing scene data before running simulations. -It supports: +This module provides two approaches for preparing scene data before simulations: -1. Basic preprocessing - Build trajdata cache for all scenes in a dataset -2. YAML config preprocessing - Batch process specific scenes based on YAML configs -3. Central token mode - Process scenes around specific central tokens (NuPlan) +1. **User Config Path** (Recommended for complex scenarios): + - Load configuration from YAML file + - Supports multiple data sources with individual settings + - Automatic NuPlan YAML batch processing when config_dir is provided + - Full control over per-dataset parameters -Main functions: - - preprocess_basic: Basic preprocessing for all scenes - - preprocess_from_yaml_configs: Batch preprocessing from YAML configs - - load_yaml_configs: Load YAML configuration files +2. **CLI Path** (For simple, quick preprocessing): + - Specify parameters via command line + - Uniform parameters applied to all datasets + - Good for testing or simple caching tasks -Example usage: +Main exports: + - preprocess_basic: Unified preprocessing function (handles both basic and NuPlan YAML modes) + - process_nuplan_yaml_configs: Process NuPlan YAML configs into central_tokens format + - load_yaml_configs: Load and parse YAML configuration files + - PrepareDataConfig: Configuration class for CLI mode + - main: CLI entry point - from alpasim_runtime.prepare_data import preprocess_basic, preprocess_from_yaml_configs +Example usage (programmatic): - # Basic preprocessing - preprocess_basic( - desired_data=["nuplan_test"], - data_dirs={"nuplan_test": "/path/to/nuplan"}, - cache_location="/path/to/cache", - ) + from alpasim_runtime.prepare_data import preprocess_basic, PrepareDataConfig - # YAML config batch preprocessing - preprocess_from_yaml_configs( - config_dir=Path("/path/to/configs"), + # Simple preprocessing with CLI config + config = PrepareDataConfig( + desired_data=["waymo"], + data_dirs={"waymo": "/path/to/waymo"}, cache_location="/path/to/cache", - data_dirs={"nuplan_test": "/path/to/nuplan"}, - env_name="nuplan_test", - num_timesteps_before=30, - num_timesteps_after=80, + desired_dt=0.1, ) + preprocess_basic(config, verbose=True) + + # Or use user config (recommended for production) + from alpasim_runtime.config import typed_parse_config, UserSimulatorConfig + user_config = typed_parse_config("user.yaml", UserSimulatorConfig) + preprocess_basic(user_config.data_source, verbose=True) CLI usage: - python -m alpasim_runtime.prepare_data --help + # Simple mode + python -m alpasim_runtime.prepare_data \\ + --desired-data waymo \\ + --data-dir /path/to/waymo \\ + --cache-location /path/to/cache + + # Complex mode with user config + python -m alpasim_runtime.prepare_data \\ + --user-config user.yaml \\ + --rebuild-cache """ from alpasim_runtime.prepare_data.__main__ import ( + PrepareDataConfig, load_yaml_configs, main, preprocess_basic, - preprocess_from_yaml_configs, + process_nuplan_yaml_configs, ) __all__ = [ "preprocess_basic", - "preprocess_from_yaml_configs", + "process_nuplan_yaml_configs", "load_yaml_configs", + "PrepareDataConfig", "main", ] diff --git a/src/runtime/alpasim_runtime/prepare_data/__main__.py b/src/runtime/alpasim_runtime/prepare_data/__main__.py index 6ea73f0f..fba88dc8 100644 --- a/src/runtime/alpasim_runtime/prepare_data/__main__.py +++ b/src/runtime/alpasim_runtime/prepare_data/__main__.py @@ -4,26 +4,30 @@ """ Data preprocessing CLI for building trajdata cache. -This module provides a command-line interface for preparing scene data before -running simulations. It supports: +This module provides two clear paths for data preprocessing: -1. Basic preprocessing - Build trajdata cache for all scenes in a dataset -2. YAML config preprocessing - Batch process specific scenes based on YAML configs -3. Central token mode - Process scenes around specific central tokens (NuPlan) +1. **User Config Path** (--user-config): For complex scenarios + - Load full configuration from YAML file + - Supports multiple data sources, hierarchical config + - Supports YAML batch mode (NuPlan central_tokens) + - CLI overrides limited to: --rebuild-cache, --rebuild-maps, --verbose + +2. **CLI Path**: For simple, quick preprocessing + - Specify all parameters via command line + - Single dataset preprocessing only + - Basic preprocessing mode only (no YAML batch mode) + - Good for testing or simple caching tasks Usage Examples: - # Basic preprocessing using user-config - python -m alpasim_runtime.prepare_data --user-config user.yaml + # Complex: Use user config with optional overrides + python -m alpasim_runtime.prepare_data --user-config user.yaml --rebuild-cache - # Basic preprocessing with explicit parameters + # Simple: Direct CLI parameters for basic preprocessing python -m alpasim_runtime.prepare_data \\ --desired-data nuplan_test \\ --data-dir /path/to/nuplan \\ --cache-location /path/to/cache - - # Rebuild cache even if it exists - python -m alpasim_runtime.prepare_data --user-config user.yaml --rebuild-cache """ from __future__ import annotations @@ -33,15 +37,58 @@ import sys import time from collections import defaultdict +from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional import yaml +from alpasim_runtime.config import UserSimulatorConfig, typed_parse_config +from omegaconf import OmegaConf from trajdata.dataset import UnifiedDataset logger = logging.getLogger(__name__) +@dataclass +class PrepareDataConfig: + """Configuration for CLI-based data preprocessing. + + This is used ONLY for CLI mode. User config mode uses DataSourceConfig directly. + + Note: CLI mode only supports basic preprocessing. For YAML batch mode (NuPlan), + use user-config files with config_dir in extra_params. + """ + + # Data source parameters (required) + desired_data: List[str] + data_dirs: Dict[str, str] + cache_location: str + + # Optional preprocessing parameters + rebuild_cache: bool = False + rebuild_maps: bool = False + incl_vector_map: bool = True + desired_dt: float = 0.1 + num_workers: int = 1 + + def to_trajdata_params(self) -> dict: + """Convert to flat parameters for trajdata's UnifiedDataset. + + Returns: + Dictionary with keys expected by UnifiedDataset constructor + """ + return { + "desired_data": self.desired_data, + "data_dirs": self.data_dirs, + "cache_location": self.cache_location, + "rebuild_cache": self.rebuild_cache, + "rebuild_maps": self.rebuild_maps, + "num_workers": self.num_workers, + "desired_dt": self.desired_dt, + "incl_vector_map": self.incl_vector_map, + } + + def load_yaml_configs(config_dir: Path) -> Dict[str, List[Dict[str, str]]]: """ Load all yaml configuration files and group them by central_log. @@ -114,21 +161,14 @@ def python_tuple_constructor(loader, tag_suffix, node): ) continue - # Extract first central_token - configs_by_log[central_log].append( - { - "central_token": central_tokens[0], - "logfile": central_log, - "yaml_file": str(yaml_file), - } - ) - # Extract every central_token - # for token in central_tokens: - # configs_by_log[central_log].append({ - # 'central_token': token, - # 'logfile': central_log, - # 'yaml_file': str(yaml_file), - # }) + for token in central_tokens: + configs_by_log[central_log].append( + { + "central_token": token, + "logfile": central_log, + "yaml_file": str(yaml_file), + } + ) except Exception as e: logger.error(f"Failed to load {yaml_file.name}: {e}") @@ -143,50 +183,29 @@ def python_tuple_constructor(loader, tag_suffix, node): return dict(configs_by_log) -def preprocess_from_yaml_configs( - config_dir: Path, - cache_location: str, - data_dirs: Dict[str, str], - env_name: str = "nuplan_test", - rebuild_cache: bool = True, - rebuild_maps: bool = False, - num_workers: int = 1, - desired_dt: float = 0.5, - num_timesteps_before: int = 30, - num_timesteps_after: int = 80, - verbose: bool = True, -) -> bool: - """ - Batch preprocess data based on YAML configuration files. - - This function reads YAML config files containing central_log and central_tokens, - and processes only those specific scenes. This is useful for processing - specific scenarios without loading the entire dataset. +def process_nuplan_yaml_configs( + dataset_name: str, extra_params: Dict[str, Any] +) -> Optional[Dict[str, Any]]: + """Process NuPlan YAML configuration files into central_tokens_config format. Args: - config_dir: Directory containing YAML configuration files. - cache_location: Path to cache directory. - data_dirs: Dictionary of dataset name to data directory paths. - env_name: Environment name (e.g., "nuplan_test"). - rebuild_cache: Whether to rebuild cache. - rebuild_maps: Whether to rebuild maps. - num_workers: Number of worker processes. - desired_dt: Desired timestep duration in seconds. - num_timesteps_before: Number of timesteps before the central token. - num_timesteps_after: Number of timesteps after the central token. - verbose: Whether to show verbose logs. + dataset_name: Name of the NuPlan dataset (e.g., 'nuplan_mini', 'nuplan_test') + extra_params: Dictionary containing 'config_dir' and optional timestep parameters Returns: - True if successful, False otherwise. + Processed dataset kwargs with central_tokens_config, or None if no configs found """ - # Load all YAML configs + logger.info(f"Processing NuPlan YAML configs for {dataset_name}") + config_dir = Path(extra_params["config_dir"]) + + # Load YAML configs configs_by_log = load_yaml_configs(config_dir) if not configs_by_log: - logger.error("No valid YAML configuration files found.") - return False + logger.warning(f"No valid YAML configs found in {config_dir}") + return None - # Merge all configs into a single central_tokens_config list + # Build central_tokens_config list all_central_tokens_config: List[Dict[str, Any]] = [] for _, configs in configs_by_log.items(): for cfg in configs: @@ -194,122 +213,72 @@ def preprocess_from_yaml_configs( { "central_token": cfg["central_token"], "logfile": cfg["logfile"], - "num_timesteps_before": num_timesteps_before, - "num_timesteps_after": num_timesteps_after, } ) - logger.info(f"Total {len(all_central_tokens_config)} central tokens to process") - - try: - # Create cache directory - Path(cache_location).mkdir(parents=True, exist_ok=True) - - # Create UnifiedDataset (this triggers cache building) - logger.info("Creating UnifiedDataset with YAML configs...") - start_time = time.perf_counter() + logger.info(f" Found {len(all_central_tokens_config)} central tokens") - dataset = UnifiedDataset( - dataset_kwargs={ - "central_tokens_config": all_central_tokens_config, - "num_timesteps_before": num_timesteps_before, - "num_timesteps_after": num_timesteps_after, - }, - desired_data=[env_name], - cache_location=cache_location, - rebuild_cache=rebuild_cache, - rebuild_maps=rebuild_maps, - require_map_cache=False, - num_workers=num_workers, - desired_dt=desired_dt, - verbose=verbose, - data_dirs=data_dirs, - ) - - elapsed = time.perf_counter() - start_time - - # Get both scene index count and total dataset length - num_scenes = dataset.num_scenes() - logger.info("=" * 80) - logger.info("Preprocessing completed!") - logger.info(f" Num Scenes: {num_scenes}") - logger.info(f" Time elapsed: {elapsed:.2f} seconds") - logger.info("=" * 80) - - return True - - except Exception as e: - logger.error(f"Error during preprocessing: {e}") - if verbose: - import traceback - - traceback.print_exc() - return False + # Return processed config + return { + "central_tokens_config": all_central_tokens_config, + "num_timesteps_before": extra_params.get("num_timesteps_before", 30), + "num_timesteps_after": extra_params.get("num_timesteps_after", 80), + } -def preprocess_basic( - desired_data: List[str], - data_dirs: Dict[str, str], - cache_location: str, - rebuild_cache: bool = False, - rebuild_maps: bool = False, - incl_vector_map: bool = True, - desired_dt: float = 0.1, - num_workers: int = 1, - verbose: bool = True, - list_scenes: bool = False, - smooth_trajectories: bool = True, -) -> bool: - """ - Basic preprocessing - build trajdata cache for all scenes. +def preprocess_basic(config: Any, verbose: bool = True) -> bool: + """Basic preprocessing - build trajdata cache for all scenes. Args: - desired_data: List of dataset names to load. - data_dirs: Dict mapping dataset names to their data directories. - cache_location: Path to trajdata cache directory. - rebuild_cache: Whether to force rebuild cache. - rebuild_maps: Whether to rebuild map cache. - incl_vector_map: Whether to include vector maps. - desired_dt: Desired time delta between frames in seconds. - num_workers: Number of workers for data loading. - verbose: Whether to show verbose output. - list_scenes: Whether to list all available scenes after preparation. + config: Configuration (supports both PrepareDataConfig from CLI and + DataSourceConfig from user config). + verbose: Whether to show verbose output from trajdata. Returns: True if successful, False otherwise. """ - # Convert config to flat params for UnifiedDataset + params = config.to_trajdata_params() logger.info("Data source configuration:") - logger.info(f" desired_data: {desired_data}") - logger.info(f" data_dirs: {data_dirs}") - logger.info(f" cache_location: {cache_location}") - logger.info(f" rebuild_cache: {rebuild_cache}") - logger.info(f" desired_dt: {desired_dt}") - logger.info(f" smooth_trajectories: {smooth_trajectories}") + logger.info(f" cache_location: {params['cache_location']}") + logger.info(f" desired_dt: {params['desired_dt']}") + logger.info(f" rebuild_cache: {params['rebuild_cache']}") + logger.info(f" rebuild_maps: {params['rebuild_maps']}") + logger.info(f" desired_data: {params['desired_data']}") + logger.info(f" data_dirs: {params['data_dirs']}") + + # Process NuPlan-specific YAML configs if present + dataset_kwargs = params.get("dataset_kwargs", {}) + if dataset_kwargs: + for dataset_name, extra_params in dataset_kwargs.items(): + # Check if this is a NuPlan dataset (nuplan, nuplan_mini, nuplan_test, etc.) + if "nuplan" in dataset_name.lower() and "config_dir" in extra_params: + processed_config = process_nuplan_yaml_configs( + dataset_name, extra_params + ) + if processed_config: + dataset_kwargs[dataset_name] = processed_config # Create cache directory - cache_path = Path(cache_location) + cache_path = Path(params["cache_location"]) cache_path.mkdir(parents=True, exist_ok=True) - logger.info(f"Cache directory: {cache_path}") # Build UnifiedDataset (this triggers cache building) - logger.info("Creating UnifiedDataset (this may take a while)...") + logger.info("Creating UnifiedDataset...") start_time = time.perf_counter() try: dataset = UnifiedDataset( - desired_data=desired_data, - data_dirs=data_dirs, - cache_location=cache_location, - incl_vector_map=incl_vector_map, - rebuild_cache=rebuild_cache, - rebuild_maps=rebuild_maps, - # require_map_cache=False, - desired_dt=desired_dt, - num_workers=num_workers, + desired_data=params["desired_data"], + data_dirs=params["data_dirs"], + cache_location=params["cache_location"], + incl_vector_map=params["incl_vector_map"], + rebuild_cache=params["rebuild_cache"], + rebuild_maps=params["rebuild_maps"], + desired_dt=params["desired_dt"], + num_workers=params["num_workers"], + dataset_kwargs=dataset_kwargs, verbose=verbose, - dataset_kwargs={"smooth_trajectories": smooth_trajectories}, ) except Exception as e: logger.error(f"Failed to create UnifiedDataset: {e}") @@ -321,24 +290,10 @@ def preprocess_basic( elapsed = time.perf_counter() - start_time logger.info(f"UnifiedDataset created in {elapsed:.2f} seconds") - # Get both scene index count and total dataset length + # Get scene count num_scenes = dataset.num_scenes() logger.info(f"Scene files (logs): {num_scenes}") - # List scenes if requested - if list_scenes and num_scenes > 0: - logger.info("Available scenes:") - max_display = min(num_scenes, 100) - for i in range(max_display): - try: - scene = dataset.get_scene(i) - logger.info(f" [{i}] {scene.name}") - except Exception as e: - logger.warning(f" [{i}] (failed to load: {e})") - - if num_scenes > 100: - logger.info(f" ... and {num_scenes - 100} more scenes") - logger.info("Data preparation complete!") return True @@ -346,7 +301,12 @@ def preprocess_basic( def create_arg_parser() -> argparse.ArgumentParser: """Create argument parser for prepare_data CLI.""" parser = argparse.ArgumentParser( - description="Prepare scene data and build trajdata cache for alpasim simulations.", + description=( + "Prepare scene data and build trajdata cache for alpasim simulations.\n\n" + "Two modes:\n" + " 1. User config (--user-config): Complex scenarios with full YAML config\n" + " 2. CLI mode (--desired-data + --data-dir): Simple, direct preprocessing" + ), formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) @@ -409,41 +369,9 @@ def create_arg_parser() -> argparse.ArgumentParser: action="store_true", help="Exclude vector maps (default: include)", ) - preprocess_group.add_argument( - "--smooth-trajectories", - type=str, - choices=["true", "false", "True", "False"], - default=None, - help="Enable/disable cubic spline smoothing for trajectories (default: True)", - ) - - # YAML config mode options (NuPlan specific) - yaml_group = parser.add_argument_group("YAML Config Mode (NuPlan)") - yaml_group.add_argument( - "--num-timesteps-before", - type=int, - default=30, - help="Number of timesteps before central token (default: 30)", - ) - yaml_group.add_argument( - "--num-timesteps-after", - type=int, - default=80, - help="Number of timesteps after central token (default: 80)", - ) # Output options output_group = parser.add_argument_group("Output Options") - output_group.add_argument( - "--validate-only", - action="store_true", - help="Only validate configuration without building cache", - ) - output_group.add_argument( - "--list-scenes", - action="store_true", - help="List all available scenes after preparation", - ) output_group.add_argument( "--log-level", type=str, @@ -453,57 +381,29 @@ def create_arg_parser() -> argparse.ArgumentParser: ) output_group.add_argument( "--verbose", - action="store_true", + action=argparse.BooleanOptionalAction, default=True, - help="Show verbose output (default: True)", - ) - output_group.add_argument( - "--quiet", - action="store_true", - help="Suppress verbose output", + help="Show verbose output (default: enabled)", ) return parser -def load_config_from_file(config_path: str) -> Dict[str, Any]: - """Load data source configuration from user config file.""" - from alpasim_runtime.config import UserSimulatorConfig, typed_parse_config - - user_config = typed_parse_config(config_path, UserSimulatorConfig) - - if user_config.data_source is None: - raise ValueError( - f"No data_source configuration found in {config_path}. " - "Please add a 'data_source' section to your user config." - ) - - ds = user_config.data_source - return { - "desired_data": ds.desired_data, - "data_dirs": ds.data_dirs, - "cache_location": ds.cache_location, - "config_dir": ds.config_dir, - "incl_vector_map": ds.incl_vector_map, - "rebuild_cache": ds.rebuild_cache, - "rebuild_maps": ds.rebuild_maps, - "desired_dt": ds.desired_dt, - "num_workers": ds.num_workers, - "num_timesteps_before": ds.num_timesteps_before, - "num_timesteps_after": ds.num_timesteps_after, - "smooth_trajectories": user_config.smooth_trajectories, - } - - def parse_data_dirs( data_dirs_args: Optional[List[str]], desired_data: List[str] ) -> Dict[str, str]: - """ - Parse data directory arguments into a dict. + """Parse data directory arguments into a dict. Supports two formats: - "dataset_name=/path/to/data" - explicit mapping - "/path/to/data" - auto-map to desired_data entries in order + + Args: + data_dirs_args: List of data directory arguments + desired_data: List of dataset names + + Returns: + Dictionary mapping dataset names to data directories """ if not data_dirs_args: return {} @@ -520,66 +420,94 @@ def parse_data_dirs( if i < len(desired_data): result[desired_data[i]] = arg else: - # Use as default for remaining datasets + # Use as default for remaining datasets (with warning) for ds in desired_data[len(result) :]: if ds not in result: + logger.warning( + f"Dataset '{ds}' has no explicit data_dir, using last provided: '{arg}'" + ) result[ds] = arg return result -def run_yaml_batch_preprocessing( - config_dir: str, - desired_data: List[str], - data_dirs: Dict[str, str], - cache_location: str, - rebuild_cache: bool = False, - rebuild_maps: bool = False, - num_workers: int = 1, - desired_dt: float = 0.5, - num_timesteps_before: int = 30, - num_timesteps_after: int = 80, - verbose: bool = True, -) -> int: +def run_from_user_config(config_path: str, args: argparse.Namespace) -> bool: + """Run preprocessing from user config file with minimal CLI overrides. + + Args: + config_path: Path to user config YAML file + args: Parsed command line arguments (used only for overrides) + + Returns: + True if successful, False otherwise """ - Execute YAML batch preprocessing mode. + logger.info(f"Loading configuration from: {config_path}") + user_config = typed_parse_config(config_path, UserSimulatorConfig) + user_config = OmegaConf.to_object(user_config) - This is a unified entry point for YAML-based batch preprocessing, - whether triggered from user config file or CLI arguments. + config = user_config.data_source + + # Apply minimal CLI overrides (only top-level flags) + if args.rebuild_cache: + config.rebuild_cache = True + logger.info("CLI override: rebuild_cache=True") + if args.rebuild_maps: + config.rebuild_maps = True + logger.info("CLI override: rebuild_maps=True") + + # Use unified preprocessing (handles both basic and YAML batch mode) + return preprocess_basic(config, verbose=args.verbose) + + +def run_from_cli(args: argparse.Namespace) -> bool: + """Run preprocessing from CLI arguments directly. + + CLI mode only supports basic preprocessing with uniform parameters applied + to all datasets. For dataset-specific parameters (e.g., different + smooth_trajectories per dataset) or YAML batch mode, use --user-config. Args: - config_dir: Directory containing YAML scene config files. - desired_data: List of dataset names (first one will be used as env_name). - data_dirs: Dictionary mapping dataset names to data directories. - cache_location: Path to cache directory. - rebuild_cache: Whether to rebuild cache. - rebuild_maps: Whether to rebuild maps. - num_workers: Number of worker processes. - desired_dt: Desired timestep duration in seconds. - num_timesteps_before: Number of timesteps before central token. - num_timesteps_after: Number of timesteps after central token. - verbose: Whether to show verbose output. + args: Parsed command line arguments Returns: - Exit code (0 for success, 1 for failure). + True if successful, False otherwise """ - logger.info("Using YAML config batch preprocessing mode") + # Validate required arguments + if not args.desired_data: + logger.error("--desired-data is required when not using --user-config") + return False + if not args.data_dirs: + logger.error("--data-dir is required when not using --user-config") + return False + if not args.cache_location: + logger.error("--cache-location is required when not using --user-config") + return False + + # Build simple configuration + data_dirs = parse_data_dirs(args.data_dirs, args.desired_data) + + # Validate that all datasets have data directories + missing = set(args.desired_data) - set(data_dirs.keys()) + if missing: + logger.error(f"Missing data directories for datasets: {missing}") + logger.error("Provide --data-dir for each dataset: dataset=/path or in order") + return False + + incl_vector_map = not args.no_vector_map - success = preprocess_from_yaml_configs( - config_dir=Path(config_dir), - cache_location=cache_location, + config = PrepareDataConfig( + desired_data=args.desired_data, data_dirs=data_dirs, - env_name=desired_data[0], # Use first dataset as environment name - rebuild_cache=rebuild_cache, - rebuild_maps=rebuild_maps, - num_workers=num_workers, - desired_dt=desired_dt, - num_timesteps_before=num_timesteps_before, - num_timesteps_after=num_timesteps_after, - verbose=verbose, + cache_location=args.cache_location, + incl_vector_map=incl_vector_map, + rebuild_cache=args.rebuild_cache, + rebuild_maps=args.rebuild_maps, + desired_dt=args.desired_dt, + num_workers=args.num_workers, ) - return 0 if success else 1 + logger.info("Mode: CLI-based basic preprocessing") + return preprocess_basic(config, verbose=args.verbose) def main(arg_list: Optional[List[str]] = None) -> int: @@ -595,103 +523,23 @@ def main(arg_list: Optional[List[str]] = None) -> int: datefmt="%H:%M:%S", ) - verbose = args.verbose and not args.quiet - logger.info("=" * 60) logger.info("Alpasim Data Preparation Tool") logger.info("=" * 60) - # Determine mode and load configuration + # Route to appropriate mode try: if args.user_config: - logger.info(f"Loading configuration from: {args.user_config}") - config = load_config_from_file(args.user_config) - - # Override with command line args if provided - if args.rebuild_cache: - config["rebuild_cache"] = True - if args.rebuild_maps: - config["rebuild_maps"] = True - if args.smooth_trajectories is not None: - smooth_value = args.smooth_trajectories.lower() in ["true", "1", "yes"] - config["smooth_trajectories"] = smooth_value - logger.info( - f"Command line overrides config: smooth_trajectories={smooth_value}" - ) - - # Check if user config contains config_dir -> automatically use Mode 2 - if config.get("config_dir") is not None: - logger.info("Detected 'config_dir' in user config") - return run_yaml_batch_preprocessing( - config_dir=config["config_dir"], - desired_data=config["desired_data"], - data_dirs=config["data_dirs"], - cache_location=config["cache_location"], - rebuild_cache=config.get("rebuild_cache", False), - rebuild_maps=config.get("rebuild_maps", False), - num_workers=config.get("num_workers", 8), - desired_dt=config.get("desired_dt", 0.5), - num_timesteps_before=config.get("num_timesteps_before", 30), - num_timesteps_after=config.get("num_timesteps_after", 80), - verbose=verbose, - ) - # Otherwise continue with Mode 1 (basic mode) - + success = run_from_user_config(args.user_config, args) else: - if not args.desired_data: - logger.error("Either --user-config, or --desired-data must be provided") - return 1 - if not args.data_dirs: - logger.error("--data-dir is required when not using --user-config") - return 1 - if not args.cache_location: - logger.error( - "--cache-location is required when not using --user-config" - ) - return 1 - - data_dirs = parse_data_dirs(args.data_dirs, args.desired_data) - incl_vector_map = not args.no_vector_map - - if args.smooth_trajectories is not None: - smooth_value = args.smooth_trajectories.lower() in ["true", "1", "yes"] - else: - smooth_value = True - - config = { - "desired_data": args.desired_data, - "data_dirs": data_dirs, - "cache_location": args.cache_location, - "incl_vector_map": incl_vector_map, - "rebuild_cache": args.rebuild_cache, - "rebuild_maps": args.rebuild_maps, - "desired_dt": args.desired_dt, - "num_workers": args.num_workers, - "smooth_trajectories": smooth_value, - } - + success = run_from_cli(args) except Exception as e: - logger.error(f"Configuration error: {e}") + logger.error(f"Error during preprocessing: {e}") import traceback traceback.print_exc() return 1 - # Run basic preprocessing - success = preprocess_basic( - desired_data=config["desired_data"], - data_dirs=config["data_dirs"], - cache_location=config["cache_location"], - rebuild_cache=config.get("rebuild_cache", False), - rebuild_maps=config.get("rebuild_maps", False), - incl_vector_map=config.get("incl_vector_map", True), - desired_dt=config.get("desired_dt", 0.1), - num_workers=config.get("num_workers", 8), - verbose=verbose, - list_scenes=args.list_scenes, - smooth_trajectories=config.get("smooth_trajectories", True), - ) - return 0 if success else 1 From 74e7d18e8bd8ac3e8d9ed9fc040c530baa299f19 Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Thu, 26 Mar 2026 09:20:22 +0000 Subject: [PATCH 16/17] refactor: address PR review feedback for trajdata_data_source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves all 10 code quality issues raised in PR review: 1. Decompose map property (330→30 lines, -90%) - Extract helper methods: _transform_map_points, _apply_coordinate_transform_to_map - Split loading paths: _load_map_from_scene_data, _load_map_from_dataset_api - Separate concerns: transform, verify, fix datatypes 2. Reduce defensive hasattr checks (18→1, -94%) - Remove checks for guaranteed attributes based on trajdata source code - Keep only dynamic attribute: scene.map_data - Fix attribute naming: left_boundary/right_boundary → left_edge/right_edge 3. Simplify state accessor pattern - Add _get_state_value() helper to centralize state extraction - Remove unnecessary hasattr and array wrapping/unwrapping - Use StateArray.get_attr() directly (always available) 5. Remove coordinate transformation duplication - Inline simple translations (positions + translation) - Keep map-specific logic in _transform_map_points 6. Remove dead code - Delete from_agent_batch and _load_from_batch (NotImplementedError stub) 7. Move csaps import to top - Remove deferred import from loop body - Make trajdata and csaps required dependencies (remove try-except) 8. Fix non-deterministic metadata - Use SHA256 hash instead of random UUID for dataset_hash - Use fixed training_date instead of datetime.now() 9. Explicit dependency management - Add _ensure_rig_loaded() method with clear naming - Document property loading order in class docstring - Add comments at call sites 10. Simplify scene_id property - Remove unused setter (scene_id should be immutable) - Remove fallback paths (guaranteed at construction) - 12 lines → 3 lines Additional improvements: - Fix circular import: create daemon/exceptions.py module - Fix OmegaConf access: add OmegaConf.to_object() conversions - Fix asset_base_path logic: map by dataset name to avoid conflicts - Remove has_rotation checks: world_to_nre only contains translation - Pass map_api to workers: enable all datasets to load maps lazily - Improve SceneLoader: build asset_base_path_map at init Code quality metrics: - Lines: 1054 → 925 (-12.2%) - hasattr: 18 → 1 (-94%) - map property: 330 → 30 lines (-90%) - All pre-commit checks passing Co-Authored-By: Claude Sonnet 4.5 (1M context) --- .../alpasim_runtime/daemon/__init__.py | 3 + src/runtime/alpasim_runtime/daemon/engine.py | 15 +- .../alpasim_runtime/daemon/exceptions.py | 18 + .../alpasim_runtime/runtime_context.py | 4 +- src/runtime/alpasim_runtime/scene_loader.py | 28 +- src/runtime/tests/test_daemon_main.py | 2 +- .../tests/test_daemon_request_plumbing.py | 6 +- .../tests/test_trajdata_integration.py | 7 +- .../alpasim_utils/trajdata_data_source.py | 746 ++++++++---------- 9 files changed, 361 insertions(+), 468 deletions(-) create mode 100644 src/runtime/alpasim_runtime/daemon/exceptions.py diff --git a/src/runtime/alpasim_runtime/daemon/__init__.py b/src/runtime/alpasim_runtime/daemon/__init__.py index 5d7ab7b6..40324f86 100644 --- a/src/runtime/alpasim_runtime/daemon/__init__.py +++ b/src/runtime/alpasim_runtime/daemon/__init__.py @@ -2,6 +2,7 @@ # Copyright (c) 2026 NVIDIA Corporation from alpasim_runtime.daemon.engine import DaemonEngine +from alpasim_runtime.daemon.exceptions import InvalidRequestError, UnknownSceneError from alpasim_runtime.daemon.request_store import RequestStore from alpasim_runtime.daemon.scheduler import DaemonScheduler, DaemonUnavailableError @@ -10,4 +11,6 @@ "DaemonScheduler", "DaemonUnavailableError", "RequestStore", + "InvalidRequestError", + "UnknownSceneError", ] diff --git a/src/runtime/alpasim_runtime/daemon/engine.py b/src/runtime/alpasim_runtime/daemon/engine.py index 11e05b85..ead01497 100644 --- a/src/runtime/alpasim_runtime/daemon/engine.py +++ b/src/runtime/alpasim_runtime/daemon/engine.py @@ -10,6 +10,7 @@ from alpasim_grpc.v0 import logging_pb2, runtime_pb2 from alpasim_runtime.address_pool import AddressPool +from alpasim_runtime.daemon.exceptions import InvalidRequestError from alpasim_runtime.daemon.scheduler import DaemonScheduler, DaemonUnavailableError from alpasim_runtime.runtime_context import ( build_runtime_context, @@ -100,20 +101,6 @@ def build_simulation_return( ) -class InvalidRequestError(ValueError): - """Raised when a simulation request contains invalid parameters.""" - - pass - - -class UnknownSceneError(InvalidRequestError): - """Raised when a simulation request references a scene_id with no known data source.""" - - def __init__(self, scene_id: str): - super().__init__(f"No data source found for scene_id: {scene_id}") - self.scene_id = scene_id - - def build_pending_jobs_from_request( request: runtime_pb2.SimulationRequest, get_data_source: Callable[[str], SceneDataSource], diff --git a/src/runtime/alpasim_runtime/daemon/exceptions.py b/src/runtime/alpasim_runtime/daemon/exceptions.py new file mode 100644 index 00000000..2beaed16 --- /dev/null +++ b/src/runtime/alpasim_runtime/daemon/exceptions.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA Corporation + +"""Exception classes for runtime daemon.""" + + +class InvalidRequestError(ValueError): + """Raised when a simulation request contains invalid parameters.""" + + pass + + +class UnknownSceneError(InvalidRequestError): + """Raised when a simulation request references a scene_id with no known data source.""" + + def __init__(self, scene_id: str): + super().__init__(f"No data source found for scene_id: {scene_id}") + self.scene_id = scene_id diff --git a/src/runtime/alpasim_runtime/runtime_context.py b/src/runtime/alpasim_runtime/runtime_context.py index 8ecddd6e..589199a4 100644 --- a/src/runtime/alpasim_runtime/runtime_context.py +++ b/src/runtime/alpasim_runtime/runtime_context.py @@ -20,6 +20,7 @@ gather_versions_from_addresses, validate_scenarios, ) +from omegaconf import OmegaConf from trajdata.dataset import UnifiedDataset from eval.schema import EvalConfig @@ -190,8 +191,7 @@ async def build_runtime_context( # Create UnifiedDataset and build scene_id to data source mapping logger.info("Creating UnifiedDataset from config") - - data_source_config = config.user.data_source + data_source_config = OmegaConf.to_object(config.user.data_source) trajdata_params = data_source_config.to_trajdata_params() dataset = UnifiedDataset(**trajdata_params) logger.info( diff --git a/src/runtime/alpasim_runtime/scene_loader.py b/src/runtime/alpasim_runtime/scene_loader.py index 1893a6cc..4c038f71 100644 --- a/src/runtime/alpasim_runtime/scene_loader.py +++ b/src/runtime/alpasim_runtime/scene_loader.py @@ -8,9 +8,10 @@ import logging from alpasim_runtime.config import SimulatorConfig -from alpasim_runtime.daemon.scheduler import UnknownSceneError +from alpasim_runtime.daemon.exceptions import UnknownSceneError from alpasim_utils.scene_data_source import SceneDataSource from alpasim_utils.trajdata_data_source import TrajdataDataSource +from omegaconf import OmegaConf from trajdata.dataset import UnifiedDataset logger = logging.getLogger(__name__) @@ -28,6 +29,7 @@ class SceneLoader: _scene_id_to_idx: Mapping from scene IDs to dataset indices _config: Simulator configuration for scene parameters _cache: Cache of loaded SceneDataSource objects + _asset_base_path_map: Mapping from dataset name to asset_base_path """ def __init__( @@ -48,6 +50,19 @@ def __init__( self._config = config self._cache: dict[str, SceneDataSource] = {} + # Build dataset_name -> asset_base_path mapping + self._asset_base_path_map: dict[str, str] = {} + data_source_config = OmegaConf.to_object(config.user.data_source) + for dataset_name, source in data_source_config.sources.items(): + if "asset_base_path" in source.extra_params: + self._asset_base_path_map[dataset_name] = source.extra_params[ + "asset_base_path" + ] + logger.info( + f"Registered asset_base_path for {dataset_name}: " + f"{source.extra_params['asset_base_path']}" + ) + def get_data_source(self, scene_id: str) -> SceneDataSource: """Get or create a data source for the given scene_id. @@ -80,10 +95,12 @@ def get_data_source(self, scene_id: str) -> SceneDataSource: if scene is None: raise UnknownSceneError(scene_id) - # Get asset_base_path from USDZ config if available - asset_base_path = None - if self._config.user.data_source.usdz is not None: - asset_base_path = self._config.user.data_source.usdz.asset_base_path + # Get asset_base_path for this scene's dataset + # Use scene.env_name to lookup the correct asset_base_path + asset_base_path = self._asset_base_path_map.get(scene.env_name) + + # Get map_api for lazy loading (lightweight, can be passed to workers) + map_api = getattr(self._dataset, "_map_api", None) # Create scene_cache (pre-create to avoid pickle errors) scene_cache = self._dataset.cache_class( @@ -95,6 +112,7 @@ def get_data_source(self, scene_id: str) -> SceneDataSource: data_source = TrajdataDataSource.from_trajdata_scene( scene=scene, dataset=None, # Don't pass dataset to avoid pickle errors + map_api=map_api, # Pass map_api scene_cache=scene_cache, smooth_trajectories=self._config.user.smooth_trajectories, asset_base_path=asset_base_path, diff --git a/src/runtime/tests/test_daemon_main.py b/src/runtime/tests/test_daemon_main.py index 8923b463..53139618 100644 --- a/src/runtime/tests/test_daemon_main.py +++ b/src/runtime/tests/test_daemon_main.py @@ -11,7 +11,7 @@ import pytest from alpasim_grpc.v0 import common_pb2, runtime_pb2 from alpasim_runtime.daemon.app import RuntimeDaemonApp -from alpasim_runtime.daemon.engine import UnknownSceneError +from alpasim_runtime.daemon.exceptions import UnknownSceneError from alpasim_runtime.daemon.servicer import RuntimeDaemonServicer from alpasim_runtime.simulate.__main__ import _serve, create_arg_parser, run_simulation diff --git a/src/runtime/tests/test_daemon_request_plumbing.py b/src/runtime/tests/test_daemon_request_plumbing.py index 78ecba01..88a0d10c 100644 --- a/src/runtime/tests/test_daemon_request_plumbing.py +++ b/src/runtime/tests/test_daemon_request_plumbing.py @@ -7,10 +7,8 @@ import pytest from alpasim_grpc.v0 import runtime_pb2 -from alpasim_runtime.daemon.engine import ( - UnknownSceneError, - build_pending_jobs_from_request, -) +from alpasim_runtime.daemon.engine import build_pending_jobs_from_request +from alpasim_runtime.daemon.exceptions import UnknownSceneError def test_adapter_expands_nr_rollouts() -> None: diff --git a/src/runtime/tests/test_trajdata_integration.py b/src/runtime/tests/test_trajdata_integration.py index 77185f18..3a547d50 100644 --- a/src/runtime/tests/test_trajdata_integration.py +++ b/src/runtime/tests/test_trajdata_integration.py @@ -9,11 +9,8 @@ import pytest from alpasim_grpc.v0 import runtime_pb2 -from alpasim_runtime.daemon.engine import ( - DaemonEngine, - UnknownSceneError, - build_pending_jobs_from_request, -) +from alpasim_runtime.daemon.engine import DaemonEngine, build_pending_jobs_from_request +from alpasim_runtime.daemon.exceptions import UnknownSceneError @pytest.fixture diff --git a/src/utils/alpasim_utils/trajdata_data_source.py b/src/utils/alpasim_utils/trajdata_data_source.py index 5cc31216..fac76fa2 100644 --- a/src/utils/alpasim_utils/trajdata_data_source.py +++ b/src/utils/alpasim_utils/trajdata_data_source.py @@ -32,29 +32,14 @@ from __future__ import annotations import copy +import hashlib import logging import os from dataclasses import dataclass from typing import Optional +import csaps import numpy as np -from scipy.spatial.transform import Rotation as R - -try: - from trajdata import AgentBatch - from trajdata.caching import EnvCache - from trajdata.data_structures.agent import AgentMetadata - from trajdata.data_structures.scene_metadata import Scene - from trajdata.dataset import UnifiedDataset - from trajdata.maps import VectorMap -except ImportError: - AgentBatch = None - Scene = None - UnifiedDataset = None - EnvCache = None - AgentMetadata = None - VectorMap = None - from alpasim_utils.artifact import Metadata from alpasim_utils.geometry import Trajectory from alpasim_utils.scenario import ( @@ -66,6 +51,12 @@ VehicleConfig, ) from alpasim_utils.scene_data_source import SceneDataSource +from scipy.spatial.transform import Rotation as R +from trajdata.caching import EnvCache +from trajdata.data_structures.agent import AgentMetadata +from trajdata.data_structures.scene_metadata import Scene +from trajdata.dataset import UnifiedDataset +from trajdata.maps import VectorMap logger = logging.getLogger(__name__) @@ -77,11 +68,22 @@ class TrajdataDataSource(SceneDataSource): This class implements the SceneDataSource protocol, allowing direct loading from trajdata Scene or AgentBatch objects without requiring USDZ format. + + Property loading order dependencies: + - rig: No dependencies (loads first, sets world_to_nre transformation) + - traffic_objects: Requires rig (uses world_to_nre for coordinate conversion) + - map: Requires rig (uses world_to_nre for coordinate conversion) + - metadata: Requires rig (uses trajectory time range) + + All properties use lazy loading and caching for efficiency. """ _scene: Scene | None = None _scene_cache: EnvCache | None = None _dataset: UnifiedDataset | None = None + _map_api: Optional[object] = ( + None # MapAPI for loading maps (lightweight, can be passed separately) + ) _rig: Rig | None = None _traffic_objects: TrafficObjects | None = None _map: VectorMap | None = None @@ -95,6 +97,7 @@ def from_trajdata_scene( cls, scene: Scene, dataset: Optional[UnifiedDataset] = None, + map_api=None, scene_cache: Optional[EnvCache] = None, scene_id: Optional[str] = None, smooth_trajectories: bool = True, @@ -107,6 +110,7 @@ def from_trajdata_scene( Args: scene: trajdata Scene object dataset: UnifiedDataset instance (for getting scene_cache and map) + map_api: MapAPI instance for loading maps (lightweight alternative to dataset) scene_cache: Optional EnvCache instance (if not provided, will be created from dataset) scene_id: Optional scene ID (if not provided, uses scene.name) smooth_trajectories: Whether to smooth trajectories @@ -116,14 +120,10 @@ def from_trajdata_scene( Returns: TrajdataDataSource instance """ - if Scene is None: - raise ImportError( - "trajdata is not installed. Please install it to use TrajdataDataSource." - ) - data_source = cls( _scene=scene, _dataset=dataset, + _map_api=map_api, _scene_cache=scene_cache, _scene_id=scene_id or scene.name, _smooth_trajectories=smooth_trajectories, @@ -132,59 +132,10 @@ def from_trajdata_scene( data_source._base_timestamp_us = base_timestamp_us return data_source - @classmethod - def from_agent_batch( - cls, - batch: AgentBatch, - scene_id: str, - smooth_trajectories: bool = True, - ) -> TrajdataDataSource: - """ - Create TrajdataDataSource from trajdata AgentBatch object. - - Note: This method requires the batch to contain complete scene information. - It is generally recommended to use from_trajdata_scene instead. - - Args: - batch: trajdata AgentBatch object - scene_id: Scene ID - smooth_trajectories: Whether to smooth trajectories - - Returns: - TrajdataDataSource instance - """ - if AgentBatch is None: - raise ImportError( - "trajdata is not installed. Please install it to use TrajdataDataSource." - ) - - data_source = cls(scene_id=scene_id) - data_source._smooth_trajectories = smooth_trajectories - # Extract data from batch - data_source._load_from_batch(batch) - return data_source - - def _load_from_batch(self, batch: AgentBatch) -> None: - """Load data from AgentBatch (internal method)""" - # Need to extract data based on batch structure - # Specific implementation depends on your trajdata data format - raise NotImplementedError( - "from_agent_batch needs to be implemented based on your trajdata data format. " - "It is recommended to use from_trajdata_scene method instead." - ) - @property def scene_id(self) -> str: - """Scene ID""" - if self._scene_id: - return self._scene_id - if self._scene is not None: - return self._scene.name - raise ValueError("scene_id is not set and cannot be obtained from scene") - - @scene_id.setter - def scene_id(self, value: str) -> None: - self._scene_id = value + """Scene ID (immutable identifier).""" + return self._scene_id @property def asset_path(self) -> str | None: @@ -223,8 +174,8 @@ def _extract_asset_folder_name(self) -> str: Asset folder name (defaults to scene_id if no specific name found) """ # Try to get from scene data_access_info - if self._scene is not None and hasattr(self._scene, "data_access_info"): - data_access_info = self._scene.data_access_info or {} + if self._scene is not None and self._scene.data_access_info is not None: + data_access_info = self._scene.data_access_info # USDZ: Use usdz_stem (filename without .usdz extension) if "usdz_stem" in data_access_info: @@ -279,6 +230,47 @@ def _get_scene_cache(self) -> EnvCache: "or TrajdataDataSource.from_trajdata_scene(scene, scene_cache=your_cache)" ) + @staticmethod + def _get_state_value(state, attr_name: str, default=None): + """Extract scalar value from state attribute using StateArray.get_attr(). + + Args: + state: StateArray object from trajdata cache (from get_raw_state) + attr_name: Name of attribute to extract (e.g., "x", "y", "h") + default: Default value if attribute doesn't exist + + Returns: + Scalar float value + + Raises: + AttributeError: If attribute doesn't exist and no default provided + """ + try: + value = state.get_attr(attr_name) + + if value is None: + if default is not None: + return default + raise AttributeError(f"Attribute {attr_name} is None") + + # Convert to scalar if needed + if isinstance(value, np.ndarray): + return float(value.flat[0]) + return float(value) + except (KeyError, AttributeError, TypeError, IndexError) as e: + if default is not None: + return default + raise AttributeError(f"Failed to get {attr_name} from state: {e}") + + def _ensure_rig_loaded(self) -> None: + """Ensure rig is loaded before accessing world_to_nre. + + This method is called by properties that depend on world_to_nre + transformation (traffic_objects, map helper methods). + """ + if self._rig is None: + _ = self.rig + def _extract_agent_trajectory( self, agent: AgentMetadata, @@ -301,33 +293,11 @@ def _extract_agent_trajectory( try: state = scene_cache.get_raw_state(agent.name, ts) - # Get position and orientation - x = state.get_attr("x") if hasattr(state, "get_attr") else state.x - y = state.get_attr("y") if hasattr(state, "get_attr") else state.y - z = ( - state.get_attr("z") - if hasattr(state, "get_attr") - else (state.z if hasattr(state, "z") else 0.0) - ) - heading = ( - state.get_attr("h") if hasattr(state, "get_attr") else state.h - ) - - # Convert to numpy array (handle scalar case) - if isinstance(x, (int, float)): - x = np.array([x]) - if isinstance(y, (int, float)): - y = np.array([y]) - if isinstance(z, (int, float)): - z = np.array([z]) - if isinstance(heading, (int, float)): - heading = np.array([heading]) - - # Take first element (if array) - x_val = float(x[0] if x.ndim > 0 else x) - y_val = float(y[0] if y.ndim > 0 else y) - z_val = float(z[0] if z.ndim > 0 else z) - heading_val = float(heading[0] if heading.ndim > 0 else heading) + # Get position and orientation using helper + x_val = self._get_state_value(state, "x") + y_val = self._get_state_value(state, "y") + z_val = self._get_state_value(state, "z", default=0.0) + heading_val = self._get_state_value(state, "h") # Calculate timestamp if base_timestamp_us is None: @@ -359,16 +329,14 @@ def _extract_agent_trajectory( ) # Create VehicleConfig (extract from extent) - vehicle_config = None - if hasattr(agent.extent, "length"): - vehicle_config = VehicleConfig( - aabb_x_m=agent.extent.length, - aabb_y_m=agent.extent.width, - aabb_z_m=agent.extent.height, - aabb_x_offset_m=-agent.extent.length / 2, - aabb_y_offset_m=0.0, - aabb_z_offset_m=-agent.extent.height / 2, - ) + vehicle_config = VehicleConfig( + aabb_x_m=agent.extent.length, + aabb_y_m=agent.extent.width, + aabb_z_m=agent.extent.height, + aabb_x_offset_m=-agent.extent.length / 2, + aabb_y_offset_m=0.0, + aabb_z_offset_m=-agent.extent.height / 2, + ) return trajectory, vehicle_config @@ -468,12 +436,9 @@ def _extract_camera_info_from_scene(self) -> tuple[list[CameraId], dict]: return camera_ids, camera_calibrations # Check if sensor_calibration information exists - if ( - not hasattr(self._scene, "data_access_info") - or not self._scene.data_access_info - ): + if not self._scene.data_access_info: logger.warning( - "scene.data_access_info does not exist, skipping camera information extraction" + "scene.data_access_info is empty, skipping camera information extraction" ) return camera_ids, camera_calibrations @@ -569,6 +534,41 @@ def _is_static_object( avg_velocity = np.mean(velocities) return avg_velocity < velocity_threshold + def _transform_map_points( + self, + points: np.ndarray, + translation_xy: np.ndarray, + first_traj_z: float, + ) -> np.ndarray: + """Transform map points to local coordinates (translation only). + + Args: + points: Nx3 array of map points + translation_xy: 2-element XY translation + first_traj_z: Z coordinate of first trajectory point + + Returns: + Transformed points + """ + if ( + points is None + or len(points) == 0 + or points.ndim != 2 + or points.shape[1] < 3 + ): + return points + + points_copy = points.copy() + + # Apply XY translation + points_copy[:, 0] = points_copy[:, 0] + translation_xy[0] + points_copy[:, 1] = points_copy[:, 1] + translation_xy[1] + + # Align Z coordinate to trajectory baseline + points_copy[:, 2] = points_copy[:, 2] + first_traj_z + + return points_copy + @property def traffic_objects(self) -> TrafficObjects: """Load and return traffic objects""" @@ -600,9 +600,8 @@ def traffic_objects(self) -> TrafficObjects: continue # Convert trajectory to local coordinates (NRE) - use rig's world_to_nre - if self._rig is None: - # If rig is not loaded yet, load it first - _ = self.rig + # Explicit dependency: need world_to_nre from rig + self._ensure_rig_loaded() world_to_nre = self._rig.world_to_nre translation = world_to_nre[:3, 3] @@ -617,8 +616,6 @@ def traffic_objects(self) -> TrafficObjects: # Smooth if needed if self._smooth_trajectories: try: - import csaps - css = csaps.CubicSmoothingSpline( trajectory.timestamps_us / 1e6, trajectory.positions.T, @@ -638,23 +635,19 @@ def traffic_objects(self) -> TrafficObjects: positions=filtered_positions.astype(np.float32), quaternions=trajectory.quaternions.copy(), ) - except ImportError: - logger.warning("csaps not installed, skipping trajectory smoothing") + except Exception as e: + logger.warning(f"Failed to smooth trajectory: {e}") # Get AABB - if hasattr(agent.extent, "length"): - aabb = AABB( - x=agent.extent.length, y=agent.extent.width, z=agent.extent.height - ) - else: - # Default AABB - aabb = AABB(x=4.5, y=1.8, z=1.5) + aabb = AABB( + x=agent.extent.length, y=agent.extent.width, z=agent.extent.height + ) # Determine if static object is_static = self._is_static_object(trajectory) # Get category label - label_class = agent.type.name if hasattr(agent.type, "name") else "UNKNOWN" + label_class = getattr(agent.type, "name", "UNKNOWN") traffic_dict[agent.name] = TrafficObject( track_id=agent.name, @@ -667,337 +660,210 @@ def traffic_objects(self) -> TrafficObjects: self._traffic_objects = TrafficObjects(**traffic_dict) return self._traffic_objects - @property - def map(self) -> Optional[VectorMap]: - """Load and return VectorMap (obtained from dataset._map_api or scene.map_data)""" - if self._map is not None: - return self._map + def _apply_coordinate_transform_to_map(self, vec_map: VectorMap) -> None: + """Apply world_to_nre coordinate transformation to map in-place. - if VectorMap is None: - logger.warning("trajdata is not installed, cannot load map") - return None + Note: world_to_nre only contains translation (no rotation), as the local + coordinate frame maintains the same orientation as the world frame (ENU). - if self._scene is None: - logger.warning("Cannot load map: scene is not set") - return None + Args: + vec_map: VectorMap to transform + """ + world_to_nre = self._rig.world_to_nre + translation = world_to_nre[:3, 3] + translation_xy = translation[:2] + first_traj_z = ( + self.rig.trajectory.positions[0][2] if len(self.rig.trajectory) > 0 else 0.0 + ) - # First, try to get map from scene.map_data (for USDZ and other datasets that attach map directly) - if hasattr(self._scene, "map_data") and self._scene.map_data is not None: - logger.info(f"Loading map from scene.map_data for {self.scene_id}") - # Make a deep copy to avoid modifying shared map object - self._map = copy.deepcopy(self._scene.map_data) + logger.info( + f"Map coordinate transformation: " + f"translation_xy={translation_xy}, " + f"first_traj_z={first_traj_z:.2f}m" + ) - # Apply coordinate transformation if needed - if self._rig is None: - # If rig is not loaded yet, load it first (this will set world_to_nre) - _ = self.rig + # Transform all lane points + if vec_map.lanes is None: + return - world_to_nre = self._rig.world_to_nre + for lane in vec_map.lanes: + # Transform center (always exists) + lane.center.points = self._transform_map_points( + lane.center.points, + translation_xy, + first_traj_z, + ) - # Check if transformation is needed (if world_to_nre is not identity) - if world_to_nre is not None and not np.allclose(world_to_nre, np.eye(4)): - translation = world_to_nre[:3, 3] - logger.info( - f"Transforming map to local coordinates with translation: {translation}" + # Transform left_edge (optional) + if lane.left_edge is not None and lane.left_edge.points is not None: + lane.left_edge.points = self._transform_map_points( + lane.left_edge.points, + translation_xy, + first_traj_z, ) - # Transform map coordinates - self._map.translate(translation[0], translation[1], translation[2]) + # Transform right_edge (optional) + if lane.right_edge is not None and lane.right_edge.points is not None: + lane.right_edge.points = self._transform_map_points( + lane.right_edge.points, + translation_xy, + first_traj_z, + ) - logger.info( - f"Successfully loaded map from scene.map_data for {self.scene_id}" - ) - return self._map + def _fix_map_datatypes(self, vec_map: VectorMap) -> None: + """Fix lane connectivity data types (convert lists to sets). + + This is needed because some trajdata loaders may incorrectly create + these as lists instead of sets. + + Args: + vec_map: VectorMap to fix + """ + if vec_map.lanes is None: + return + + for lane in vec_map.lanes: + # Convert to set if they are lists (defensive, but based on observed issues) + if isinstance(lane.next_lanes, list): + lane.next_lanes = set(lane.next_lanes) + if isinstance(lane.prev_lanes, list): + lane.prev_lanes = set(lane.prev_lanes) + if isinstance(lane.adj_lanes_right, list): + lane.adj_lanes_right = set(lane.adj_lanes_right) + if isinstance(lane.adj_lanes_left, list): + lane.adj_lanes_left = set(lane.adj_lanes_left) + + def _verify_map_transformation(self, vec_map: VectorMap) -> None: + """Verify map coordinate transformation is correct. + + Args: + vec_map: Transformed VectorMap + """ + if vec_map.lanes is None or len(vec_map.lanes) == 0: + return + + # Get first lane and points (center always exists in RoadLane) + first_lane = vec_map.lanes[0] + first_map_point = first_lane.center.points[0, :3] + first_traj_point = self.rig.trajectory.positions[0] + + distance_xy = np.linalg.norm(first_map_point[:2]) + z_diff = abs(first_map_point[2] - first_traj_point[2]) + + logger.info( + f"Map transformation verification: " + f"first lane center: {first_map_point}, " + f"first trajectory: {first_traj_point}, " + f"XY distance: {distance_xy:.2f}m, " + f"Z difference: {z_diff:.2f}m" + ) - # Otherwise, try to get map from dataset._map_api (for datasets with map_api) - if self._dataset is None: + if z_diff > 10.0: logger.warning( - "Cannot load map: dataset is not set and scene.map_data is not available" + f"Map Z coordinate may not be correctly aligned. " + f"Map Z={first_map_point[2]:.2f}m, Traj Z={first_traj_point[2]:.2f}m" ) + + def _load_map_from_scene_data(self) -> Optional[VectorMap]: + """Load map from scene.map_data (USDZ). + + Returns: + VectorMap if available, None otherwise + """ + if not hasattr(self._scene, "map_data") or self._scene.map_data is None: return None - # Get map from dataset._map_api (refer to trajdata_artifact_converter.py) - try: - # Check if dataset includes map support - if ( - not hasattr(self._dataset, "incl_vector_map") - or not self._dataset.incl_vector_map - or not hasattr(self._dataset, "_map_api") - or self._dataset._map_api is None - ): - logger.warning( - "Dataset does not have map support enabled or map_api is unavailable" - ) - return None + logger.info(f"Loading map from scene.map_data for {self.scene_id}") + vec_map = copy.deepcopy(self._scene.map_data) - # Build map name: "{env_name}:{location}" - if not hasattr(self._scene, "location") or not self._scene.location: - logger.warning( - f"Scene {self.scene_id} has no location information, cannot load map" - ) - return None + # Ensure rig is loaded (need world_to_nre for transformation) + self._ensure_rig_loaded() - map_name = f"{self._scene.env_name}:{self._scene.location}" + # Apply coordinate transformation + self._apply_coordinate_transform_to_map(vec_map) - # Get vector_map_params (if exists) - vector_map_params = {} - if hasattr(self._dataset, "vector_map_params"): - vector_map_params = self._dataset.vector_map_params + # Fix datatypes and verify + self._fix_map_datatypes(vec_map) + self._verify_map_transformation(vec_map) - # Get map from map_api - vec_map = self._dataset._map_api.get_map(map_name, **vector_map_params) + logger.info("Successfully loaded map from scene.map_data") + return vec_map - if vec_map is None: - logger.debug( - f"Scene {self.scene_id} (map_name: {map_name}) has no map data" - ) - return None + def _load_map_from_dataset_api(self) -> Optional[VectorMap]: + """Load map from dataset._map_api (datasets with map API). - # Create an independent copy of VectorMap for current scene to avoid modifying - # map objects in shared cache. This allows continued use of MapAPI's disk cache - # and index loading capabilities while preventing coordinate system pollution - # between multiple scenes through shared VectorMap instances. - self._map = copy.deepcopy(vec_map) + Returns: + VectorMap if available, None otherwise + """ + # Try to get map_api from either self._map_api or self._dataset + map_api = self._map_api + if map_api is None and self._dataset is not None: + map_api = getattr(self._dataset, "_map_api", None) - # Important: Transform map to local coordinate system (NRE) - # Since trajectories are already converted to local coordinates, map also needs - # to be converted to match. This is consistent with USDZ format handling: - # both map and trajectories need to be converted to the same coordinate system - if self._rig is None: - # If rig is not loaded yet, load it first (this will set world_to_nre) - _ = self.rig + if map_api is None: + logger.warning("Cannot load map: map_api not available") + return None - world_to_nre = self._rig.world_to_nre + # Get vector_map_params from dataset if available + vector_map_params = {} + if self._dataset is not None: + vector_map_params = getattr(self._dataset, "vector_map_params", {}) - # Check if world_to_nre contains rotation - rotation_matrix = world_to_nre[:3, :3] - translation = world_to_nre[:3, 3] - has_rotation = not np.allclose(rotation_matrix, np.eye(3)) + # Build map name + if not self._scene.location: + logger.warning(f"Scene {self.scene_id} has no location, cannot load map") + return None - if has_rotation: - logger.warning( - "world_to_nre contains rotation. Map transformation may need rotation handling. " - "Currently only applying translation." - ) + map_name = f"{self._scene.env_name}:{self._scene.location}" - # Transform all points in map (center.points, left_boundary.points, right_boundary.points, etc.) - # Note: Must transform before finalize, as finalize may rebuild certain data structures - # Important: Only transform X, Y coordinates, align Z coordinate to trajectory Z baseline - # Because map Z usually represents height relative to ground (usually 0), while trajectory Z - # represents altitude. After transformation, map Z should align with trajectory Z - # (both relative to first trajectory point's Z, usually 0 after transformation) - translation_xy = translation[:2] # Only use X, Y translation - - # Get Z coordinate of first trajectory point (transformed baseline, usually 0) - first_traj_z = ( - self.rig.trajectory.positions[0][2] - if len(self.rig.trajectory) > 0 - else 0.0 - ) + try: + vec_map = map_api.get_map(map_name, **vector_map_params) + if vec_map is None: + logger.debug(f"Scene {self.scene_id} (map: {map_name}) has no map data") + return None - logger.info( - f"Map coordinate transformation: " - f"translation_xy={translation_xy}, " - f"first_traj_z={first_traj_z:.2f}m, " - f"map Z will be aligned to trajectory Z baseline" - ) + # Deep copy to avoid modifying shared cache + vec_map = copy.deepcopy(vec_map) - def transform_map_points(points: np.ndarray) -> np.ndarray: - """Transform map points: only transform X, Y, align Z to trajectory Z baseline""" - if ( - points is None - or len(points) == 0 - or points.ndim != 2 - or points.shape[1] < 3 - ): - return points - - points_copy = points.copy() - - # Transform X, Y coordinates - if has_rotation: - # If rotation exists, need to rotate X, Y - xy_rotated = ( - points_copy[:, :2] @ rotation_matrix[:2, :2].T - ) + translation_xy - points_copy[:, 0] = xy_rotated[:, 0] - points_copy[:, 1] = xy_rotated[:, 1] - else: - # Only translate X, Y - points_copy[:, 0] = points_copy[:, 0] + translation_xy[0] - points_copy[:, 1] = points_copy[:, 1] + translation_xy[1] - - # Z coordinate alignment: align map Z to trajectory Z baseline - # Map Z is usually height relative to ground (usually 0), - # after transformation should align with trajectory Z (both relative to first - # trajectory point's Z, usually 0 after transformation) - # So: new_z = original_z + first_traj_z - # If map Z=0 (ground), after transformation it becomes first_traj_z (usually 0) - points_copy[:, 2] = points_copy[:, 2] + first_traj_z - - return points_copy - - if hasattr(self._map, "lanes"): - for lane_idx, lane in enumerate(self._map.lanes): - # Transform center.points - if hasattr(lane, "center") and hasattr(lane.center, "points"): - points = lane.center.points - if points is not None and len(points) > 0: - try: - transformed_points = transform_map_points(points) - lane.center.points = transformed_points - except Exception as e: - logger.warning( - f"Failed to transform lane {lane_idx} center.points: {e}" - ) - - # Transform left_boundary.points - if hasattr(lane, "left_boundary") and hasattr( - lane.left_boundary, "points" - ): - points = lane.left_boundary.points - if points is not None and len(points) > 0: - try: - transformed_points = transform_map_points(points) - lane.left_boundary.points = transformed_points - except Exception as e: - logger.warning( - f"Failed to transform lane {lane_idx} left_boundary.points: {e}" - ) - - # Transform right_boundary.points - if hasattr(lane, "right_boundary") and hasattr( - lane.right_boundary, "points" - ): - points = lane.right_boundary.points - if points is not None and len(points) > 0: - try: - transformed_points = transform_map_points(points) - lane.right_boundary.points = transformed_points - except Exception as e: - logger.warning( - f"Failed to transform lane {lane_idx} right_boundary.points: {e}" - ) - - # Transform other possible point attributes (if any) - for attr_name in ["intersections", "crosswalks", "stop_lines"]: - if hasattr(lane, attr_name): - attr_value = getattr(lane, attr_name) - if isinstance(attr_value, list): - for item in attr_value: - if hasattr(item, "points"): - points = item.points - if points is not None and len(points) > 0: - try: - transformed_points = ( - transform_map_points(points) - ) - item.points = transformed_points - except Exception as e: - logger.debug( - f"Failed to transform {attr_name} points: {e}" - ) - - # If map needs finalize, call it (after transformation) - # Note: finalize may rebuild certain indices but won't change point coordinates - # But for safety, we verify transformation again after finalize - if hasattr(self._map, "__post_init__"): - self._map.__post_init__() - if hasattr(self._map, "compute_search_indices"): - self._map.compute_search_indices() - - # Verify again: check if first point's coordinates are still correct after finalize - if hasattr(self._map, "lanes") and len(self._map.lanes) > 0: - first_lane = self._map.lanes[0] - if hasattr(first_lane, "center") and hasattr( - first_lane.center, "points" - ): - first_map_point_after_finalize = ( - first_lane.center.points[0, :3] - if len(first_lane.center.points) > 0 - else None - ) - if first_map_point_after_finalize is not None: - # Check if Z coordinate aligns with trajectory (should both be first_traj_z, usually 0) - actual_z = first_map_point_after_finalize[2] - if abs(actual_z - first_traj_z) > 1.0: # Allow 1 meter error - logger.warning( - f"Map Z coordinate may have been reset after finalize. " - f"Expected Z≈{first_traj_z:.2f}m (aligned to trajectory Z baseline), " - f"got Z={actual_z:.2f}m. " - f"This may cause coordinate misalignment." - ) - # If alignment is correct, no need to log (reduce noise) - - # Fix data types (if needed) - if hasattr(self._map, "lanes"): - for lane in self._map.lanes: - if hasattr(lane, "next_lanes") and isinstance( - lane.next_lanes, list - ): - lane.next_lanes = set(lane.next_lanes) - if hasattr(lane, "prev_lanes") and isinstance( - lane.prev_lanes, list - ): - lane.prev_lanes = set(lane.prev_lanes) - if hasattr(lane, "adj_lanes_right") and isinstance( - lane.adj_lanes_right, list - ): - lane.adj_lanes_right = set(lane.adj_lanes_right) - if hasattr(lane, "adj_lanes_left") and isinstance( - lane.adj_lanes_left, list - ): - lane.adj_lanes_left = set(lane.adj_lanes_left) - - # Verify map transformation: check if first lane's first point is within reasonable range - if hasattr(self._map, "lanes") and len(self._map.lanes) > 0: - first_lane = self._map.lanes[0] - if hasattr(first_lane, "center") and hasattr( - first_lane.center, "points" - ): - first_map_point = ( - first_lane.center.points[0, :3] - if len(first_lane.center.points) > 0 - else None - ) - if first_map_point is not None: - # Map point should be near trajectory (within hundreds of meters) - distance_from_origin_xy = np.linalg.norm( - first_map_point[:2] - ) # Only check X, Y - distance_from_origin_xyz = np.linalg.norm( - first_map_point - ) # Check X, Y, Z - - # Get first trajectory point for comparison - first_traj_point = self.rig.trajectory.positions[0] - - logger.info( - f"Map transformation verification: " - f"first lane center point: {first_map_point}, " - f"first trajectory point: {first_traj_point}, " - f"distance (X,Y): {distance_from_origin_xy:.2f}m, " - f"distance (X,Y,Z): {distance_from_origin_xyz:.2f}m, " - f"Z difference: {abs(first_map_point[2] - first_traj_point[2]):.2f}m" - ) + # Ensure rig is loaded (need world_to_nre for transformation) + self._ensure_rig_loaded() - # Warn if Z coordinate is too far from trajectory Z - z_diff = abs(first_map_point[2] - first_traj_point[2]) - if z_diff > 10.0: # Z coordinate difference exceeds 10 meters - logger.warning( - f"Map Z coordinate may not be correctly aligned with trajectory. " - f"Map Z={first_map_point[2]:.2f}m, Trajectory Z={first_traj_point[2]:.2f}m, " - f"difference={z_diff:.2f}m. This may cause route generation to fail." - ) + # Apply coordinate transformation + self._apply_coordinate_transform_to_map(vec_map) - logger.info( - f"Successfully loaded map: {map_name} (transformed to local coordinate system)" - ) - return self._map + # Finalize map + vec_map.__post_init__() + vec_map.compute_search_indices() + + # Fix datatypes and verify + self._fix_map_datatypes(vec_map) + self._verify_map_transformation(vec_map) + + logger.info(f"Successfully loaded map: {map_name}") + return vec_map except Exception as e: - logger.error(f"Error loading map: {e}", exc_info=True) + logger.error(f"Error loading map from dataset API: {e}", exc_info=True) return None + @property + def map(self) -> Optional[VectorMap]: + """Load and return VectorMap.""" + if self._map is not None: + return self._map + + if self._scene is None: + logger.warning("Cannot load map: scene is not set") + return None + + # Try scene.map_data first (simpler path) + self._map = self._load_map_from_scene_data() + if self._map is not None: + return self._map + + # Fallback to dataset._map_api + self._map = self._load_map_from_dataset_api() + return self._map + @property def metadata(self) -> Metadata: """Create and return Metadata object""" @@ -1028,16 +894,22 @@ def metadata(self) -> Metadata: time_range_start = float(rig.trajectory.time_range_us.start) / 1e6 time_range_end = float(rig.trajectory.time_range_us.stop) / 1e6 - # Create metadata - import uuid - from datetime import datetime + # Generate deterministic IDs based on scene identifiers + # Create hash from scene_id and time range for reproducibility + hash_input = f"{scene_id}_{time_range_start}_{time_range_end}" + dataset_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:16] + uuid_str = hashlib.sha256(f"{hash_input}_uuid".encode()).hexdigest()[:32] + # Use fixed training date instead of datetime.now() for determinism + training_date = "trajdata-generated" + + # Create metadata self._metadata = Metadata( scene_id=scene_id, version_string="trajdata_direct", - training_date=datetime.now().strftime("%Y-%m-%d"), - dataset_hash=str(uuid.uuid4()), - uuid=str(uuid.uuid4()), + training_date=training_date, + dataset_hash=dataset_hash, + uuid=uuid_str, is_resumable=False, sensors=Metadata.Sensors( camera_ids=camera_id_names, From 3dbe2c84c0baeac270daf890e2b7151aa826640d Mon Sep 17 00:00:00 2001 From: Caojun Wang Date: Fri, 24 Apr 2026 12:16:34 +0000 Subject: [PATCH 17/17] fix: align data_source config structure between YAML and DataSourceConfig schema --- src/runtime/alpasim_runtime/config.py | 6 +- .../alpasim_runtime/runtime_context.py | 5 +- src/runtime/alpasim_runtime/scene_loader.py | 11 ++-- src/runtime/tests/test_config.py | 52 ++++++++------- src/wizard/configs/base_config.yaml | 63 +++++++------------ 5 files changed, 67 insertions(+), 70 deletions(-) diff --git a/src/runtime/alpasim_runtime/config.py b/src/runtime/alpasim_runtime/config.py index 6a1199da..81fce70b 100644 --- a/src/runtime/alpasim_runtime/config.py +++ b/src/runtime/alpasim_runtime/config.py @@ -40,7 +40,7 @@ class GenericSourceConfig: """ enabled: bool = True - data_dir: str = MISSING + data_dir: Optional[str] = None extra_params: Dict[str, Any] = field(default_factory=dict) @@ -86,6 +86,10 @@ def to_trajdata_params(self) -> dict: for dataset_name, source in self.sources.items(): if source.enabled: + if source.data_dir is None: + raise ValueError( + f"data_source.sources.{dataset_name}.data_dir is required when enabled" + ) desired_data.append(dataset_name) data_dirs[dataset_name] = source.data_dir if source.extra_params: diff --git a/src/runtime/alpasim_runtime/runtime_context.py b/src/runtime/alpasim_runtime/runtime_context.py index 589199a4..47244e6d 100644 --- a/src/runtime/alpasim_runtime/runtime_context.py +++ b/src/runtime/alpasim_runtime/runtime_context.py @@ -3,8 +3,8 @@ from __future__ import annotations -import logging import copy +import logging import math from dataclasses import dataclass @@ -193,6 +193,9 @@ async def build_runtime_context( logger.info("Creating UnifiedDataset from config") data_source_config = OmegaConf.to_object(config.user.data_source) trajdata_params = data_source_config.to_trajdata_params() + dataset_kwargs = trajdata_params.pop("dataset_kwargs", None) + if dataset_kwargs: + trajdata_params["dataset_kwargs"] = dataset_kwargs dataset = UnifiedDataset(**trajdata_params) logger.info( f"Created UnifiedDataset with {dataset.num_scenes()} scenes, " diff --git a/src/runtime/alpasim_runtime/scene_loader.py b/src/runtime/alpasim_runtime/scene_loader.py index 4c038f71..8164bb9c 100644 --- a/src/runtime/alpasim_runtime/scene_loader.py +++ b/src/runtime/alpasim_runtime/scene_loader.py @@ -53,14 +53,13 @@ def __init__( # Build dataset_name -> asset_base_path mapping self._asset_base_path_map: dict[str, str] = {} data_source_config = OmegaConf.to_object(config.user.data_source) + for dataset_name, source in data_source_config.sources.items(): - if "asset_base_path" in source.extra_params: - self._asset_base_path_map[dataset_name] = source.extra_params[ - "asset_base_path" - ] + asset_base_path = source.extra_params.get("asset_base_path") + if asset_base_path is not None: + self._asset_base_path_map[dataset_name] = asset_base_path logger.info( - f"Registered asset_base_path for {dataset_name}: " - f"{source.extra_params['asset_base_path']}" + f"Registered asset_base_path for {dataset_name}: {asset_base_path}" ) def get_data_source(self, scene_id: str) -> SceneDataSource: diff --git a/src/runtime/tests/test_config.py b/src/runtime/tests/test_config.py index 37fdc87e..9fe691b7 100644 --- a/src/runtime/tests/test_config.py +++ b/src/runtime/tests/test_config.py @@ -53,20 +53,18 @@ def test_typed_parse_config_invalid_yaml(tmp_path): def test_data_source_config_defaults(): - """Test that DataSourceConfig has sensible defaults for hierarchical structure.""" + """Test that DataSourceConfig has sensible defaults.""" cfg = config.DataSourceConfig( cache_location="/tmp/cache", - usdz=config.USDZSourceConfig(data_dir="/data/usdz"), + sources={"usdz": config.GenericSourceConfig(data_dir="/data/usdz")}, ) assert cfg.cache_location == "/tmp/cache" assert cfg.rebuild_cache is False assert cfg.rebuild_maps is False - assert cfg.num_workers == 4 - assert cfg.usdz is not None - assert cfg.usdz.enabled is True - assert cfg.usdz.data_dir == "/data/usdz" - assert cfg.usdz.desired_dt == 0.1 - assert cfg.usdz.incl_vector_map is True + assert cfg.num_workers == 1 + assert "usdz" in cfg.sources + assert cfg.sources["usdz"].data_dir == "/data/usdz" + assert cfg.sources["usdz"].extra_params == {} def test_data_source_config_to_trajdata_params(): @@ -75,11 +73,14 @@ def test_data_source_config_to_trajdata_params(): cache_location="/tmp/cache", rebuild_cache=True, num_workers=8, - usdz=config.USDZSourceConfig( - data_dir="/data/usdz", - desired_dt=0.05, - asset_base_path="/assets", - ), + desired_dt=0.05, + incl_vector_map=True, + sources={ + "usdz": config.GenericSourceConfig( + data_dir="/data/usdz", + extra_params={"asset_base_path": "/assets"}, + ) + }, ) params = cfg.to_trajdata_params() @@ -90,33 +91,40 @@ def test_data_source_config_to_trajdata_params(): assert params["num_workers"] == 8 assert params["desired_dt"] == 0.05 assert params["incl_vector_map"] is True + assert params["dataset_kwargs"] == {"usdz": {"asset_base_path": "/assets"}} def test_data_source_config_multiple_sources(): """Test configuration with multiple data sources enabled.""" cfg = config.DataSourceConfig( cache_location="/tmp/cache", - usdz=config.USDZSourceConfig(data_dir="/data/usdz"), - nuplan=config.NuPlanSourceConfig( - enabled=True, - data_dir="/data/nuplan", - config_dir="/configs", - ), + sources={ + "usdz": config.GenericSourceConfig(data_dir="/data/usdz"), + "nuplan": config.GenericSourceConfig( + data_dir="/data/nuplan", + extra_params={ + "config_dir": "/configs", + "num_timesteps_before": 30, + "num_timesteps_after": 80, + }, + ), + }, ) params = cfg.to_trajdata_params() assert set(params["desired_data"]) == {"usdz", "nuplan"} assert params["data_dirs"]["usdz"] == "/data/usdz" assert params["data_dirs"]["nuplan"] == "/data/nuplan" - assert params["config_dir"] == "/configs" - assert params["num_timesteps_before"] == 30 + assert params["dataset_kwargs"]["nuplan"]["config_dir"] == "/configs" + assert params["dataset_kwargs"]["nuplan"]["num_timesteps_before"] == 30 + assert params["dataset_kwargs"]["nuplan"]["num_timesteps_after"] == 80 def test_data_source_config_no_sources_enabled(): """Test that error is raised when no data sources are enabled.""" cfg = config.DataSourceConfig( cache_location="/tmp/cache", - nuplan=config.NuPlanSourceConfig(enabled=False), + # All sources are None, so no sources enabled ) with pytest.raises(ValueError, match="No data sources enabled"): diff --git a/src/wizard/configs/base_config.yaml b/src/wizard/configs/base_config.yaml index ad6f2af4..e5114e77 100644 --- a/src/wizard/configs/base_config.yaml +++ b/src/wizard/configs/base_config.yaml @@ -33,6 +33,7 @@ defines: trajdata_cache: "${defines.filesystem}/trajdata_cache" # Trajdata unified cache location sensorsim_entrypoint: "/app/internal/scripts/pycena/runtime/pycena_run" + nre_cache_size: 4 # Default; topology configs override based on concurrent rollouts helper: scripts vscode: sources/remote-vscode-server physics_cache_size: 16 # should match or exceed concurrent scenes to avoid cache thrashing @@ -230,58 +231,40 @@ runtime: # nr_workers and endpoints.*.n_concurrent_rollouts are set by topology configs. # Unified data source configuration using trajdata - # Hierarchical structure separates common settings from source-specific config + # Common settings apply to all sources; source-specific config lives under sources: data_source: # Common configuration (applies to all data sources) cache_location: "${defines.trajdata_cache}" # Shared trajdata cache + desired_dt: 0.1 # 10 Hz sampling rate for trajectories + incl_vector_map: true # Include vector map data (roads, lanes, etc.) rebuild_cache: false # Set to true to force rebuild cache rebuild_maps: false # Set to true to force rebuild maps num_workers: 4 # Parallel workers for cache creation - # USDZ data source configuration (NuRec artifacts) - usdz: - enabled: true - # Use wizard's dynamic sceneset path. The wizard creates a sceneset directory - # based on selected scenes and sets sceneset_path at runtime. - # For manual runtime usage without wizard, set sceneset_path in scenes config. - data_dir: "${scenes.scene_cache}/${scenes.sceneset_path}" - desired_dt: 0.1 # 10 Hz sampling rate for trajectories - incl_vector_map: true # Include vector map data (roads, lanes, etc.) - asset_base_path: null # Optional: Base path for MTGS rendering assets - - # NuPlan data source configuration (disabled by default) - nuplan: - enabled: false - data_dir: null - config_dir: null # Set to enable YAML batch preprocessing mode - num_timesteps_before: 30 # Timesteps before central token - num_timesteps_after: 80 # Timesteps after central token - desired_dt: 0.1 - incl_vector_map: true + sources: + # USDZ data source configuration (NuRec artifacts) + usdz: + enabled: true + # Use wizard's dynamic sceneset path. The wizard creates a sceneset directory + # based on selected scenes and sets sceneset_path at runtime. + # For manual runtime usage without wizard, set sceneset_path in scenes config. + data_dir: "${scenes.scene_cache}/${scenes.sceneset_path}" + extra_params: + asset_base_path: null # Optional: Base path for MTGS rendering assets + + # NuPlan data source configuration (disabled by default) + nuplan: + enabled: false + data_dir: null + extra_params: + config_dir: null # Set to enable YAML batch preprocessing mode + num_timesteps_before: 30 # Timesteps before central token + num_timesteps_after: 80 # Timesteps after central token # Enable cubic spline smoothing for trajectories smooth_trajectories: true endpoints: - sensorsim: - # how many rollouts can run on a single worker at once - n_concurrent_rollouts: 2 - - driver: - n_concurrent_rollouts: 14 - - physics: - n_concurrent_rollouts: 14 - skip: false # physics supports skipping - - controller: - n_concurrent_rollouts: 14 - skip: false # controller supports skipping - - trafficsim: - n_concurrent_rollouts: 14 - skip: true # trafficsim is not included yet - # shut down the system after simulation is finished. without this flag the microservice servers # will remain on forever requring a manual interrupt (useful for debugging) do_shutdown: true