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]

nit: for the new files, should be 2026 only. Applies elsewhere, too

[WCJ-BERT]

Sure, will apply.

[mwatson-nvidia]

The comment seems to indicate that this needs to match the control_timestep_us. Is this the case, or is this just a conveniently high frequency? I suspect it might just be the latter, in which case I would drop this part of the comment as it may be misleading.

If this does need to match the control timestep, we might want to see about some derived parameter

[WCJ-BERT]

Good catch on the misleading comment. I'll remove it.

Regarding whether desired_dt should match control_timestep_us: they serve different purposes and should remain independent:

• desired_dt: Data-layer parameter that controls trajectory sampling rate during cache preparation (via trajdata). This happens before simulation runs.
• control_timestep_us: Runtime parameter that controls the simulation time step and control command frequency.

[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]

nit: I suspect the reason you added this was to make the lines shorter? but now we have some new variable that readers of the code need to internalize to understand the meaning. Suggest just using data_source.rig.trajectory.time_range_us.start and data_source.rig.trajectory.time_range_us.stop rather than introducing a new variable on the stack here

[WCJ-BERT]

Good suggestion - will change this.

[mwatson-nvidia]

is this still used?

[WCJ-BERT]

Not used. Will delete.

[mwatson-nvidia]

This comment seems inaccurate (specifically the word "need"). As I understand it, this is an optional pre-step for optimization?

[WCJ-BERT]

Yeah, i'll change the word.

[mwatson-nvidia]

we have 3 fields: self._dataset, self._scene_id_to_idx, and self._scene_id_to_data_source. I wonder if it would be better to package up these three concepts into a single class so we can pass a single argument? it could also help with readabiltiy

[WCJ-BERT]

Good suggestion - I prepare to refactor the three related fields (_dataset, _scene_id_to_idx, _scene_id_to_data_source) into a single SceneLoader class. This improves encapsulation and readability.

The SceneLoader is initialized from RuntimeContext and provides a single get_data_source(scene_id) method for lazy loading.

[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?