# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json from collections import defaultdict from pathlib import Path from random import randint import imageio import numpy as np import pandas as pd import torch from einops import rearrange from PIL import Image from pydantic import BaseModel, ValidationError from torch.utils.data import Dataset from tqdm import tqdm from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import EmbodimentTag from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.data.schema import ( DatasetMetadata, DatasetStatisticalValues, LeRobotModalityMetadata, LeRobotStateActionMetadata, ) # from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.data.transform. import ComposedModalityTransform from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.data.transform.base import ComposedModalityTransform from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.utils.video import ( get_all_frames, get_frames_by_timestamps, ) # from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.groot_configs import construct_modality_config_and_transforms LE_ROBOT_MODALITY_FILENAME = "meta/modality.json" LE_ROBOT_EPISODE_FILENAME = "meta/episodes.jsonl" LE_ROBOT_TASKS_FILENAME = "meta/tasks.jsonl" LE_ROBOT_INFO_FILENAME = "meta/info.json" LE_ROBOT_STATS_FILENAME = "meta/stats.json" LE_ROBOT_DATA_FILENAME = "data/*/*.parquet" def calculate_dataset_statistics(parquet_paths: list[Path]) -> dict: """Calculate the dataset statistics of all columns for a list of parquet files.""" # Dataset statistics all_low_dim_data_list = [] # Collect all the data for parquet_path in tqdm( sorted(list(parquet_paths)), desc="Collecting all parquet files...", ): # Load the parquet file parquet_data = pd.read_parquet(parquet_path) parquet_data = parquet_data all_low_dim_data_list.append(parquet_data) all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) # Compute dataset statistics dataset_statistics = {} for le_modality in all_low_dim_data.columns: print(f"Computing statistics for {le_modality}...") np_data = np.vstack([np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]]) dataset_statistics[le_modality] = { "mean": np.mean(np_data, axis=0).tolist(), "std": np.std(np_data, axis=0).tolist(), "min": np.min(np_data, axis=0).tolist(), "max": np.max(np_data, axis=0).tolist(), "q01": np.quantile(np_data, 0.01, axis=0).tolist(), "q99": np.quantile(np_data, 0.99, axis=0).tolist(), } return dataset_statistics class ModalityConfig(BaseModel): """Configuration for a modality.""" delta_indices: list[int] """Delta indices to sample relative to the current index. The returned data will correspond to the original data at a sampled base index + delta indices.""" modality_keys: list[str] """The keys to load for the modality in the dataset.""" class LeRobotSingleDataset(Dataset): """ Base dataset class for LeRobot that supports sharding. """ def __init__( self, dataset_path: Path | str, modality_configs: dict[str, ModalityConfig], embodiment_tag: str | EmbodimentTag, video_backend: str = "decord", video_backend_kwargs: dict | None = None, transforms: ComposedModalityTransform | None = None, single_base_index: bool = False, ): """ Initialize the dataset. Args: dataset_path (Path | str): The path to the dataset. modality_configs (dict[str, ModalityConfig]): The configuration for each modality. The keys are the modality names, and the values are the modality configurations. See `ModalityConfig` for more details. video_backend (str): Backend for video reading. video_backend_kwargs (dict): Keyword arguments for the video backend when initializing the video reader. transforms (ComposedModalityTransform): The transforms to apply to the dataset. embodiment_tag (EmbodimentTag): Overload the embodiment tag for the dataset. e.g. define it as "new_embodiment" """ # first check if the path directory exists if not Path(dataset_path).exists(): raise FileNotFoundError(f"Dataset path {dataset_path} does not exist") self.modality_configs = modality_configs self.video_backend = video_backend self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {} self.transforms = transforms if transforms is not None else ComposedModalityTransform(transforms=[]) self._dataset_path = Path(dataset_path) self._dataset_name = self._dataset_path.name if isinstance(embodiment_tag, EmbodimentTag): self.tag = embodiment_tag.value else: self.tag = embodiment_tag self._metadata = self._get_metadata(EmbodimentTag(self.tag)) self._trajectory_ids, self._trajectory_lengths = self._get_trajectories() self._all_steps = self._get_all_steps(single_base_index=single_base_index) self._modality_keys = self._get_modality_keys() self._delta_indices = self._get_delta_indices() self.set_transforms_metadata(self.metadata) self.set_epoch(0) print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}") # LeRobot-specific config self._lerobot_modality_meta = self._get_lerobot_modality_meta() self._lerobot_info_meta = self._get_lerobot_info_meta() self._data_path_pattern = self._get_data_path_pattern() self._video_path_pattern = self._get_video_path_pattern() self._chunk_size = self._get_chunk_size() self._tasks = self._get_tasks() self.curr_traj_data = None self.curr_traj_id = None # Check if the dataset is valid self._check_integrity() @property def dataset_path(self) -> Path: """The path to the dataset that contains the METADATA_FILENAME file.""" return self._dataset_path @property def metadata(self) -> DatasetMetadata: """The metadata for the dataset, loaded from metadata.json in the dataset directory""" return self._metadata @property def trajectory_ids(self) -> np.ndarray: """The trajectory IDs in the dataset, stored as a 1D numpy array of strings.""" return self._trajectory_ids @property def trajectory_lengths(self) -> np.ndarray: """The trajectory lengths in the dataset, stored as a 1D numpy array of integers. The order of the lengths is the same as the order of the trajectory IDs. """ return self._trajectory_lengths @property def all_steps(self) -> list[tuple[int, int]]: """The trajectory IDs and base indices for all steps in the dataset. Example: self.trajectory_ids: [0, 1, 2] self.trajectory_lengths: [3, 2, 4] return: [ ("traj_0", 0), ("traj_0", 1), ("traj_0", 2), ("traj_1", 0), ("traj_1", 1), ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3) ] """ return self._all_steps @property def modality_keys(self) -> dict: """The modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality. Example: { "video": ["video.image_side_0", "video.image_side_1"], "state": ["state.eef_position", "state.eef_rotation"], "action": ["action.eef_position", "action.eef_rotation"], "language": ["language.human.task"], "timestamp": ["timestamp"], "reward": ["reward"], } """ return self._modality_keys @property def delta_indices(self) -> dict[str, np.ndarray]: """The delta indices for the dataset. The keys are the modality.key, and the values are the delta indices for each modality.key.""" return self._delta_indices @property def dataset_name(self) -> str: """The name of the dataset.""" return self._dataset_name @property def lerobot_modality_meta(self) -> LeRobotModalityMetadata: """The metadata for the LeRobot dataset.""" return self._lerobot_modality_meta @property def lerobot_info_meta(self) -> dict: """The metadata for the LeRobot dataset.""" return self._lerobot_info_meta @property def data_path_pattern(self) -> str: """The path pattern for the LeRobot dataset.""" return self._data_path_pattern @property def video_path_pattern(self) -> str: """The path pattern for the LeRobot dataset.""" return self._video_path_pattern @property def chunk_size(self) -> int: """The chunk size for the LeRobot dataset.""" return self._chunk_size @property def tasks(self) -> pd.DataFrame: """The tasks for the dataset.""" return self._tasks def _get_metadata(self, embodiment_tag: EmbodimentTag) -> DatasetMetadata: """Get the metadata for the dataset. Returns: dict: The metadata for the dataset. """ # 1. Modality metadata modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME if not (modality_meta_path.exists()): modality_meta_path = Path( "/mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/modality.json" ) print( "WARNING: Could not find modality.json in dataset path, falling back to /mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/modality.json" ) assert modality_meta_path.exists(), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" # 1.1. State and action modalities simplified_modality_meta: dict[str, dict] = {} with open(modality_meta_path, "r") as f: le_modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) for modality in ["state", "action"]: simplified_modality_meta[modality] = {} le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr(le_modality_meta, modality) for subkey in le_state_action_meta: state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype) if np.issubdtype(state_action_dtype, np.floating): continuous = True else: continuous = False simplified_modality_meta[modality][subkey] = { "absolute": le_state_action_meta[subkey].absolute, "rotation_type": le_state_action_meta[subkey].rotation_type, "shape": [le_state_action_meta[subkey].end - le_state_action_meta[subkey].start], "continuous": continuous, } # 1.2. Video modalities le_info_path = self.dataset_path / LE_ROBOT_INFO_FILENAME assert le_info_path.exists(), f"Please provide a {LE_ROBOT_INFO_FILENAME} file in {self.dataset_path}" with open(le_info_path, "r") as f: le_info = json.load(f) simplified_modality_meta["video"] = {} for new_key in le_modality_meta.video: original_key = le_modality_meta.video[new_key].original_key if original_key is None: original_key = new_key if original_key in le_info["features"]: le_video_meta = le_info["features"][original_key] else: le_video_meta = le_info["features"][f"observation.images.{original_key}"] # try: # le_video_meta = le_info["features"][original_key] # except: # le_video_meta = le_info["features"][f"observation.images.{original_key}"] height = le_video_meta["shape"][le_video_meta["names"].index("height")] width = le_video_meta["shape"][le_video_meta["names"].index("width")] # NOTE(FH): different lerobot dataset versions have different keys for the number of channels and fps try: channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] fps = le_video_meta["video_info"]["video.fps"] except ValueError: channels = le_video_meta["shape"][le_video_meta["names"].index("channels")] fps = le_video_meta["info"]["video.fps"] simplified_modality_meta["video"][new_key] = { "resolution": [width, height], "channels": channels, "fps": fps, } # 2. Dataset statistics stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME if "agibot" in str(stats_path): # if not stats_path.exists(): # print("WARNING: Could not find stats.json in dataset path, falling back to /mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/stats.json") print( "NOTE: Using standard action normalization at /mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/stats.json" ) stats_path = Path("/mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/stats.json") try: with open(stats_path, "r") as f: le_statistics = json.load(f) for stat in le_statistics.values(): if isinstance(stat, int): continue DatasetStatisticalValues.model_validate(stat) except (FileNotFoundError, ValidationError) as e: print(f"Failed to load dataset statistics: {e}") print(f"Calculating dataset statistics for {self.dataset_name}") # Get all parquet files in the dataset paths parquet_files = list((self.dataset_path).glob(LE_ROBOT_DATA_FILENAME)) le_statistics = calculate_dataset_statistics(parquet_files) dataset_statistics = {} for our_modality in ["state", "action"]: dataset_statistics[our_modality] = {} for subkey in simplified_modality_meta[our_modality]: dataset_statistics[our_modality][subkey] = {} state_action_meta = le_modality_meta.get_key_meta(f"{our_modality}.{subkey}") assert isinstance(state_action_meta, LeRobotStateActionMetadata) le_modality = state_action_meta.original_key for stat_name in le_statistics[le_modality]: indices = np.arange( state_action_meta.start, state_action_meta.end, ) stat = np.array(le_statistics[le_modality][stat_name]) dataset_statistics[our_modality][subkey][stat_name] = stat[indices].tolist() # 3. Full dataset metadata metadata = DatasetMetadata( statistics=dataset_statistics, # type: ignore modalities=simplified_modality_meta, # type: ignore embodiment_tag=embodiment_tag, ) return metadata def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: """Get the trajectories in the dataset.""" # Get trajectory lengths, IDs, and whitelist from dataset metadata episode_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME with open(episode_path, "r") as f: episode_metadata = [json.loads(line) for line in f] trajectory_ids = [] trajectory_lengths = [] for episode in episode_metadata: trajectory_ids.append(episode["episode_index"]) trajectory_lengths.append(episode["length"]) return np.array(trajectory_ids), np.array(trajectory_lengths) def _get_all_steps(self, single_base_index=False) -> list[tuple[int, int]]: """Get the trajectory IDs and base indices for all steps in the dataset. Returns: list[tuple[str, int]]: A list of (trajectory_id, base_index) tuples. Example: self.trajectory_ids: [0, 1, 2] self.trajectory_lengths: [3, 2, 4] return: [ ("traj_0", 0), ("traj_0", 1), ("traj_0", 2), ("traj_1", 0), ("traj_1", 1), ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3) ] """ all_steps: list[tuple[int, int]] = [] for trajectory_id, trajectory_length in zip(self.trajectory_ids, self.trajectory_lengths): if single_base_index: all_steps.append((trajectory_id, 0)) else: for base_index in range(trajectory_length): all_steps.append((trajectory_id, base_index)) return all_steps def _get_modality_keys(self) -> dict: """Get the modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality. See property `modality_keys` for the expected format. """ modality_keys = defaultdict(list) for modality, config in self.modality_configs.items(): modality_keys[modality] = config.modality_keys return modality_keys def _get_delta_indices(self) -> dict[str, np.ndarray]: """Restructure the delta indices to use modality.key as keys instead of just the modalities.""" delta_indices: dict[str, np.ndarray] = {} for config in self.modality_configs.values(): for key in config.modality_keys: delta_indices[key] = np.array(config.delta_indices) return delta_indices def _get_lerobot_modality_meta(self) -> LeRobotModalityMetadata: """Get the metadata for the LeRobot dataset.""" modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME if not (modality_meta_path.exists()): modality_meta_path = Path( "/mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/modality.json" ) print( "WARNING: Could not find modality.json in dataset path, falling back to /mnt/amlfs-03/shared/datasets/agibot-beta-converted-0512/agibotworld/modality.json" ) assert modality_meta_path.exists(), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" with open(modality_meta_path, "r") as f: modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) return modality_meta def _get_lerobot_info_meta(self) -> dict: """Get the metadata for the LeRobot dataset.""" info_meta_path = self.dataset_path / LE_ROBOT_INFO_FILENAME with open(info_meta_path, "r") as f: info_meta = json.load(f) return info_meta def _get_data_path_pattern(self) -> str: """Get the data path pattern for the LeRobot dataset.""" return self.lerobot_info_meta["data_path"] def _get_video_path_pattern(self) -> str: """Get the video path pattern for the LeRobot dataset.""" return self.lerobot_info_meta["video_path"] def _get_chunk_size(self) -> int: """Get the chunk size for the LeRobot dataset.""" return self.lerobot_info_meta["chunks_size"] def _get_tasks(self) -> pd.DataFrame: """Get the tasks for the dataset.""" tasks_path = self.dataset_path / LE_ROBOT_TASKS_FILENAME with open(tasks_path, "r") as f: tasks = [json.loads(line) for line in f] df = pd.DataFrame(tasks) return df.set_index("task_index") def _check_integrity(self): """Use the config to check if the keys are valid and detect silent data corruption.""" ERROR_MSG_HEADER = f"Error occurred in initializing dataset {self.dataset_name}:\n" for modality_config in self.modality_configs.values(): for key in modality_config.modality_keys: if key == "lapa_action" or key == "dream_actions": continue # no need for any metadata for lapa actions because it comes normalized # Check if the key is valid try: self.lerobot_modality_meta.get_key_meta(key) except Exception as e: raise ValueError(ERROR_MSG_HEADER + f"Unable to find key {key} in modality metadata:\n{e}") def set_transforms_metadata(self, metadata: DatasetMetadata): """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" self.transforms.set_metadata(metadata) def set_epoch(self, epoch: int): """Set the epoch for the dataset. Args: epoch (int): The epoch to set. """ self.epoch = epoch def __len__(self) -> int: """Get the total number of data points in the dataset. Returns: int: the total number of data points in the dataset. """ return len(self.all_steps) def __str__(self) -> str: """Get the description of the dataset.""" return f"{self.dataset_name} ({len(self)} steps)" def __getitem__(self, index: int) -> dict: """Get the data for a single step in a trajectory. Args: index (int): The index of the step to get. Returns: dict: The data for the step. """ trajectory_id, base_index = self.all_steps[index] return self.transforms(self.get_step_data(trajectory_id, base_index)) def get_step_data(self, trajectory_id: int, base_index: int) -> dict: """Get the RAW data for a single step in a trajectory. No transforms are applied. Args: trajectory_id (int): The name of the trajectory. base_index (int): The base step index in the trajectory. Returns: dict: The RAW data for the step. Example return: { "video": { "video.image_side_0": [B, T, H, W, C], "video.image_side_1": [B, T, H, W, C], }, "state": { "state.eef_position": [B, T, state_dim], "state.eef_rotation": [B, T, state_dim], }, "action": { "action.eef_position": [B, T, action_dim], "action.eef_rotation": [B, T, action_dim], }, } """ data = {} # Get the data for all modalities self.curr_traj_data = self.get_trajectory_data(trajectory_id) for modality in self.modality_keys: # Get the data corresponding to each key in the modality for key in self.modality_keys[modality]: data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) return data def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: """Get the data for a trajectory.""" if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: return self.curr_traj_data else: chunk_index = self.get_episode_chunk(trajectory_id) parquet_path = self.dataset_path / self.data_path_pattern.format( episode_chunk=chunk_index, episode_index=trajectory_id ) assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" return pd.read_parquet(parquet_path) def get_trajectory_index(self, trajectory_id: int) -> int: """Get the index of the trajectory in the dataset by the trajectory ID. This is useful when you need to get the trajectory length or sampling weight corresponding to the trajectory ID. Args: trajectory_id (str): The ID of the trajectory. Returns: int: The index of the trajectory in the dataset. """ trajectory_indices = np.where(self.trajectory_ids == trajectory_id)[0] if len(trajectory_indices) != 1: raise ValueError(f"Error finding trajectory index for {trajectory_id}, found {trajectory_indices=}") return trajectory_indices[0] def get_episode_chunk(self, ep_index: int) -> int: """Get the chunk index for an episode index.""" return ep_index // self.chunk_size def retrieve_data_and_pad( self, array: np.ndarray, step_indices: np.ndarray, max_length: int, padding_strategy: str = "first_last", ) -> np.ndarray: """Retrieve the data from the dataset and pad it if necessary. Args: array (np.ndarray): The array to retrieve the data from. step_indices (np.ndarray): The step indices to retrieve the data for. max_length (int): The maximum length of the data. padding_strategy (str): The padding strategy, either "first" or "last". """ # Get the padding indices front_padding_indices = step_indices < 0 end_padding_indices = step_indices >= max_length padding_positions = np.logical_or(front_padding_indices, end_padding_indices) # Retrieve the data with the non-padding indices # If there exists some padding, Given T step_indices, the shape of the retrieved data will be (T', ...) where T' < T raw_data = array[step_indices[~padding_positions]] assert isinstance(raw_data, np.ndarray), f"{type(raw_data)=}" # This is the shape of the output, (T, ...) if raw_data.ndim == 1: expected_shape = (len(step_indices),) else: expected_shape = (len(step_indices), *array.shape[1:]) # Pad the data output = np.zeros(expected_shape) # Assign the non-padded data output[~padding_positions] = raw_data # If there exists some padding, pad the data if padding_positions.any(): if padding_strategy == "first_last": # Use first / last step data to pad front_padding_data = array[0] end_padding_data = array[-1] output[front_padding_indices] = front_padding_data output[end_padding_indices] = end_padding_data elif padding_strategy == "zero": # Use zero padding output[padding_positions] = 0 else: raise ValueError(f"Invalid padding strategy: {padding_strategy}") return output def get_video_path(self, trajectory_id: int, key: str) -> Path: chunk_index = self.get_episode_chunk(trajectory_id) original_key = self.lerobot_modality_meta.video[key].original_key if original_key is None: original_key = key video_filename = self.video_path_pattern.format( episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key ) if not (self.dataset_path / video_filename).exists(): original_key = f"observation.images.{original_key}" video_filename = self.video_path_pattern.format( episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key ) return self.dataset_path / video_filename def get_video( self, trajectory_id: int, key: str, base_index: int, ) -> np.ndarray: """Get the video frames for a trajectory by a base index. Args: dataset (BaseSingleDataset): The dataset to retrieve the data from. trajectory_id (str): The ID of the trajectory. key (str): The key of the video. base_index (int): The base index of the trajectory. Returns: np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) """ # Get the step indices step_indices = self.delta_indices[key] + base_index # print(f"{step_indices=}") # Get the trajectory index trajectory_index = self.get_trajectory_index(trajectory_id) # Ensure the indices are within the valid range # This is equivalent to padding the video with extra frames at the beginning and end step_indices = np.maximum(step_indices, 0) # if step_indices[-1] >= self.trajectory_lengths[trajectory_index]: # step_indices -= (self.trajectory_lengths[trajectory_index] - step_indices[-1] + 1) step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) assert key.startswith("video."), f"Video key must start with 'video.', got {key}" # Get the sub-key key = key.replace("video.", "") video_path = self.get_video_path(trajectory_id, key) # Get the action/state timestamps for each frame in the video assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" assert "timestamp" in self.curr_traj_data.columns, f"No timestamp found in {trajectory_id=}" timestamp: np.ndarray = self.curr_traj_data["timestamp"].to_numpy() # Get the corresponding video timestamps from the step indices video_timestamp = timestamp[step_indices] try: return get_frames_by_timestamps( video_path.as_posix(), video_timestamp, video_backend=self.video_backend, video_backend_kwargs=self.video_backend_kwargs, ) except Exception: self.video_backend = "torchvision_av" return get_frames_by_timestamps( video_path.as_posix(), video_timestamp, video_backend=self.video_backend, video_backend_kwargs=self.video_backend_kwargs, ) def get_state_or_action( self, trajectory_id: int, modality: str, key: str, base_index: int, ) -> np.ndarray: """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 data is stored in absolute format, pad with the first or last step data; otherwise, pad with zero. Args: dataset (BaseSingleDataset): The dataset to retrieve the data from. trajectory_id (int): The ID of the trajectory. modality (str): The modality of the data. key (str): The key of the data. base_index (int): The base index of the trajectory. Returns: np.ndarray: The data for the trajectory and step indices. """ # Get the step indices step_indices = self.delta_indices[key] + base_index # Get the trajectory index trajectory_index = self.get_trajectory_index(trajectory_id) # Get the maximum length of the trajectory max_length = self.trajectory_lengths[trajectory_index] assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" # Get the sub-key, e.g. state.joint_angles -> joint_angles key = key.replace(modality + ".", "") # Get the lerobot key le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) le_key = le_state_or_action_cfg[key].original_key if le_key is None: le_key = key # Get the data array, shape: (T, D) assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" le_indices = np.arange( le_state_or_action_cfg[key].start, le_state_or_action_cfg[key].end, ) data_array = data_array[:, le_indices] # Get the state or action configuration state_or_action_cfg = getattr(self.metadata.modalities, modality)[key] # Pad the data return self.retrieve_data_and_pad( array=data_array, step_indices=step_indices, max_length=max_length, padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", ) def get_language( self, trajectory_id: int, key: str, base_index: int, ) -> list[str]: """Get the language annotation data for a trajectory by step indices. Args: dataset (BaseSingleDataset): The dataset to retrieve the data from. trajectory_id (int): The ID of the trajectory. key (str): The key of the annotation. base_index (int): The base index of the trajectory. Returns: list[str]: The annotation data for the trajectory and step indices. If no matching data is found, return empty strings. """ assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" # Get the step indices step_indices = self.delta_indices[key] + base_index # Get the trajectory index trajectory_index = self.get_trajectory_index(trajectory_id) # Get the maximum length of the trajectory max_length = self.trajectory_lengths[trajectory_index] # Get the end times corresponding to the closest indices step_indices = np.maximum(step_indices, 0) step_indices = np.minimum(step_indices, max_length - 1) # Get the annotations task_indices: list[int] = [] assert key.startswith("annotation."), f"Language key must start with 'annotation.', got {key}" subkey = key.replace("annotation.", "") annotation_meta = self.lerobot_modality_meta.annotation assert annotation_meta is not None, f"Annotation metadata is None for {subkey}" assert subkey in annotation_meta, ( f"Annotation key {subkey} not found in metadata, available annotation keys: {annotation_meta.keys()}" ) subkey_meta = annotation_meta[subkey] original_key = subkey_meta.original_key if original_key is None: original_key = key for i in range(len(step_indices)): task_indices.append(self.curr_traj_data[original_key][step_indices[i]].item()) return self.tasks.loc[task_indices]["task"].tolist() def get_data_by_modality( self, trajectory_id: int, modality: str, key: str, base_index: int, ): """Get the data corresponding to the modality for a trajectory by a base index. This method will call the corresponding helper method based on the modality. See the helper methods for more details. NOTE: For the language modality, the data is padded with empty strings if no matching data is found. Args: dataset (BaseSingleDataset): The dataset to retrieve the data from. trajectory_id (int): The ID of the trajectory. modality (str): The modality of the data. key (str): The key of the data. base_index (int): The base index of the trajectory. """ if modality == "video": return self.get_video(trajectory_id, key, base_index) elif modality == "state" or modality == "action": return self.get_state_or_action(trajectory_id, modality, key, base_index) elif modality == "language": return self.get_language(trajectory_id, key, base_index) else: raise ValueError(f"Invalid modality: {modality}") class CachedLeRobotSingleDataset(LeRobotSingleDataset): def __init__(self, img_resize: tuple[int, int] | None = None, *args, **kwargs): """ This class caches the video frames for each trajectory and key. It is recommended to use this class if the video frames need to be accessed multiple times. Args: resize_img (tuple[int, int], optional): The size to resize the video frames to reduce memory usage. """ # Convert img_resize to tuple if it is not already if img_resize is not None and not isinstance(img_resize, tuple): img_resize = tuple(img_resize) assert len(img_resize) == 2, f"Expected tuple of length 2, got {img_resize}" self.img_resize = img_resize # Initialize img_resize attribute first to ensure it exists super().__init__(*args, **kwargs) cached_frames: dict[str, np.ndarray] = {} for key in self.modality_keys["video"]: all_frames = [] key = key.replace("video.", "") for trajectory_id, trajectory_length in tqdm( zip(self.trajectory_ids, self.trajectory_lengths), total=len(self.trajectory_ids), desc=f"Caching {key} frames", ): video_path = self.get_video_path(trajectory_id, key) frames = get_all_frames( video_path.as_posix(), video_backend=self.video_backend, video_backend_kwargs=self.video_backend_kwargs, resize_size=img_resize, ) assert frames.ndim == 4, f"Expected 4D array, got {frames.shape} array" assert frames.shape[3] == 3, f"Expected 3 channels, got {frames.shape[3]} channels" # assert ( # frames.shape[0] == trajectory_length # ), f"Expected {trajectory_length} frames, got {frames.shape[0]} frames" all_frames.append(frames) cached_frames[key] = np.concatenate(all_frames, axis=0) print(f"{key}: {cached_frames[key].shape}") self.cached_frames = cached_frames self.start_indices = np.cumsum(self.trajectory_lengths) - self.trajectory_lengths def get_video(self, trajectory_id: int, key: str, base_index: int) -> np.ndarray: step_indices = self.delta_indices[key] + base_index # Get the trajectory index trajectory_index = self.get_trajectory_index(trajectory_id) # Ensure the indices are within the valid range # This is equivalent to padding the video with extra frames at the beginning and end step_indices = np.maximum(step_indices, 0) step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) assert key.startswith("video."), f"Video key must start with 'video.', got {key}" # Get the sub-key key = key.replace("video.", "") # Calculate the absolute indices absolute_indices = self.start_indices[trajectory_index] + step_indices return self.cached_frames[key][absolute_indices] def get_step_data(self, trajectory_id: int, base_index: int) -> dict: """Get the RAW data for a single step. No transforms are applied. Args: trajectory_id (str): The ID of the trajectory. base_index (int): The base index of the step. Returns: dict: The data for the step. """ data = {} self.curr_traj_data = self.get_trajectory_data(trajectory_id) # Get the data for all modalities for modality in self.modality_keys: # Get the data corresponding to each key in the modality for key in self.modality_keys[modality]: data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) return data def set_transforms_metadata(self, metadata: DatasetMetadata): """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" if self.img_resize is not None: all_video_keys = [key for key in self.modality_keys["video"]] for key in metadata.modalities.video: if key in all_video_keys: metadata.modalities.video[key].resolution = self.img_resize super().set_transforms_metadata(metadata) class WrappedLeRobotSingleDataset(LeRobotSingleDataset): def __init__(self, *args, data_split="full", **kwargs): super().__init__(*args, **kwargs) if data_split == "full": pass elif data_split == "train": self._all_steps = self._all_steps[: -len(self) // 20] elif data_split == "test": self._all_steps = self._all_steps[-len(self) // 20 :] print(f"Dataset is split into {data_split} data, with {len(self._all_steps)} steps.") def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: """Get the trajectories in the dataset.""" # Get trajectory lengths, IDs, and whitelist from dataset metadata episode_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME with open(episode_path, "r") as f: episode_metadata = [json.loads(line) for line in f] trajectory_ids = [] trajectory_lengths = [] for episode in episode_metadata: trajectory_ids.append(episode["episode_index"]) trajectory_lengths.append(episode["length"]) return np.array(trajectory_ids), np.array(trajectory_lengths) def __getitem__(self, index: int) -> dict: """Get the data for a single step in a trajectory. Args: index (int): The index of the step to get. Returns: dict: The data for the step. """ try: # pdb.set_trace() trajectory_id, base_index = self.all_steps[index] original_outputs = self.transforms(self.get_step_data(trajectory_id, base_index)) # delta_actions = original_outputs["action"][1:] - original_outputs["action"][:-1] # delta_actions = original_outputs["action"][1:] - original_outputs["action"][[0]] # delta_actions /= torch.linspace(1, len(delta_actions), steps=len(delta_actions))[:, None] def printvideo(videos, filename): t_videos = rearrange(videos, "c f h w -> f h w c") t_videos = t_videos.detach().to(dtype=torch.uint8).cpu().contiguous().numpy() writer = imageio.get_writer(filename, fps=5) for frame in t_videos: writer.append_data(frame) frames = torch.from_numpy(original_outputs["video"]) # frames = torch.from_numpy(original_outputs["video"]) frames = torch.clamp(frames * 255.0, 0, 255).to(torch.uint8) frames = frames.squeeze(1).transpose(0, 1) # printvideo(frames, "example.mp4") text = "" if "annotation.human.coarse_action" in original_outputs: text = original_outputs["annotation.human.coarse_action"][0].split(":")[-1].strip() video_path = { key: self.get_video_path(trajectory_id, key.replace("video.", "")) for key in self.modality_keys["video"] } data = { "__key__": original_outputs["state"], "action": original_outputs["action"], # "action": original_outputs["action"][:-1], # "action": delta_actions, # "action": torch.zeros_like(delta_actions), "video": frames, # "video_path": video_path, "ai_caption": "", "text": text, "t5_text_embeddings": torch.zeros(512, 1024, dtype=torch.bfloat16).cuda(), "t5_text_mask": torch.ones(512, dtype=torch.int64).cuda(), "fps": 4, "image_size": 256 * torch.ones(4).cuda(), "num_frames": 13, "padding_mask": torch.zeros(1, 256, 256).cuda(), } return data except Exception as e: print(f"Error occurred while getting item {index}: {e}") print("Retrying with a random index...") return self.__getitem__(randint(0, len(self) - 1)) class LeRobotDataset(torch.utils.data.Dataset): def __init__( self, num_frames=81, time_division_factor=4, time_division_remainder=1, max_pixels=1920 * 1080, data_file_keys=("video",), image_file_extension=("jpg", "jpeg", "png", "webp"), video_file_extension=("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), repeat=1, args=None, dataset_path=None, data_split="train", embodiment=None, downscaled_res=False, ): if args is not None: # height = args.height # width = args.width max_pixels = args.max_pixels num_frames = args.num_frames data_file_keys = args.data_file_keys.split(",") repeat = args.dataset_repeat dataset_path = args.dataset_path embodiment = args.embodiment downscaled_res = args.downscaled_res self.num_frames = num_frames self.time_division_factor = time_division_factor self.time_division_remainder = time_division_remainder self.max_pixels = max_pixels # self.height = height # self.width = width # self.height_division_factor = height_division_factor # self.width_division_factor = width_division_factor self.data_file_keys = data_file_keys self.image_file_extension = image_file_extension self.video_file_extension = video_file_extension self.repeat = repeat # from gr00t_dreams.data.dataset import WrappedLeRobotSingleDataset # from gr00t_dreams.groot_configs import construct_modality_config_and_transforms from cosmos_policy._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( construct_modality_config_and_transforms, ) self.dataset_path = [] for p in dataset_path.split(","): if (Path(p) / "data").exists() and (Path(p) / "meta").exists() and (Path(p) / "videos").exists(): self.dataset_path.append(p) else: # This is not a LeRobot dataset, assume it's a directory of LeRobot datasets for sub_p in Path(p).iterdir(): if sub_p.is_dir(): self.dataset_path.append(str(sub_p)) self.lerobot_datasets = [] for p in self.dataset_path: config, train_transform, test_transform = construct_modality_config_and_transforms( num_frames=(num_frames + 1), embodiment=embodiment, downscaled_res=downscaled_res ) # Add an additional prefix frame as baseline to compute delta actions self.lerobot_datasets.append( WrappedLeRobotSingleDataset( dataset_path=p, modality_configs=config, transforms=train_transform if data_split == "train" or data_split == "full" else test_transform, embodiment_tag="gr1_unified" if "gr1" in embodiment else embodiment, data_split=data_split, ) ) print(f"Loaded lerobot {data_split} dataset from {self.dataset_path} with {len(self)} samples.") # if height is not None and width is not None: # print("Height and width are fixed. Setting `dynamic_resolution` to False.") # self.dynamic_resolution = False # elif height is None and width is None: # print("Height and width are none. Setting `dynamic_resolution` to True.") # self.dynamic_resolution = True def __getitem__(self, data_id): data_id %= len(self) for dataset in self.lerobot_datasets: if data_id < len(dataset): break data_id -= len(dataset) lerobot_data = dataset[data_id] prompt = lerobot_data["text"] video = lerobot_data["video"] video_frames = [] for i in range(1, video.shape[1]): # Skip first frame (used only as action baseline) frame = video[:, i, :, :] frame = Image.fromarray(frame.permute(1, 2, 0).numpy()) video_frames.append(frame) if len(video_frames) != self.num_frames: print( f"Warning: Expected {self.num_frames} frames, but got {len(video_frames)} frames. Randomly sampling an item instead." ) return self.__getitem__(random.randint(0, len(self) - 1)) # noqa: F821 video_frames = np.stack([np.array(frame, dtype=np.uint8) for frame in video_frames]) # Cumulative baselined delta actions (old version) # NOTE: Need to tweak this after (num_frames + 1) change # delta_actions = lerobot_data["action"][1:] - lerobot_data["action"][[0]] # Chunked cumulative baselined delta actions (for chunked action architecture) actions = lerobot_data["action"] delta_actions = [] for t in range(1, len(actions) - 1, self.time_division_factor): delta_actions.append(actions[t : t + self.time_division_factor] - actions[t - 1]) delta_actions = torch.cat(delta_actions, dim=0) data = { "prompt": prompt, "video": torch.from_numpy(video_frames).permute(3, 0, 1, 2), # "action": torch.from_numpy(delta_actions), "action": (delta_actions), "ai_caption": "", "text": prompt, "t5_text_embeddings": torch.zeros(512, 1024, dtype=torch.bfloat16).cuda(), "t5_text_mask": torch.ones(512, dtype=torch.int64).cuda(), "fps": 4, "image_size": 256 * torch.ones(4).cuda(), "num_frames": 13, "padding_mask": torch.zeros(1, 256, 256).cuda(), "__key__": lerobot_data["__key__"], } # data = { # "prompt": prompt, # "video": video_frames, # "action": delta_actions, # } return data def __len__(self): return sum([len(d) for d in self.lerobot_datasets]) * self.repeat