Dataloader¶

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 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 # just for action base data
    self.curr_traj_data = self.get_trajectory_data(trajectory_id)
    # TODO: The logic below is poorly implemented. Data reading should be directly based on curr_traj_data.
    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)
    data = self._apply_action_mode(data)
    return data

def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame:
    """Get the data for a trajectory."""
    if self._lerobot_version == "v2.0":

        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}"
            df = _safe_read_parquet(parquet_path)
            # Apply filter (e.g., remove is_demo=True rows)
            if trajectory_id in self._traj_row_masks:
                mask = self._traj_row_masks[trajectory_id]
                df = df[mask].reset_index(drop=True)
            return df
    elif self._lerobot_version == "v3.0":
        return self.get_trajectory_data_lerobot_v3(trajectory_id)

def get_trajectory_data_lerobot_v3(self, trajectory_id: int) -> pd.DataFrame:
    """Get the data for a trajectory from lerobot v3."""
    if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None:
        return self.curr_traj_data
    else: #TODO check detail later
        episode_meta = self.trajectory_ids_to_metadata[trajectory_id]
        chunk_index = episode_meta["data/chunk_index"]
        file_index = self.get_episode_file_index(trajectory_id)
        # file_from_index = self.get_episode_file_from_index(trajectory_id)


        parquet_path = self.dataset_path / self.data_path_pattern.format(
            chunk_index=chunk_index, file_index=file_index
        )
        assert parquet_path.exists(), f"Parquet file not found at {parquet_path}"
        file_data = _safe_read_parquet(parquet_path)

        # filter by trajectory_id
        episode_data = file_data.loc[file_data["episode_index"] == trajectory_id].copy()
        return episode_data

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 get_episode_file_index(self, ep_index: int) -> int:
    """Get the file index for an episode index."""
    episode_meta = self.trajectory_ids_to_metadata[ep_index]
    return episode_meta["data/file_index"]

def get_episode_file_from_index(self, ep_index: int) -> int:
    """Get the file from index for an episode index."""
    episode_meta = self.trajectory_ids_to_metadata[ep_index]
    return episode_meta["data/file_from_index"]

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(
    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)
    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.", "")

    # Check if this is image-type data (PNG bytes in parquet)
    if key in self._image_type_video_keys:
        assert self.curr_traj_data is not None
        return self._read_images_from_parquet(trajectory_id, key, step_indices)

    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]
    if self._lerobot_version == "v3.0":
        episode_meta = self.trajectory_ids_to_metadata.get(trajectory_id, {})
        from_timestamps = episode_meta.get("videos/from_timestamps", {})
        original_video_key = self.lerobot_modality_meta.video[key].original_key
        if original_video_key is None:
            original_video_key = key
        from_timestamp = float(from_timestamps.get(original_video_key, 0.0))
        video_timestamp = video_timestamp + from_timestamp

    # V2: batch-read next frame in same video file open
    next_frame_offset = None
    if (getattr(self, '_v2_next_frame_index', None) is not None
            and not key.startswith("wrist")):
        next_step = np.array([self._v2_next_frame_index])
        next_step = np.maximum(next_step, 0)
        next_step = np.minimum(next_step, self.trajectory_lengths[trajectory_index] - 1)
        next_ts = timestamp[next_step]
        if self._lerobot_version == "v3.0" and 'from_timestamp' in dir():
            next_ts = next_ts + from_timestamp
        next_frame_offset = len(video_timestamp)
        video_timestamp = np.concatenate([video_timestamp, next_ts])

    frames = get_frames_by_timestamps(
        video_path.as_posix(),
        video_timestamp,
        video_backend=self.video_backend,
        video_backend_kwargs=self.video_backend_kwargs,
    )

    if next_frame_offset is not None:
        self._v2_cached_next_frame = frames[next_frame_offset:]
        return frames[:next_frame_offset]

    return frames

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 key {le_key} is{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",
        # padding_strategy="zero",           # HACK for realdata
    )