### Initialize and Step Through a Simulation Scene

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Initialize a simulation scene from a dataset, reset it to get initial agent states, and step through timesteps by providing agent's next states. This example replays ground truth motions.

```python
from typing import Dict

import numpy as np

from trajdata import AgentBatch, UnifiedDataset
from trajdata.data_structures.scene_metadata import Scene
from trajdata.simulation import SimulationScene

dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
    },
)

desired_scene: Scene = dataset.get_scene(scene_idx=0)
sim_scene = SimulationScene(
    env_name="nusc_mini_sim",
    scene_name="sim_scene",
    scene=desired_scene,
    dataset=dataset,
    init_timestep=0,
    freeze_agents=True,
)

obs: AgentBatch = sim_scene.reset()
for t in range(1, sim_scene.scene.length_timesteps):
    new_xyh_dict: Dict[str, np.ndarray] = dict()

    for idx, agent_name in enumerate(obs.agent_name):
        curr_yaw = obs.curr_agent_state[idx, -1]
        curr_pos = obs.curr_agent_state[idx, :2]

        next_state = np.zeros((3,))
        next_state[:2] = curr_pos
        next_state[2] = curr_yaw
        new_xyh_dict[agent_name] = next_state

    obs = sim_scene.step(new_xyh_dict)
```

--------------------------------

### SimulationScene Initialization and Setup

Source: https://context7.com/nvlabs/trajdata/llms.txt

Demonstrates initializing a SimulationScene from a dataset and setting up metrics and statistics. Requires a configured UnifiedDataset.

```python
from collections import defaultdict
from typing import Dict
import numpy as np
from trajdata import AgentBatch, AgentType, UnifiedDataset
from trajdata.simulation import SimulationScene, sim_metrics, sim_stats
from trajdata.data_structures.state import StateArray

# Setup dataset for simulation
dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    only_types=[AgentType.VEHICLE],
    agent_interaction_distances=defaultdict(lambda: 50.0),
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

# Initialize simulation from a real scene
scene = dataset.get_scene(scene_idx=0)
sim_scene = SimulationScene(
    env_name="nusc_mini_sim",
    scene_name="sim_scene",
    scene=scene,
    dataset=dataset,
    init_timestep=0,
    freeze_agents=True,  # Keep same agents throughout simulation
)

# Setup metrics

ade = sim_metrics.ADE()
fde = sim_metrics.FDE()

# Setup statistics histograms
vel_hist = sim_stats.VelocityHistogram(bins=np.linspace(0, 40, 41))
jerk_hist = sim_stats.JerkHistogram(bins=np.linspace(0, 40, 41), dt=scene.dt)

# Reset and get initial observation
obs: AgentBatch = sim_scene.reset()
```

--------------------------------

### Install trajdata with pip

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Install the trajdata library using pip. This is the basic installation command.

```sh
pip install trajdata
```

--------------------------------

### Scene-Centric Batching Setup

Source: https://context7.com/nvlabs/trajdata/llms.txt

Illustrates the setup for loading data using scene-centric batching, where all agents in a scene are loaded at each timestep. This requires importing `SceneBatch` and `UnifiedDataset`.

```python
from torch.utils.data import DataLoader
from trajdata import SceneBatch, AgentType, UnifiedDataset
```

--------------------------------

### Install trajdata with Dataset Support

Source: https://context7.com/nvlabs/trajdata/llms.txt

Install trajdata using pip, with optional dependencies for specific dataset support or all dataset dependencies.

```bash
pip install trajdata
```

```bash
pip install "trajdata[nusc]"      # nuScenes
```

```bash
pip install "trajdata[lyft]"      # Lyft Level 5
```

```bash
pip install "trajdata[waymo]"     # Waymo Open Motion
```

```bash
pip install "trajdata[interaction]"  # INTERACTION Dataset
```

```bash
pip install "trajdata[vod]"       # View-of-Delft
```

```bash
pip install "trajdata[nusc,lyft,waymo,interaction,vod]"
```

--------------------------------

### MapAPI and VectorMap Initialization and Usage

Source: https://context7.com/nvlabs/trajdata/llms.txt

Shows how to initialize the MapAPI and load VectorMap data. Includes examples for querying lanes and rasterizing the map. Requires a cache path and map environment:location identifier.

```python
from pathlib import Path
from trajdata import MapAPI, VectorMap

# Initialize Map API with cache path
cache_path = Path("~/.unified_data_cache").expanduser()
map_api = MapAPI(cache_path)

# Load a vector map by environment:location identifier
vec_map: VectorMap = map_api.get_map("nusc_mini:boston-seaport")

print(f"Map: {vec_map.env_name}, {vec_map.map_name}")
print(f"Number of lanes: {len(vec_map.lanes)}")
print(f"Map extent: {vec_map.extent}")  # [min_x, min_y, min_z, max_x, max_y, max_z]
```

```python
# Query closest lane to a point
import numpy as np
query_point = np.array([500.0, 1000.0, 0.0])
closest_lane = vec_map.get_closest_lane(query_point)
print(f"Closest lane ID: {closest_lane.id}")
```

```python
# Get all lanes within a radius
nearby_lanes = vec_map.get_lanes_within(query_point, radius=30.0)
print(f"Lanes within 30m: {len(nearby_lanes)}")
```

```python
# Lane centerline interpolation
lane = vec_map.lanes[0]
interpolated = lane.center.interpolate(max_dist=0.5)  # Max 0.5m between points
print(f"Original points: {len(lane.center.points)}, "
      f"Interpolated: {len(interpolated.points)}")
```

```python
# Project points onto lane centerline
random_points = lane.center.midpoint + np.random.uniform(-2, 2, size=(10, 4))
projected = lane.center.project_onto(random_points)
```

```python
# Rasterize map for visualization
map_img, raster_from_world = vec_map.rasterize(
    resolution=2,              # pixels per meter
    return_tf_mat=True,
    incl_centerlines=True,
    area_color=(255, 255, 255),
    edge_color=(0, 0, 0),
)
```

--------------------------------

### Install trajdata with dataset extras

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Install trajdata along with specific dataset dependencies. Use this command to include support for datasets like nuScenes, Lyft, Waymo, INTERACTION, or View-of-Delft.

```sh
# For nuScenes
pip install "trajdata[nusc]"

# For Lyft
pip install "trajdata[lyft]"

# For Waymo
pip install "trajdata[waymo]"

# For INTERACTION
pip install "trajdata[interaction]"

# For View-of-Delft 
pip install "trajdata[vod]"

# All
pip install "trajdata[nusc,lyft,waymo,interaction,vod]"
```

--------------------------------

### Package Developer Installation

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Install trajdata in editable mode for development. This involves installing requirements and then the package itself.

```sh
pip install -r requirements.txt
pip install -e .
```

--------------------------------

### Initialize UnifiedDataset with Configuration

Source: https://context7.com/nvlabs/trajdata/llms.txt

Create a UnifiedDataset instance with options for data sources, batching, time horizons, agent filtering, map inclusion, and standardization. Requires specifying data directories.

```python
from collections import defaultdict
from torch.utils.data import DataLoader
from trajdata import AgentBatch, AgentType, UnifiedDataset

# Create dataset with various configuration options
dataset = UnifiedDataset(
    desired_data=["nusc_mini-mini_train"],  # Dataset and split specification
    centric="agent",                         # "agent" or "scene" batching
    desired_dt=0.1,                          # Resample data to 10Hz
    history_sec=(3.2, 3.2),                  # (min, max) history in seconds
    future_sec=(4.8, 4.8),                   # (min, max) future in seconds
    only_predict=[AgentType.VEHICLE],        # Only predict vehicles
    agent_interaction_distances=defaultdict(lambda: 30.0),  # Neighbor distance
    incl_robot_future=False,                 # Include ego future trajectory
    incl_raster_map=True,                    # Include rasterized map patches
    raster_map_params={
        "px_per_m": 2,                       # Map resolution
        "map_size_px": 224,                  # Map patch size
        "offset_frac_xy": (-0.5, 0.0),       # Agent offset in patch
    },
    incl_vector_map=False,                   # Include vector map data
    standardize_data=True,                   # Normalize to agent frame
    num_workers=4,                           # Parallel preprocessing workers
    verbose=True,
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
    },
)

print(f"Total samples: {len(dataset):,}")

# Create PyTorch DataLoader with custom collate function
dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    collate_fn=dataset.get_collate_fn(),
    num_workers=4,
)

# Iterate over batches
batch: AgentBatch
for batch in dataloader:
    # batch.agent_hist: (B, T_hist, state_dim) - agent history trajectories
    # batch.agent_fut: (B, T_fut, state_dim) - agent future trajectories
    # batch.neigh_hist: (B, N, T_hist, state_dim) - neighbor histories
    # batch.maps: (B, C, H, W) - rasterized map patches
    # batch.curr_agent_state: (B, state_dim) - current agent state
    print(f"Batch shape: {batch.agent_hist.shape}")
    break

```

--------------------------------

### AgentBatch Data Structure and Operations

Source: https://context7.com/nvlabs/trajdata/llms.txt

Demonstrates the structure of the AgentBatch dataclass and common operations like moving to GPU and filtering by agent type. Assumes batch is loaded from a dataloader.

```python
from trajdata import AgentBatch

# AgentBatch fields (after loading from dataloader)
batch: AgentBatch

# Core trajectory data
batch.agent_name        # List[str]: Agent identifiers
batch.agent_type        # Tensor: AgentType enum values
batch.curr_agent_state  # StateTensor: Current state (B, state_dim)
batch.agent_hist        # StateTensor: History (B, T_hist, obs_dim)
batch.agent_hist_len    # Tensor: Valid history length per sample
batch.agent_fut         # StateTensor: Future ground truth (B, T_fut, obs_dim)
batch.agent_fut_len     # Tensor: Valid future length per sample

# Neighbor data
batch.num_neigh         # Tensor: Number of neighbors per sample
batch.neigh_hist        # StateTensor: Neighbor histories (B, N, T_hist, obs_dim)
batch.neigh_fut         # StateTensor: Neighbor futures (B, N, T_fut, obs_dim)
batch.neigh_types       # Tensor: Neighbor agent types

# Map data (if incl_raster_map=True)
batch.maps              # Tensor: Rasterized maps (B, C, H, W)
batch.map_names         # List[str]: Map location names
batch.rasters_from_world_tf  # Tensor: World to raster transforms

# Coordinate transforms
batch.agents_from_world_tf  # Tensor: World to agent frame transforms
batch.dt                # Tensor: Time delta per sample
```

```python
# Move batch to GPU
batch.to("cuda")
```

```python
# Filter batch by agent type
vehicle_batch = batch.for_agent_type(AgentType.VEHICLE)
```

--------------------------------

### Run Simulation Loop

Source: https://context7.com/nvlabs/trajdata/llms.txt

Iterates through simulation timesteps, updates agent states based on ground truth future, and collects metrics. Requires a sim_scene object and observation data.

```python
for t in range(1, sim_scene.scene.length_timesteps):
    new_states: Dict[str, StateArray] = {}

    for idx, agent_name in enumerate(obs.agent_name):
        curr_pos = obs.curr_agent_state[idx].position.numpy()
        curr_yaw = obs.curr_agent_state[idx].heading.item()

        # Example: replay ground truth future
        if obs.agent_fut_len[idx] >= 1:
            world_from_agent = np.array([
                [np.cos(curr_yaw), np.sin(curr_yaw)],
                [-np.sin(curr_yaw), np.cos(curr_yaw)],
            ])
            next_pos = obs.agent_fut[idx, 0].position.numpy() @ world_from_agent + curr_pos
            next_yaw = curr_yaw + obs.agent_fut[idx, 0].heading.item()
        else:
            next_pos = curr_pos
            next_yaw = curr_yaw

        next_state = np.array([next_pos[0], next_pos[1], 0.0, next_yaw])
        new_states[agent_name] = StateArray.from_array(next_state, "x,y,z,h")

    # Step simulation
    obs = sim_scene.step(new_states)

    # Get metrics at each step
    metrics = sim_scene.get_metrics([ade, fde])
    print(f"Step {t}: ADE={list(metrics['ade'].values())[0]:.3f}")

# Get final statistics
stats = sim_scene.get_stats([vel_hist, jerk_hist])

# Finalize and save simulation results
sim_scene.finalize()
sim_scene.save()
```

--------------------------------

### Visualize Agent Trajectories with Plotting Utilities

Source: https://context7.com/nvlabs/trajdata/llms.txt

Sets up a UnifiedDataset for agent-centric visualization, including raster maps, and then uses plotting utilities to visualize agent batches statically and interactively.

```python
from torch.utils.data import DataLoader
from trajdata import AgentBatch, AgentType, UnifiedDataset
from trajdata.visualization.vis import plot_agent_batch
from trajdata.visualization.interactive_vis import plot_agent_batch_interactive

dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    centric="agent",
    desired_dt=0.1,
    only_predict=[AgentType.VEHICLE],
    incl_raster_map=True,
    raster_map_params={
        "px_per_m": 2,
        "map_size_px": 224,
        "offset_frac_xy": (-0.5, 0.0),
    },
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

dataloader = DataLoader(
    dataset,
    batch_size=4,
    shuffle=True,
    collate_fn=dataset.get_collate_fn(),
)

for batch in dataloader:
    # Static visualization
    plot_agent_batch(batch, batch_idx=0)

    # Interactive visualization with Bokeh
    plot_agent_batch_interactive(
        batch,
        batch_idx=0,
        cache_path=dataset.cache_path
    )
    break
```

--------------------------------

### Create Scene-Centric Dataset

Source: https://context7.com/nvlabs/trajdata/llms.txt

Initializes a UnifiedDataset for scene-centric batching with specified data, time steps, agent types, and history/future durations. It then creates a DataLoader for iterating through batches.

```python
dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    centric="scene",           # Scene-centric batching
    desired_dt=0.1,
    history_sec=(2.0, 2.0),
    future_sec=(4.0, 4.0),
    only_types=[AgentType.VEHICLE],
    max_agent_num=32,          # Max agents per scene sample
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

dataloader = DataLoader(
    dataset,
    batch_size=8,
    shuffle=True,
    collate_fn=dataset.get_collate_fn(),
)

batch: SceneBatch
for batch in dataloader:
    # batch.num_agents: (B,) - number of agents per scene
    # batch.agent_hist: (B, N, T, D) - all agent histories
    # batch.agent_fut: (B, N, T, D) - all agent futures
    print(f"Agents per scene: {batch.num_agents}")
    print(f"Agent histories shape: {batch.agent_hist.shape}")

    # Convert to agent-centric batch for specific agents
    agent_indices = batch.num_agents.new_zeros(batch.num_agents.shape[0])
    agent_batch = batch.to_agent_batch(agent_indices)
    break
```

--------------------------------

### Filter Agents Using AgentType Enumeration

Source: https://context7.com/nvlabs/trajdata/llms.txt

Demonstrates how to use the AgentType enumeration to filter agents by type when creating a UnifiedDataset. It also shows how to access agent types within a batch.

```python
from trajdata import AgentType, UnifiedDataset

# Available agent types
print(f"UNKNOWN: {AgentType.UNKNOWN}")      # 0
print(f"VEHICLE: {AgentType.VEHICLE}")      # 1
print(f"PEDESTRIAN: {AgentType.PEDESTRIAN}")  # 2
print(f"BICYCLE: {AgentType.BICYCLE}")      # 3
print(f"MOTORCYCLE: {AgentType.MOTORCYCLE}")  # 4

# Filter to specific agent types
dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    only_types=[AgentType.VEHICLE, AgentType.PEDESTRIAN],  # Include only these
    only_predict=[AgentType.VEHICLE],  # But only predict vehicles
    no_types=[AgentType.UNKNOWN],      # Exclude unknown agents
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

# Access agent types in batch
for batch in dataset:
    unique_types = batch.agent_types()
    print(f"Agent types in batch: {[t.name for t in unique_types]}")
    break
```

--------------------------------

### Preprocess and Cache Data

Source: https://context7.com/nvlabs/trajdata/llms.txt

Preprocesses multiple datasets and regenerates map data, forcing a rebuild of the cache. This speeds up subsequent data loading by preparing the data beforehand.

```python
import os
from trajdata import UnifiedDataset

# Preprocess multiple datasets in parallel
dataset = UnifiedDataset(
    desired_data=["nusc_mini", "lyft_sample", "nuplan_mini"],
    rebuild_cache=True,    # Force reprocessing
    rebuild_maps=True,     # Force map regeneration
    num_workers=os.cpu_count(),
    verbose=True,
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
        "lyft_sample": "~/datasets/lyft/scenes/sample.zarr",
        "nuplan_mini": "~/datasets/nuplan/dataset/nuplan-v1.1",
    },
)

print(f"Preprocessed {len(dataset):,} total samples")
print(f"Cache location: {dataset.cache_path}")
```

--------------------------------

### Load Multiple Datasets with UnifiedDataset

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Load data from multiple datasets by specifying their names and local directories. Ensure to set `desired_dt` for consistent timestep interpolation if datasets have different annotation frequencies.

```python
dataset = UnifiedDataset(
    desired_data=["nusc_mini", "eupeds_eth"],
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
        "eupeds_eth": "~/datasets/eth_ucy_peds"
    },
    desired_dt=0.1,
)
```

--------------------------------

### Filter and Cache Dataset

Source: https://context7.com/nvlabs/trajdata/llms.txt

Demonstrates applying a custom filter to a dataset and caching the results for faster loading. The `filter_busy_scenes` function keeps samples with more than 5 neighbors. Caching can be done with or without a filter.

```python
from collections import defaultdict
from trajdata import AgentType, UnifiedDataset
from trajdata.data_structures.batch_element import AgentBatchElement

# Create dataset
dataset = UnifiedDataset(
    desired_data=["nusc_mini-mini_val"],
    centric="agent",
    desired_dt=0.5,
    history_sec=(2.0, 2.0),
    future_sec=(4.0, 4.0),
    only_predict=[AgentType.VEHICLE],
    agent_interaction_distances=defaultdict(lambda: 30.0),
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

print(f"Original dataset size: {len(dataset):,}")

# Define custom filter function
def filter_busy_scenes(element: AgentBatchElement) -> bool:
    """Keep only samples with more than 5 neighbors."""
    return element.num_neighbors > 5

# Apply filter (modifies dataset in-place)
dataset.apply_filter(
    filter_fn=filter_busy_scenes,
    num_workers=4,
    max_count=1000,  # Optional: stop after finding 1000 matches
)

print(f"Filtered dataset size: {len(dataset):,}")

# Cache entire dataset to disk for fast loading
cache_path = "./dataset_cache.dill"
dataset.load_or_create_cache(
    cache_path=cache_path,
    num_workers=4,
    filter_fn=filter_busy_scenes,
)

# Subsequent runs will load from cache automatically
```

--------------------------------

### Access Vector Map Data with MapAPI

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Initialize MapAPI with a cache path to access raw vector map information from supported datasets. Retrieve a specific vector map using its dataset and scene identifier.

```python
from pathlib import Path
from trajdata import MapAPI, VectorMap

cache_path = Path("~/.unified_data_cache").expanduser()
map_api = MapAPI(cache_path)

vector_map: VectorMap = map_api.get_map("nusc_mini:boston-seaport")
```

--------------------------------

### Load Trajectory Data with UnifiedDataset

Source: https://github.com/nvlabs/trajdata/blob/main/README.md

Load trajectory data using UnifiedDataset and DataLoader. Ensure to configure `data_dirs` to match your filesystem. The `get_collate_fn()` is essential for batching.

```python
import os
from torch.utils.data import DataLoader
from trajdata import AgentBatch, UnifiedDataset

# See below for a list of already-supported datasets and splits.
dataset = UnifiedDataset(
    desired_data=["nusc_mini"],
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes"
    },
)

dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    collate_fn=dataset.get_collate_fn(),
    num_workers=os.cpu_count(), # This can be set to 0 for single-threaded loading, if desired.
)

batch: AgentBatch
for batch in dataloader:
    # Train/evaluate/etc.
    pass
```

--------------------------------

### Load Multiple Datasets Simultaneously

Source: https://context7.com/nvlabs/trajdata/llms.txt

Combine data from multiple trajectory datasets into a single UnifiedDataset, interpolating to a common frequency. Requires specifying data directories for each dataset.

```python
from trajdata import UnifiedDataset

# Load multiple datasets simultaneously
dataset = UnifiedDataset(
    desired_data=["nusc_mini", "eupeds_eth", "lyft_sample"],
    desired_dt=0.1,  # Common dt for all datasets (required when mixing)
    data_dirs={
        "nusc_mini": "~/datasets/nuScenes",
        "eupeds_eth": "~/datasets/eth_ucy_peds",
        "lyft_sample": "~/datasets/lyft/scenes/sample.zarr",
    },
)

print(f"Combined samples from all datasets: {len(dataset):,}")

# Access individual scenes
for scene_idx in range(min(5, dataset.num_scenes())):
    scene = dataset.get_scene(scene_idx)
    print(f"Scene: {scene.name}, Env: {scene.env_name}, "
          f"Length: {scene.length_timesteps} timesteps")

```

--------------------------------

### Apply Data Augmentation

Source: https://context7.com/nvlabs/trajdata/llms.txt

Applies noise augmentation to agent histories during data loading. The `NoiseHistories` class adds Gaussian noise with a specified standard deviation to agent positions.

```python
from collections import defaultdict
from trajdata import AgentType, UnifiedDataset
from trajdata.augmentation import NoiseHistories

# Create noise augmentation for agent histories
noise_augmentation = NoiseHistories(
    mean=0.0,
    std=0.1,  # Add Gaussian noise with std=0.1 to positions
)

dataset = UnifiedDataset(
    desired_data=["nusc_mini-mini_train"],
    centric="agent",
    desired_dt=0.1,
    history_sec=(3.2, 3.2),
    future_sec=(4.8, 4.8),
    only_predict=[AgentType.VEHICLE],
    augmentations=[noise_augmentation],  # Apply augmentations
    data_dirs={"nusc_mini": "~/datasets/nuScenes"},
)

# Augmentations are applied automatically during data loading
for batch in dataset:
    print(f"Augmented history shape: {batch.agent_hist.shape}")
    break
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.