Filter the shards sample schedule for each worker. Returns: list[tuple[int, int]]: The shards to sample, in (dataset_index, shard

Get the cached shard.

q: [B, Lq, Nq, C1]. k: [B, Lk, Nk, C1]. v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. q_lens

1. Force lowercase 2. Remove all punctuations

Format the count as a string in millions, e.g. 11M or 5.5M.

Simulate python f-string but without `f`

Generate modality.json file defining field mappings for the dataset.

Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. Args: num_steps_per_shard (int):

Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. Args: num_steps_per_shard (int):

Get the action data for a trajectory by a base index. If the step indices are out of range, pad with the data: if the data is stor

Get all frames from a video. Returns: tuple[np.ndarray, np.ndarray]: Frames and timestamps.

Generate a map of video frame indices to trajectory indices.

Get the parquet paths for all trajectories. Returns: dict[int, Path]: The parquet paths for all trajectories.

Get the parquet paths for all trajectories. Returns: dict[int, Path]: The parquet paths for all trajectories.

Get the video paths for all trajectories and all views. Returns: dict[int, dict[str, Path]]: The video paths for all trajectories

Get the video paths for all trajectories and all views. Returns: dict[int, dict[str, Path]]: The video paths for all trajectories

Args: x: can be any arbitrary nested structure of np array and torch tensor strict: True to check all batch sizes are the same

First try to find the class in the registry first, if it doesn't exist, use importlib to locate it

Get the data corresponding to the modality for a trajectory by step indices. This method dispatches to the appropriate specialized method bas

Get the data corresponding to the modality for a trajectory by step indices. This method dispatches to the appropriate specialized method bas

Get the embodiment tag from the data.

Read one video frame → (H, W, 3) uint8 RGB.

Split the flat GT action vector into per-key arrays.

Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py

Args: seed: "system": use scrambled int based on system time None or int < 0: invalid seed values, see `handle_invalid_seed

Get the state data for a trajectory by a base index. If the step indices are out of range, pad with the data: if the data is store

Get the state or action data for a trajectory by a base index. If the step indices are out of range, pad with the data: if the dat

These einops patterns will be used to aggregate batches and compute statistics. Note: We assume the images are in channel first format

Get the RAW data for a single step in a trajectory. No transforms are applied. Args: trajectory_id (int): The name of the traject

Get task language instruction from new format - use first episode's data_info.json

Get task language instruction from old format

Returns the training [`~torch.utils.data.DataLoader`]. Will use no sampler if `train_dataset` does not implement `__len__`, a random

Get the trajectories in a shard.

Get the data for a trajectory.

Get the video frames for a trajectory by a base index. Args: dataset (BaseSingleDataset): The dataset to retrieve the data from.

Get the video frames from cached shards for a trajectory by uniformly sampling from language-consistent ranges. Args: trajectory_

Download a file from the Wan2.1-I2V-14B-480P HuggingFace repo to HF cache.

Get host name, alias with ``socket.gethostname()``

Infer the next action from the policy in a server-client setup

Initialize distributed runtime for single-process or torchrun execution.

r""" Initialize model parameters using Xavier initialization.

r""" Initialize model parameters using Xavier initialization.

Returns: string

Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB).

Load all video frames for each camera from the debug_image/ directory. Returns: Dict mapping roboarena camera keys to (N, H, W, 3) uint8

Load local dataset from new format and return a dict with observations and actions

Load local dataset from old format and return a dict with observations and actions

Load the model to GPU for inference.

Load VLM model to GPU and keep it there.