Feat: Unified Scene Data Flow

[WCJ-BERT]

Summary

This PR migrates AlpaSim runtime from USDZ artifact-based scene loading to trajdata UnifiedDataset. It introduces a unified interface for loading scenes from multiple autonomous driving datasets (e.g., USDZ, NuPlan, Waymo) and improves memory efficiency through lazy loading.

Key Changes

Breaking Changes

• RuntimeContext no longer stores full SceneDataSource objects; it now uses a lightweight scene_id_to_idx mapping.
• DaemonEngine now loads scenes on demand in each worker.
• build_pending_jobs_from_request() now takes a callable instead of a dict.

New Functionality

• Added a SceneDataSource protocol to abstract scene data access.
• Added TrajdataDataSource for direct trajdata-based scene loading.
• Added a prepare_data CLI tool for preprocessing trajdata caches.
• Enabled lazy scene loading in workers to reduce memory usage.

Code and Documentation Updates

• Updated existing tests for the new data flow and added integration tests for trajdata functionality.
• Added data preparation instructions to README and TUTORIAL.md.
• Documented the trajdata-alpasim dependency in ONBOARDING.
• Updated wizard config defaults for data_source.

[mwatson-nvidia]

Thanks for the submission!

I have not tested this internally, and that is something that we'll want to do, I think, before proceeding with this PR. However, before I do this, I think it's worth addressing/discussing this first round of comments.

[mwatson-nvidia]

It looks like this data structure is pretty flat, meaning that, as we add additional data source types, this may grow unwieldy and it will be hard for readers to understand which fields are needed for which data sources and which fields can affect which types of data sources. I wonder if it makes sense to try and switch this to a more hierarchical structure now. Do you have any thoughts on this?

[WCJ-BERT]

What about this structure?

runtime:
   data_source:
     cache_location: "/path/to/cache"
     incl_vector_map: true
     rebuild_cache: false
     rebuild_maps: false
     num_workers: 4

     sources:
       usdz:
         enabled: true
         data_dir: "/path/to/usdz"
         desired_dt: 0.1
         extra_params:
              smooth_trajectories: true

       nuplan:
         enabled: false
         data_dir: "/path/to/nuplan"
         extra_params:
               config_dir: "/path/to/yaml"
               asset_base_path: null 
               num_timesteps_before: 30
               num_timesteps_after: 80
 
       waymo:
         enabled: false
         data_dir: "/path/to/waymo"

Also, I introduce a GenericSourceConfig to support hierarchical structure in config.py.

@dataclass
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 dataset directory
        desired_dt: Desired time delta between trajectory frames in seconds
        incl_vector_map: Whether to load vector maps (roads, lanes, etc.)
        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
    extra_params: Dict[str, Any] = field(default_factory=dict)

[mwatson-nvidia]

This is closer but still seems to have the issue where there is a bunch of information that might be present but not used/misleading. What are your thoughts on something closer to:

  @dataclass                                                                                                                                                                                  
  class DataSourceConfig:  
      cache_location: str = MISSING
      desired_dt: float = 0.1
      incl_vector_map: bool = True                                                                                                                                                            
      rebuild_cache: bool = False
      rebuild_maps: bool = False                                                                                                                                                              
      num_workers: int = 1 
                                                                                                                                                                                              
      # Only one active at a time (or a list of these)                                                                                                                                        
      usdz: Optional[UsdzSourceConfig] = None                                                                                                                                                 
      nuplan: Optional[NuPlanSourceConfig] = None                                                                                                                                             
      waymo: Optional[WaymoSourceConfig] = None

[WCJ-BERT]

Thanks for the review! I've carefully considered both approaches and implemented the explicit field structure as requested. However, after implementation, I'd like to discuss which approach better serves our long-term goals.

Two Approaches Comparison

I've explored both structures in detail. Here's what I found:

Approach 1: Explicit Fields

config.py

@dataclass
class USDZSourceConfig:
    data_dir: str = MISSING
    extra_params: Dict[str, Any] = field(default_factory=dict)

@dataclass
class NuPlanSourceConfig:
    data_dir: str = MISSING
    extra_params: Dict[str, Any] = field(default_factory=dict)

@dataclass
class DataSourceConfig:
    cache_location: str = MISSING
    desired_dt: float = 0.1
    incl_vector_map: bool = True
    rebuild_cache: bool = False
    rebuild_maps: bool = False
    num_workers: int = 1

    # Only one active at a time (or a list of these)
    usdz: Optional[USDZSourceConfig] = None
    nuplan: Optional[NuPlanSourceConfig] = None
    waymo: Optional[WaymoSourceConfig] = None

YAML Config:

data_source:
  cache_location: "/path/to/cache"
  desired_dt: 0.1
  incl_vector_map: true
  
  nuplan:
    data_dir: "/path/to/nuplan"
    extra_params:
      config_dir: "/path/to/configs"
      num_timesteps_before: 30
      num_timesteps_after: 80

Approach 2: Generic Sources (Current Implementation)

config.py:

@dataclass
class GenericSourceConfig:
    data_dir: str = MISSING
    extra_params: Dict[str, Any] = field(default_factory=dict)

@dataclass
class DataSourceConfig:
    cache_location: str = MISSING
    desired_dt: float = 0.1
    incl_vector_map: bool = True
    rebuild_cache: bool = False
    rebuild_maps: bool = False
    num_workers: int = 1
    
    # Extensible: supports any trajdata dataset
    sources: Dict[str, GenericSourceConfig] = field(default_factory=dict)

YAML Config:

data_source:
  cache_location: "/path/to/cache"
  desired_dt: 0.1
  incl_vector_map: true
  
  sources:
    nuplan:
      data_dir: "/path/to/nuplan"
      extra_params:
        config_dir: "/path/to/configs"
        num_timesteps_before: 30
        num_timesteps_after: 80
    
    nuscenes:  # Easy to add without code changes
      data_dir: "/path/to/nuscenes"

My Recommendation

I believe Approach 2 better aligns with our architecture goals:

1. trajdata's UnifiedDataset is designed for extensibility - it accepts arbitrary dataset names and parameters through its API. Our config should mirror this flexibility.
2. Maintenance burden - With explicit fields, every new dataset or trajdata API change requires coordinated updates across config.py, to_trajdata_params(), and scene_loader.py.
3. Real-world usage - In practice, all dataset-specific parameters go in extra_params anyway, so we don't actually gain type safety from separate Config classes.
4. Proven pattern - The original design successfully supported multiple datasets and has been working well.

What do you think? I'm happy to keep the explicit field approach if you feel the type distinction is important, but I wanted to share these considerations.

[mwatson-nvidia]

I think I agree with you about approach 2 being a bit better. I'd say go ahead and implement that

[mwatson-nvidia]

Concerns from Claude:

Several concerns:                                                                                                                                                                           
                                                                                                                                                                                              
  1. load_config_from_file unpacks a dataclass into a dict, then everything uses config["key"] / config.get("key", default). This throws away the type safety from DataSourceConfig. The      
  function should just return the DataSourceConfig (or UserSimulatorConfig) and the callers should use it directly. The dict intermediary means defaults are re-specified in multiple places  
  (e.g., num_workers defaults to 8 on lines 653 and 710, but 1 in the dataclass).                                                                                                             
                                                                                                                                                                                              
  2. Inconsistent defaults across paths. desired_dt defaults to 0.1 via CLI (line 415), 0.5 in preprocess_from_yaml_configs (line 164) and run_yaml_batch_preprocessing (line 555), and 0.1 in
   DataSourceConfig. num_workers defaults to 1 in the dataclass, 1 in CLI, but 8 in the config.get() calls (lines 653, 710). These will silently produce different behavior depending on which
   code path you take.                                                                                                                                                                        
                                                                     
  3. Commented-out code. Lines 135-141 have a commented-out "extract every central_token" block — this is dead code that should either be implemented or removed.                             
   
  4. run_yaml_batch_preprocessing is a trivial wrapper. It just calls preprocess_from_yaml_configs and converts bool → int. It doesn't justify being a separate function.                     
                                                                     
  5. Deferred import in load_config_from_file (line 487). The codebase convention is imports at the top of the file.                                                                          
                                                                     
  6. --verbose defaults to True with store_true. This means --verbose is always true unless --quiet is passed. The flag is a no-op — you can never not be verbose except via --quiet.         
  Confusing UX.                                                      
                                                                                                                                                                                              
  7. The --smooth-trajectories flag accepts string choices ["true", "false", "True", "False"] then parses them manually. A BooleanOptionalAction or simply --smooth-trajectories /            
  --no-smooth-trajectories would be cleaner.
                                                                                                                                                                                              
  8. The whole file is ~720 lines for what amounts to "call UnifiedDataset(...) with the right args." The real work is in trajdata. Most of the complexity here is plumbing config between two
   input formats (CLI args vs. YAML) and two modes (basic vs. YAML batch), all through a lossy dict intermediate. Simplifying the config flow (just use the dataclass end-to-end) would cut
  this significantly. 

[mwatson-nvidia]

Claude comments:

  1. The map property is a 250-line monster (lines 670-999). It has two completely different code paths (scene.map_data vs. dataset._map_api) with duplicated transformation logic, followed  
  by ~100 lines of post-transform verification and data-type fixups. This should be broken into separate methods. The transform logic itself (transform_map_points) is defined as a nested    
  closure, which makes it harder to test.                                                                                                                                                     
                                                                                                                                                                                              
  2. Excessive defensive hasattr checks everywhere. The code is littered with hasattr(lane, "center"), hasattr(lane.center, "points"), hasattr(agent.extent, "length"), etc. If these are     
  typed trajdata objects, these checks shouldn't be necessary. If the schema is genuinely unstable, that's a bigger problem. This reads like the author wasn't sure what the trajdata API
  actually guarantees.                               

[mwatson]: this might actually be expected--I recall that trajdata has a lot of optional stuff. At the same time, we might do better to assume we have what we need and catch exceptions when fields don't exist? I'd like to hear your thoughts on this                                                                                                                                        
                                                                     
  3. The _extract_agent_trajectory state accessor pattern (lines 305-330) is wild. It does state.get_attr("x") if hasattr(state, "get_attr") else state.x, then checks if the result is       
  scalar, wraps it in an array, then immediately unwraps it with float(x[0] if x.ndim > 0 else x). This suggests uncertainty about what get_raw_state returns — it should be pinned down once
  and handled simply.                                                                                                                                                                         
                                                                     
  4. n/a                                                                                                                              
                                                                     
  5. Coordinate transformation is copy-pasted three times. Ego trajectory (line 424), traffic objects (line 609), and map (lines 805-838) all apply positions + translation. This should be a 
  single utility function.
                                                                                                                                                                                              
  6. from_agent_batch is dead code. It calls _load_from_batch which immediately raises NotImplementedError. This should be removed — it's not a stub, it's misleading.                        
  
  7. The smoothing code in traffic_objects (lines 618-642) has a deferred import csaps inside a loop body. This violates the codebase convention (imports at top of file), and the import     
  happens once per traffic object if smoothing is enabled.           
                                                                                                                                                                                              
  8. metadata generates random UUIDs and datetime.now() (lines 1038-1039). This makes it non-deterministic — running the same scene twice produces different metadata. dataset_hash being a   
  random UUID defeats the purpose of a hash.
                                                                                                                                                                                              
  9. Implicit ordering dependency. traffic_objects silently forces self.rig to be loaded first (line 605) to get world_to_nre. map does the same (lines 691-693, 766-768). This coupling is   
  hidden — if someone accesses map or traffic_objects before rig, it triggers a chain of lazy loading with side effects. The dependency should be explicit (e.g., require world_to_nre as a
  constructor argument, or compute it once eagerly).                                                                                                                                          
                                                                     
  10. _scene_id as a mutable dataclass field with a property setter (lines 82-187) is odd for what should be an immutable identifier. The scene_id property has three resolution paths (the   
  field, _scene.name, or raise). This complexity isn't justified.

[mwatson-nvidia]

I feel like there should probably be a tie between this and the existing Artifact. For instance, should we indicate that Artifact implements this protocol, or just keep with the current duck typed approach?

[WCJ-BERT]

Good question about the relationship between Artifact and SceneDataSource!

Design Intent:
Our goal is to unify all data sources (USDZ, NuPlan, Waymo) through trajdata, using the SceneDataSource protocol as the user-facing interface.

Current Architecture:

• SceneDataSource - Protocol defining the unified interface for Runtime
• TrajdataDataSource - Implements SceneDataSource, provides unified access to all data formats
• Artifact - Internal utility for reading USDZ files, used by trajdata during USDZ data conversion

Why Artifact Still Exists:
Artifact is still needed as an internal component. When trajdata processes USDZ data, it uses Artifact under the hood to read the USDZ file format.

In terms of protocols:

• Artifact and SceneDataSource operate at different abstraction levels
• Artifact is an implementation detail (USDZ file reader)
• SceneDataSource is the public interface (unified data access)

What do you think of this design?

[mwatson-nvidia]

I like the idea of a more generic interface/protocol with different backends. However, Artifact already structurally satisfies the SceneDataSource protocol. So it feels like a better approach might be to make that relationship explicit and have Artifact formally implement SceneDataSource. Then, we could look at the existing consumers of Artifact and see if they can be switched over to SceneDataSource so that we could more easily reuse the code for the new sources.

[mwatson-nvidia]

Additionally, I tried running this PR locally and got a string of errors running the basic tutorial uv run alpasim_wizard deploy=local topology=1gpu driver=vavam wizard.log_dir=$PWD/tutorial . Can you try running it locally and see if it works for you?

[WCJ-BERT]

Thanks for testing locally. The root cause was a mismatch between base_config.yaml and the DataSourceConfig schema — usdz/nuplan were direct keys under data_source, but the dataclass expected them under sources. This caused OmegaConf to fail at config parse time before anything else ran. I've already fixed these bugs.

One note on trajdata: the project currently pins trajdata-alpasim from the alpasim branch on GitHub. I've made some improvements to that branch (including USDZ support enhancements) and have a separate PR open for it. Running this PR requires that version of trajdata — I'll make sure to link the trajdata PR here so they can be reviewed together. NVlabs/trajdata#62

The lastest version works in my local.
Prepare Data:

Test Runtime:

Would appreciate it if you could give it a spin on your end. (◕‿◕)

[mwatson-nvidia]

Sorry for not catching this earlier. One concern I have is that we can have two artifacts with the same clip id (which is I think how the caching is done) but with different versions. The way that we get around this in the existing approach is to have a UUID for the usdz filename and an a scene database (really a csv) to disambiguate

[WCJ-BERT]

Thanks for raising this — you're right that trajdata uses the internal clip ID as the cache key, not the UUID.

If we want a stronger guarantee, we could incorporate the nre_version_string (e.g. 25.7.8-fc8b0551) into the cache_location path to fully isolate per-version caches. How do you think of this solution?

[WCJ-BERT]

Otherwise, i have a similar proposal. If we add ${scenes.sceneset_path} to base_config.yaml, is it helpful to distinguish diffferent versions' usdz? I think sceneset_path is the MD5 fingerprint of all UUIDs (sorted) selected for the run — so any version change in a scene produces a different UUID, which changes the hash, which routes to a new isolated cache directory.

runtime:
  # nr_workers and endpoints.*.n_concurrent_rollouts are set by topology configs.

  # Unified data source configuration using trajdata
  # 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}/${scenes.sceneset_path}"  # Per-sceneset cache, isolates different scene versions
    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

    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****