Load LIBERO init states compatibly across PyTorch versions.

initial_actions: list[dict[str, dict[str, np.ndarray]]] 0: (the first dataset) trajectory_name: action_key: act

Load only LoRA weights from a pretrained model without loading config.

Offload action head components to CPU.

Pack **kwargs or a single dict arg as dict def f(*args, **kwargs): kwdict = pack_kwargs(args, kwargs) # kwdict is now packed

Pack *args or a single list arg as list def f(*args): arg_list = pack_varargs(args) # arg_list is now packed as a list

Parse a JSON string like '{"joint_pos": [0, 6], "gripper": [6, 7]}'.

We rewrite this method to copy mp4 videos to the target position

Huggingface will call model.train() at each training_step. To ensure the expected behaviors for modules like dropout, batchnorm, etc.

Set per-horizon statistics for transforms that support it (e.g., PerHorizonActionTransform). Args: per_horizon_stats: Dic

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`):

Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.

Returns: numpy image of shape [H, W, C] in "auto" mode, we assume C == 3

r""" Reconstruct video tensors from patch embeddings. Args: x (Tensor): Patchified features, with shape [

Initialize the dataset. Args: dataset_path (Path | str): The path to the dataset. modality_configs (dict[str

Args: batch: video: [V, T, C, H, W] Returns: required input with the format `BatchFeature`

Broadcast batch data from rank 0 to all other ranks.

Broadcast batch data from rank 0 to all other ranks.

Calculate relative action statistics for a specific action key. Args: action_key: The action key to calculate stats for (

Use the config to check if the keys are valid and detect silent data corruption.

Convert AR_droid action dict to roboarena action array. AR_droid format: - action.joint_position: (N, 7) - ac

Convert roboarena observation format to AR_droid format. Roboarena format: - observation/exterior_image_0_left: (H, W, 3)

Convert absolute action to relative action by subtracting reference state. Args: action_data: Absolute action data, shape

Initialize a Per-GPU cross-attention cache for the Wan model.

Initialize a Per-GPU KV cache for the Wan model.

Extract all frames and their timestamps using ffmpeg.

Extract frames at specific timestamps using ffmpeg.

Copied from python standard lib shutil.copytree except that we allow exist_ok Use f_copytree as entry

r""" Forward pass through the diffusion model blocks.

r""" Forward pass through the diffusion model Args: x (List[Tensor]): List of input video tensors, each w

Get the trajectory IDs and base indices for all steps in the dataset. Returns: list[tuple[int, int]]: A list of (trajectory_id, b

Get the chunk size for the LeRobot dataset.

Get the data path pattern for the LeRobot dataset.

Restructure the delta indices to use modality.key as keys instead of just the modalities.

Get the detailed global instructions for the dataset. Loads from episodes_detail_global_instruction.jsonl if it exists.

Get language instruction from metadata files for special language keys. Supports: - detailed_global_instruction_medium: Longe

Get the metadata for the LeRobot dataset.

Get the metadata for the LeRobot dataset.

Get the per-horizon relative action stats metadata for the LeRobot dataset. Similar to _get_lerobot_relative_stats_meta but calculate

Get the relative action stats metadata for the LeRobot dataset. Returns: dict[str, DatasetStatisticalValues]: Dictionary

Get the metadata for the LeRobot dataset.

Calculate the maximum delta index across all modalities. Returns: int: The maximum delta index value.

Get the metadata for the dataset. Returns: dict: The metadata for the dataset.

Get the modality keys for the dataset. Returns: dict: Dictionary mapping modality names to their keys.

Get the step filter for the dataset.

Extract target_rotations from StateActionTransform instances in the pipeline.

Get the tasks for the dataset.

Get the trajectories in the dataset.