From ddb26b7189aa9dfa1a5e75de44700f834295a715 Mon Sep 17 00:00:00 2001 From: Jade Date: Mon, 30 Jun 2025 13:11:16 -0400 Subject: [PATCH] add multi --- lerobot/common/datasets/lerobot_dataset.py | 367 ++++++++++--- lerobot/common/datasets/utils_must.py | 581 +++++++++++++++++++++ tests/datasets/test_datasets.py | 46 +- 3 files changed, 915 insertions(+), 79 deletions(-) create mode 100644 lerobot/common/datasets/utils_must.py diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py index 425e9f7e7..47794493e 100644 --- a/lerobot/common/datasets/lerobot_dataset.py +++ b/lerobot/common/datasets/lerobot_dataset.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib +import copy import logging import shutil from pathlib import Path @@ -73,6 +74,28 @@ from lerobot.common.datasets.video_utils import ( get_video_info, ) +# mustafa stuff here +from lerobot.common.datasets.utils_must import ( + reshape_features_to_max_dim, + keep_datasets_with_valid_fps, + keep_datasets_with_the_same_features_per_robot_type, + aggregate_stats_per_robot_type, + create_padded_features, + pad_tensor, + map_dict_keys, + ROBOT_TYPE_KEYS_MAPPING, + OBS_IMAGE, + OBS_IMAGE_2, + OBS_IMAGE_3, + TASKS_KEYS_MAPPING, +) +from lerobot.common.constants import ( + ACTION, + OBS_ENV_STATE, + OBS_STATE, + +) + CODEBASE_VERSION = "v2.1" @@ -83,6 +106,7 @@ class LeRobotDatasetMetadata: root: str | Path | None = None, revision: str | None = None, force_cache_sync: bool = False, + feature_keys_mapping: dict[str, str] | None = None, ): self.repo_id = repo_id self.revision = revision if revision else CODEBASE_VERSION @@ -99,6 +123,14 @@ class LeRobotDatasetMetadata: (self.root / "meta").mkdir(exist_ok=True, parents=True) self.pull_from_repo(allow_patterns="meta/") self.load_metadata() + # added by mshukor + self.feature_keys_mapping = feature_keys_mapping.get(repo_id, None) if feature_keys_mapping else None + self.inverse_feature_keys_mapping = ( + {v: k for k, v in self.feature_keys_mapping.items() if v} if self.feature_keys_mapping else {} + ) + self.info["features"] = map_dict_keys( + self.info["features"], feature_keys_mapping=self.feature_keys_mapping + ) def load_metadata(self): self.info = load_info(self.root) @@ -177,7 +209,15 @@ class LeRobotDatasetMetadata: @property def video_keys(self) -> list[str]: """Keys to access visual modalities stored as videos.""" - return [key for key, ft in self.features.items() if ft["dtype"] == "video"] + # changed + keys = [] + for key, ft in self.features.items(): + key_ = ( + self.inverse_feature_keys_mapping.get(key, key) if self.inverse_feature_keys_mapping else key + ) + if ft["dtype"] == "video": + keys.append(key_) + return keys @property def camera_keys(self) -> list[str]: @@ -342,6 +382,19 @@ class LeRobotDataset(torch.utils.data.Dataset): force_cache_sync: bool = False, download_videos: bool = True, video_backend: str | None = None, + + # new thing by M + feature_keys_mapping: dict[str, str] | None = None, + max_action_dim: int = None, + max_state_dim: int = None, + max_num_images: int = None, + max_image_dim: int = None, + training_features: list | None = None, + discard_first_n_frames: int = 0, + discard_first_idle_frames: bool = False, + motion_threshold: float = 5e-2, + motion_window_size: int = 10, + motion_buffer: int = 3, ): """ 2 modes are available for instantiating this class, depending on 2 different use cases: @@ -455,15 +508,31 @@ class LeRobotDataset(torch.utils.data.Dataset): self.video_backend = video_backend if video_backend else get_safe_default_codec() self.delta_indices = None + # by mshukor + self.training_features = training_features + self.discard_first_n_frames = discard_first_n_frames + self.discard_first_idle_frames = discard_first_idle_frames + self.motion_threshold = motion_threshold + self.motion_window_size = motion_window_size + self.motion_buffer = motion_buffer + # Unused attributes self.image_writer = None self.episode_buffer = None self.root.mkdir(exist_ok=True, parents=True) + # more mshukor + self.feature_keys_mapping = feature_keys_mapping.get(repo_id, None) if feature_keys_mapping else None + self.inverse_feature_keys_mapping = ( + {v: k for k, v in self.feature_keys_mapping.items() if v} if self.feature_keys_mapping else {} + ) + # Load metadata + # TODO: change self.meta = LeRobotDatasetMetadata( - self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync + self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync, + feature_keys_mapping=feature_keys_mapping, ) if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"): episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes] @@ -482,17 +551,62 @@ class LeRobotDataset(torch.utils.data.Dataset): self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes) + # mustafa code + if self.discard_first_n_frames > 0: + print("Discarding first n frames:", self.discard_first_n_frames) + self.subset_frame_ids = [] + for ep_idx in range(self.num_episodes): + from_ = self.episode_data_index["from"][ep_idx] + to_ = self.episode_data_index["to"][ep_idx] + # TODO implement advanced strategy + self.subset_frame_ids += [frame_idx for frame_idx in range(from_ + int(self.fps*self.discard_first_n_frames), to_)] + elif self.discard_first_idle_frames: + print(f"Discarding first idle frames: motion_threshold={self.motion_threshold}, motion_window_size={self.motion_window_size}, motion_buffer={self.motion_buffer}") + self.robot_states = torch.stack(self.hf_dataset[OBS_ROBOT]).numpy() # shape: [T, D] + self.subset_frame_ids = [] + for ep_idx in range(self.num_episodes): + from_ = self.episode_data_index["from"][ep_idx] + to_ = self.episode_data_index["to"][ep_idx] + ep_states = self.robot_states[from_:to_] + velocities = np.linalg.norm(np.diff(ep_states, axis=0), axis=1) + velocities = np.concatenate([[0.0], velocities]) + start_idx = find_start_of_motion(velocities, self.motion_window_size, self.motion_threshold, self.motion_buffer) + self.subset_frame_ids += list(range(from_ + start_idx, to_)) + # Check timestamps - timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy() - episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy() - ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()} - check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s) + # commented TODO: check why + # timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy() + # episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy() + # ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()} + # check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s) # Setup delta_indices if self.delta_timestamps is not None: - check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s) + # TODO: check why commented + # check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s) self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps) + # Mustafa + self.meta.info["features"] = map_dict_keys( + self.meta.info["features"], feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features + ) + self.keys_to_max_dim = { + ACTION: max_action_dim, + OBS_ENV_STATE: max_state_dim, + OBS_STATE: max_state_dim, + OBS_IMAGE: max_image_dim, + OBS_IMAGE_2: max_image_dim, + OBS_IMAGE_3: max_image_dim, + } + self.meta.info["features"] = reshape_features_to_max_dim( + self.meta.info["features"], reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim + ) + self.meta.stats = map_dict_keys(self.meta.stats, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features) + self.robot_type = self.meta.info.get("robot_type", "") + # Override tasks + print(TASKS_KEYS_MAPPING.get(self.repo_id, self.meta.tasks), "previous", self.meta.tasks) + self.meta.tasks = TASKS_KEYS_MAPPING.get(self.repo_id, self.meta.tasks) + def push_to_hub( self, branch: str | None = None, @@ -647,6 +761,7 @@ class LeRobotDataset(torch.utils.data.Dataset): key: [max(ep_start.item(), min(ep_end.item() - 1, idx + delta)) for delta in delta_idx] for key, delta_idx in self.delta_indices.items() } + # FIXME(mshukor): what if we train on multiple datasets with different features padding = { # Pad values outside of current episode range f"{key}_is_pad": torch.BoolTensor( [(idx + delta < ep_start.item()) | (idx + delta >= ep_end.item()) for delta in delta_idx] @@ -670,12 +785,18 @@ class LeRobotDataset(torch.utils.data.Dataset): return query_timestamps + # TODO: changed by mustafa def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict: - return { - key: torch.stack(self.hf_dataset.select(q_idx)[key]) - for key, q_idx in query_indices.items() - if key not in self.meta.video_keys - } + queries = {} + for key, q_idx in query_indices.items(): + if key not in self.meta.video_keys and self.inverse_feature_keys_mapping.get(key, key) not in self.meta.video_keys: + key_ = ( + self.inverse_feature_keys_mapping.get(key, key) + if self.inverse_feature_keys_mapping + else key + ) + queries[key] = torch.stack(self.hf_dataset.select(q_idx)[key_]) + return queries def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]: """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function @@ -699,8 +820,12 @@ class LeRobotDataset(torch.utils.data.Dataset): def __len__(self): return self.num_frames + # changed by mshukor def __getitem__(self, idx) -> dict: + if self.discard_first_n_frames > 0 or self.discard_first_idle_frames: + idx = self.subset_frame_ids[idx] item = self.hf_dataset[idx] + item = map_dict_keys(item, feature_keys_mapping=self.feature_keys_mapping) ep_idx = item["episode_index"].item() query_indices = None @@ -717,15 +842,25 @@ class LeRobotDataset(torch.utils.data.Dataset): video_frames = self._query_videos(query_timestamps, ep_idx) item = {**video_frames, **item} - if self.image_transforms is not None: - image_keys = self.meta.camera_keys - for cam in image_keys: - item[cam] = self.image_transforms(item[cam]) - # Add task as a string task_idx = item["task_index"].item() - item["task"] = self.meta.tasks[task_idx] - + try: + item["task"] = self.meta.tasks[task_idx] + except: + print(self.meta.tasks, task_idx, self.repo_id) + if "robot_type" not in item: + item["robot_type"] = self.robot_type + item = map_dict_keys(item, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features) + # Add padded features + # item = self._add_padded_features(item, self.training_features) + if self.image_transforms is not None: + for cam in item: + if cam in self.meta.camera_keys or ("image" in cam and "is_pad" not in cam): + item[cam] = self.image_transforms(item[cam]) + # Map pad keys + # print(item.keys(), "before") + # item = map_dict_pad_keys(item, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features) + # print(item.keys()) return item def __repr__(self): @@ -1022,54 +1157,161 @@ class MultiLeRobotDataset(torch.utils.data.Dataset): tolerances_s: dict | None = None, download_videos: bool = True, video_backend: str | None = None, + + # add + sampling_weights: list[float] | None = None, + feature_keys_mapping: dict[str, dict[str, str]] | None = None, + max_action_dim: int = None, + max_state_dim: int = None, + max_num_images: int = None, + max_image_dim: int = None, + train_on_all_features: bool = False, + training_features: list | None = None, + discard_first_n_frames: int = 0, + min_fps: int = 1, + max_fps: int = 100, + discard_first_idle_frames: bool = False, + motion_threshold: float = 0.05, + motion_window_size: int = 10, + motion_buffer: int = 3, ): super().__init__() self.repo_ids = repo_ids self.root = Path(root) if root else HF_LEROBOT_HOME - self.tolerances_s = tolerances_s if tolerances_s else dict.fromkeys(repo_ids, 0.0001) + self.tolerances_s = tolerances_s if tolerances_s else {repo_id: 1e-4 for repo_id in repo_ids} # Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which # are handled by this class. - self._datasets = [ - LeRobotDataset( - repo_id, - root=self.root / repo_id, - episodes=episodes[repo_id] if episodes else None, - image_transforms=image_transforms, - delta_timestamps=delta_timestamps, - tolerance_s=self.tolerances_s[repo_id], - download_videos=download_videos, - video_backend=video_backend, - ) - for repo_id in repo_ids - ] + _datasets = [] + datasets_repo_ids = [] + self.sampling_weights = [] + self.training_features = training_features + + sampling_weights = sampling_weights if sampling_weights is not None else [1] * len(repo_ids) + assert len(sampling_weights) == len(repo_ids), ( + "The number of sampling weights must match the number of datasets. " + f"Got {len(sampling_weights)} weights for {len(repo_ids)} datasets." + ) + for i, repo_id in enumerate(repo_ids): + try: + # delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta) + _datasets.append( + LeRobotDataset( + repo_id, + root=self.root / repo_id, + episodes=episodes.get(repo_id, None) if episodes else None, + image_transforms=image_transforms, + delta_timestamps = delta_timestamps.get(repo_id, None) if delta_timestamps else None, + tolerance_s=self.tolerances_s[repo_id], + download_videos=download_videos, + video_backend=video_backend, + feature_keys_mapping=feature_keys_mapping, + training_features=training_features, + discard_first_n_frames=discard_first_n_frames, + discard_first_idle_frames=discard_first_idle_frames, + motion_threshold=motion_threshold, + motion_window_size=motion_window_size, + motion_buffer=motion_buffer, + ) + ) + datasets_repo_ids.append(repo_id) + self.sampling_weights.append(float(sampling_weights[i])) + except Exception as e: + print(f"Failed to load dataset: {repo_id} due to Exception: {e}") + print( + f"Finish loading {len(_datasets)} datasets, with sampling weights: {self.sampling_weights} corresponding to: {datasets_repo_ids}" + ) # Disable any data keys that are not common across all of the datasets. Note: we may relax this # restriction in future iterations of this class. For now, this is necessary at least for being able # to use PyTorch's default DataLoader collate function. + # FIXME(mshukor): apply mapping to unify used keys + self.train_on_all_features = train_on_all_features self.disabled_features = set() - intersection_features = set(self._datasets[0].features) - for ds in self._datasets: - intersection_features.intersection_update(ds.features) - if len(intersection_features) == 0: - raise RuntimeError( - "Multiple datasets were provided but they had no keys common to all of them. " - "The multi-dataset functionality currently only keeps common keys." - ) - for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True): - extra_keys = set(ds.features).difference(intersection_features) - logging.warning( - f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the " - "other datasets." - ) - self.disabled_features.update(extra_keys) + if not self.train_on_all_features: + intersection_features = set(_datasets[0].features) + for ds in _datasets: + intersection_features.intersection_update(ds.features) + if len(intersection_features) == 0: + raise RuntimeError( + "Multiple datasets were provided but they had no keys common to all of them. " + "The multi-dataset functionality currently only keeps common keys." + ) + for repo_id, ds in zip(repo_ids, _datasets, strict=True): + extra_keys = set(ds.features).difference(intersection_features) + logging.warning( + f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the " + "other datasets." + ) + self.disabled_features.update(extra_keys) + union_features = {} + for ds in _datasets: + for k, v in ds.features.items(): + if k not in self.disabled_features: + union_features[k] = v + + if len(union_features) == 0: + raise RuntimeError("Multiple datasets were provided, but no features were found.") self.image_transforms = image_transforms - self.delta_timestamps = delta_timestamps - # TODO(rcadene, aliberts): We should not perform this aggregation for datasets - # with multiple robots of different ranges. Instead we should have one normalization - # per robot. - self.stats = aggregate_stats([dataset.meta.stats for dataset in self._datasets]) + self.delta_timestamps = ( + delta_timestamps.get(repo_id, None) if delta_timestamps else None + ) # delta_timestamps # FIXME(mshukor): last repo? + # self.stats = aggregate_stats(self._datasets) # FIXME(mshukor): stats should be computed per robot type and then the robot type should be passed as input to the model + for ds in _datasets: + ds.meta.info["robot_type"] = ROBOT_TYPE_KEYS_MAPPING.get(ds.repo_id, ds.meta.info["robot_type"]) + ds.robot_type = ds.meta.info["robot_type"] + # In case datasets with the same robot_type have different features + _datasets = keep_datasets_with_valid_fps(_datasets, min_fps=min_fps, max_fps=max_fps) + self._datasets, datasets_maks = keep_datasets_with_the_same_features_per_robot_type(_datasets) + self.sampling_weights = [self.sampling_weights[i] for i in range(len(_datasets)) if datasets_maks[i]] + self.repo_ids = [repo_ids[i] for i in range(len(_datasets)) if datasets_maks[i]] + self.datasets_repo_ids = [datasets_repo_ids[i] for i in range(len(_datasets)) if datasets_maks[i]] + # Compute cumulative sizes for fast indexing + self.cumulative_sizes = np.array( + [0] + list(torch.cumsum(torch.tensor([len(d) for d in self._datasets]), dim=0)) + ) + self.sampling_weights = np.array(self.sampling_weights, dtype=np.float32) + self.stats = aggregate_stats_per_robot_type(self._datasets) + self.meta = copy.deepcopy(self._datasets[0].meta) # FIXME(mshukor): aggregate meta from all datasets + self.meta.info = { + repo_id: ds.meta.info for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False) + } + # self.meta.info["features"] = self._datasets[0].meta.info["features"] # Assume all datasets have the same features + # FIXME(mshukor): pad based on types in case we have more than one state? + self.keys_to_max_dim = { + ACTION: max_action_dim, + OBS_ENV_STATE: max_state_dim, + OBS_STATE: max_state_dim, + OBS_IMAGE: max_image_dim, + OBS_IMAGE_2: max_image_dim, + OBS_IMAGE_3: max_image_dim, + } + # self.meta.info["features"] = reshape_features_to_max_dim(self._datasets[0].meta.info["features"], reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim) + self.meta.info["features"] = reshape_features_to_max_dim( + union_features, reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim + ) + # reshape stats + for robot_type_, stats_ in self.stats.items(): + for feat_key, feat_stats in stats_.items(): + if feat_key in [ACTION, OBS_ENV_STATE, OBS_STATE]: + for k, v in feat_stats.items(): + if k in ["min", "mean"]: + pad_value = 0 + elif k in ["max", "std"]: + pad_value = 1 + else: + continue + self.stats[robot_type_][feat_key][k] = pad_tensor(v, max_size=self.keys_to_max_dim.get(feat_key, -1), pad_dim=-1, pad_value=pad_value) + self.meta.stats = self.stats + # self.meta.info["features"] = aggregate_features(self._datasets) + self.meta.tasks = { + repo_id: ds.meta.tasks for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False) + } + self.meta.episodes = { + repo_id: ds.meta.episodes for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False) + } + self.robot_types = [ds.meta.info["robot_type"] for ds in self._datasets] @property def repo_id_to_index(self): """Return a mapping from dataset repo_id to a dataset index automatically created by this class. @@ -1156,23 +1398,14 @@ class MultiLeRobotDataset(torch.utils.data.Dataset): def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: if idx >= len(self): raise IndexError(f"Index {idx} out of bounds.") - # Determine which dataset to get an item from based on the index. - start_idx = 0 - dataset_idx = 0 - for dataset in self._datasets: - if idx >= start_idx + dataset.num_frames: - start_idx += dataset.num_frames - dataset_idx += 1 - continue - break - else: - raise AssertionError("We expect the loop to break out as long as the index is within bounds.") - item = self._datasets[dataset_idx][idx - start_idx] + dataset_idx = np.searchsorted(self.cumulative_sizes, idx, side="right").item() - 1 + local_idx = (idx - self.cumulative_sizes[dataset_idx]).item() + item = self._datasets[dataset_idx][local_idx] item["dataset_index"] = torch.tensor(dataset_idx) - for data_key in self.disabled_features: + item = create_padded_features(item, self.meta.info["features"]) + for data_key in self.disabled_features: # FIXME(mshukor): not in getitem? if data_key in item: del item[data_key] - return item def __repr__(self): diff --git a/lerobot/common/datasets/utils_must.py b/lerobot/common/datasets/utils_must.py new file mode 100644 index 000000000..147150d40 --- /dev/null +++ b/lerobot/common/datasets/utils_must.py @@ -0,0 +1,581 @@ +""" +Utils function by Mustafa to refactor +""" +import torch +import numpy as np +from lerobot.common.datasets.compute_stats import ( + aggregate_stats +) +from collections import defaultdict +OBS_IMAGE = "observation.image" +OBS_IMAGE_2 = "observation.image2" +OBS_IMAGE_3 = "observation.image3" + +def reshape_features_to_max_dim(features: dict, reshape_dim: int = -1, keys_to_max_dim: dict = {}) -> dict: + """Reshape features to have a maximum dimension of `max_dim`.""" + reshaped_features = {} + for key in features: + if key in keys_to_max_dim and keys_to_max_dim[key] is not None: + reshaped_features[key] = features[key] + shape = list(features[key]["shape"]) + if any([k in key for k in [OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3]]): # Assume square images + shape[-3] = keys_to_max_dim[key] + shape[-2] = keys_to_max_dim[key] + else: + shape[reshape_dim] = keys_to_max_dim[key] + reshaped_features[key]["shape"] = tuple(shape) + else: + reshaped_features[key] = features[key] + return reshaped_features + +def keep_datasets_with_valid_fps( + ls_datasets: list, min_fps: int = 1, max_fps: int = 100 +) -> list: + print(f"Keeping datasets with fps between {min_fps} and {max_fps}. Considering {len(ls_datasets)} datasets.") + for ds in ls_datasets: + if ds.fps < min_fps or ds.fps > max_fps: + print(f"Dataset {ds} has invalid fps: {ds.fps}. Removing it.") + ls_datasets.remove(ds) + print(f"Keeping {len(ls_datasets)} datasets with valid fps.") + return ls_datasets + +def keep_datasets_with_the_same_features_per_robot_type( + ls_datasets: list +) -> list: + """ + Filters datasets to only keep those with consistent feature shapes per robot type. + + Args: + ls_datasets (List): List of datasets, each with a `meta.info['robot_type']` + and `meta.episodes_stats` dictionary. + + Returns: + List: Filtered list of datasets with consistent feature shapes. + """ + robot_types = {ds.meta.info["robot_type"] for ds in ls_datasets} + datasets_to_remove = set() + + for robot_type in robot_types: + # Collect all stats dicts for this robot type + stats_list = [ + ep_stats + for ds in ls_datasets if ds.meta.info["robot_type"] == robot_type + for ep_stats in ds.meta.episodes_stats.values() + ] + if not stats_list: + continue + + # Determine the most common shape for each key + all_keys = {key for stats in stats_list for key in stats} + for ds in ls_datasets: + if ds.meta.info["robot_type"] != robot_type: + continue + for key in all_keys: + shape_counter = defaultdict(int) + + for stats in stats_list: + value = stats.get(key) + if value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray)): # FIXME(mshukor): check all stats; min, mean, max + shape_counter[value["mean"].shape] += 1 + if not shape_counter: + continue + + # Identify the most frequent shape + main_shape = max(shape_counter, key=shape_counter.get) + # Flag datasets that don't match the main shape + # for ds in ls_datasets: + first_ep_stats = next(iter(ds.meta.episodes_stats.values()), None) + if not first_ep_stats: + continue + value = first_ep_stats.get(key) + if value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray)) and value["mean"].shape != main_shape: + datasets_to_remove.add(ds) + break + + # Filter out inconsistent datasets + datasets_maks = [ds not in datasets_to_remove for ds in ls_datasets] + filtered_datasets = [ds for ds in ls_datasets if ds not in datasets_to_remove] + print(f"Keeping {len(filtered_datasets)} datasets. Removed {len(datasets_to_remove)} inconsistent ones. Inconsistent datasets:\n{datasets_to_remove}") + return filtered_datasets, datasets_maks + + + +def aggregate_stats_per_robot_type(ls_datasets) -> dict[str, dict[str, torch.Tensor]]: + """Aggregate stats of multiple LeRobot datasets into multiple set of stats per robot type. + + The final stats will have the union of all data keys from each of the datasets. + + The final stats will have the union of all data keys from each of the datasets. For instance: + - new_max = max(max_dataset_0, max_dataset_1, ...) + - new_min = min(min_dataset_0, min_dataset_1, ...) + - new_mean = (mean of all data) + - new_std = (std of all data) + """ + + robot_types = {ds.meta.info["robot_type"] for ds in ls_datasets} + stats = {robot_type: {} for robot_type in robot_types} + for robot_type in robot_types: + robot_type_datasets = [] + for ds in ls_datasets: + if ds.meta.info["robot_type"] == robot_type: + robot_type_datasets.extend(list(ds.meta.episodes_stats.values())) + # robot_type_datasets = [list(ds.episodes_stats.values()) for ds in ls_datasets if ds.meta.info["robot_type"] == robot_type] + stat = aggregate_stats(robot_type_datasets) + stats[robot_type] = stat + return stats + +def str_to_torch_dtype(dtype_str): + """Convert a dtype string to a torch dtype.""" + mapping = { + "float32": torch.float32, + "int64": torch.int64, + "int16": torch.int16, + "bool": torch.bool, + "video": torch.float32, # Assuming video is stored as uint8 images + } + return mapping.get(dtype_str, torch.float32) # Default to float32 + +def create_padded_features(item: dict, features: dict = {}): + for key, ft in features.items(): + if any([k in key for k in ["cam", "effort", "absolute"]]): # FIXME(mshukor): temporary hack + continue + shape = ft["shape"] + if len(shape) == 3: # images to torch format (C, H, W) + shape = (shape[2], shape[0], shape[1]) + if len(shape) == 1 and shape[0] == 1: # ft with shape are actually tensor(ele) + shape = [] + if key not in item: + dtype = str_to_torch_dtype(ft["dtype"]) + item[key] = torch.zeros(shape, dtype=dtype) + item[f"{key}_padding_mask"] = torch.tensor(0, dtype=torch.int64) + if "image" in key: # FIXME(mshukor): support other observations + item[f"{key}_is_pad"] = torch.BoolTensor([False]) + else: + item[f"{key}_padding_mask"] = torch.tensor(1, dtype=torch.int64) + return item + +ROBOT_TYPE_KEYS_MAPPING = { + "lerobot/stanford_hydra_dataset": "static_single_arm", + "lerobot/iamlab_cmu_pickup_insert": "static_single_arm", + "lerobot/berkeley_fanuc_manipulation": "static_single_arm", + "lerobot/toto": "static_single_arm", + "lerobot/roboturk": "static_single_arm", + "lerobot/jaco_play": "static_single_arm", + "lerobot/taco_play": "static_single_arm_7statedim", +} + +def pad_tensor( + tensor: torch.Tensor, max_size: int, pad_dim: int = -1, pad_value: float = 0.0 +) -> torch.Tensor: + is_numpy = isinstance(tensor, np.ndarray) + if is_numpy: + tensor = torch.tensor(tensor) + pad = max_size - tensor.shape[pad_dim] + if pad > 0: + pad_sizes = (0, pad) # pad right + tensor = torch.nn.functional.pad(tensor, pad_sizes, value=pad_value) + return tensor.numpy() if is_numpy else tensor + +def map_dict_keys(item: dict, feature_keys_mapping: dict, training_features: list = None, pad_key: str = "is_pad") -> dict: + """Maps feature keys from the dataset to the keys used in the model.""" + if feature_keys_mapping is None: + return item + features = {} + for key in item: + if key in feature_keys_mapping: + if feature_keys_mapping[key] is not None: + if training_features is None or feature_keys_mapping[key] in training_features: + features[feature_keys_mapping[key]] = item[key] + else: + if training_features is None or key in training_features or pad_key in key: + features[key] = item[key] + return features + + +TASKS_KEYS_MAPPING = { + "pranavsaroha/so100_legos4": {0: "Pick up the LEGO block and place it in the bowl of the same color as the LEGO block."}, + "pranavsaroha/so100_onelego2": {0: "Pick up the green LEGO block and place it in the green bowl."}, + "jpata/so100_pick_place_tangerine": {0: "Pick up the tangerine and place it."}, + "pranavsaroha/so100_onelego3": {0: "Pick up the green LEGO block and place it in the green bowl."}, + "pranavsaroha/so100_carrot_2": {0: "Pick up a carrot and put it in the bin."}, + "pranavsaroha/so100_carrot_5": {0: "Pick up a carrot and put it in the bin."}, + "pandaRQ/pick_med_1": {0: "Pick up the object and place it in the box."}, + "HITHY/so100_strawberry": {0: "Grasp a strawberry and put it in the bin."}, + "vladfatu/so100_above": {0: "Pick up red object and place it in the box."}, + "koenvanwijk/orange50-1": {0: "Pick up the orange object and but it in the LEGO box. "}, + "koenvanwijk/orange50-variation-2": {0: "Pick up the orange object and but it in the LEGO box. "}, + "FeiYjf/new_GtoR": {0: "Move along the line on the paper from start to end."}, + "CSCSXX/pick_place_cube_1.18": {0: "Pick up the cube and place it in the box."}, + "vladfatu/so100_office": {0: "Pick up the red object and place it in the box."}, + "dragon-95/so100_sorting": {0: "Pick up the object from box A and place it in box B."}, + "dragon-95/so100_sorting_1": {0: "Pick up the object from box A and place it in box B."}, + "nbaron99/so100_pick_and_place4": {0: "Pick up the triangular object and place it on a green sticker."}, + "Beegbrain/pick_place_green_block": {0: "Pick up the green block and place in the red cup."}, + "Ityl/so100_recording2": {0: "Pick up the red cube and place it on top of the blue cube."}, + "dragon-95/so100_sorting_2": {0: "Pick up the object from box A and place it in box B."}, + "dragon-95/so100_sorting_3": {0: "Pick up the object from box A and place it in box B."}, + "aractingi/push_cube_offline_data": {0: "Push the green cube to the yellow sticker."}, + "HITHY/so100_peach3": {0: "Grasp a peach and put it in the bin."}, + "HITHY/so100_peach4": {0: "Grasp a peach and put it on the plate."}, + "shreyasgite/so100_legocube_50": {0: "Grasp a lego block and put it in the bin."}, + "shreyasgite/so100_base_env": {0: "Grasp a lego block and put it in the bin."}, + "triton7777/so100_dataset_mix": { + 0: "Pick up the black tape and place it inside the white tape roll.", + 1: "Pick up the gift miniatures and place them in the black box.", + 2: "Sort the mixed objects into their appropriate categories.", + 3: "Place the pens into the pen holder.", + 4: "Place pens, bottles, and any suitable items into the pen holder, as appropriate.", + 5: "Place the oranges into the yellow basket.", + 6: "Stack the plates and place the cup on top." + }, + "Deason11/Open_the_drawer_to_place_items": {0: "Put the objects in the open drawer."}, + "Deason11/PLACE_TAPE_PUSH_DRAWER": {0: "Place the tape in the drawer and close it. "}, + "NONHUMAN-RESEARCH/SOARM100_TASK_VENDA": {0: "Pick up the object and place it in the box."}, + "mikechambers/block_cup_14": {0: "Grasp a block and put it in a cup."}, + "samsam0510/tooth_extraction_3": {0: "Extract the tooth and put it somewhere."}, + "samsam0510/tooth_extraction_4": {0: "Extarct the molar and put it somewhere."}, + "samsam0510/cube_reorientation_2": {0: "Rotate the object so it aligns with the silhouette on the table."}, + "samsam0510/cube_reorientation_4": {0: "Rotate the object so it aligns with respect to the line on the table."}, + "samsam0510/glove_reorientation_1": {0: "Rotate the glove so the bottom part aligns with the line on the table."}, + "DorayakiLin/so100_pick_charger_on_tissue": {0: "Pick up the charger and put it on the white tissue."}, + "zijian2022/noticehuman3": {0: "Notice human."}, + "liuhuanjim013/so100_th": {0: "Grasp a lego figure and put it in the box."}, + "Bartm3/tape_to_bin": {0: "Grasp a tape and put it in the bin."}, + + # Community dataset v2 + "Chojins/chess_game_009_white": { + 0: "Move the blue chess pieces to the highlighted squares." + }, + "1g0rrr/sam_openpi03": { + 0: "Pick up the cube and place it in the box." + }, + "sihyun77/suho_3_17_1": { + 0: "Grasp a lego block and put it in the bin." + }, + "sihyun77/sihyun_3_17_2": { + 0: "Grasp a lego block and put it in the bin." + }, + "sihyun77/suho_3_17_3": { + 0: "Grasp a lego block and put it in the bin." + }, + "sihyun77/sihyun_3_17_5": { + 0: "Grasp a lego block and put it in the bin." + }, + "Odog16/so100_cube_drop_pick_v1": { + 0: "Pick up the orange cube, release it, and then pick it up again." + }, + "sihyun77/sihyun_main_2": { + 0: "Grasp a lego block and put it in the bin." + }, + "sihyun77/suho_main_2": { + 0: "Grasp a lego block and put it in the bin." + }, + "Bartm3/dice2": { + 0: "Grasp a dice and put it in the bin." + }, + "sihyun77/sihyun_main_3": { + 0: "Grasp a lego block and put it in the bin." + }, + "Loki0929/so100_duck": { + 0: "Grasp red, green, yellow ducks and put them in the box." + }, + "pietroom/holdthis": { + 0: "Hold the object steadily without releasing it." + }, + "pietroom/actualeasytask": { + 0: "Grasp the marker and put it in the plastic box." + }, + "Beegbrain/pick_lemon_and_drop_in_bowl": { + 0: "Pick the yellow lemon and drop it in the red bowl." + }, + "Beegbrain/sweep_tissue_cube": { + 0: "Sweep the red cubes to the right with the tissue." + }, + "zijian2022/321": { + 0: "Grasp a lego block and put it in the bin." + }, + "1g0rrr/sam_openpi_solder1": { + 0: "Bring contact to the pad on the board." + }, + "1g0rrr/sam_openpi_solder2": { + 0: "Bring contact to the pad on the board." + }, + "gxy1111/so100_pick_place": { + 0: "Grasp a toy panda and put it in the cup." + }, + "Odog16/so100_cube_stacking_v1": { + 0: "Stack the cubes in the following order from bottom to top: black, blue, then orange." + }, + "sihyun77/mond_1": { + 0: "Grasp a lego block and put it in the bin." + }, + "andlyu/so100_indoor_1": { + 0: "Locate and grasp the blueberry." + }, + "andlyu/so100_indoor_3": { + 0: "Locate and grasp the blueberry." + }, + "frk2/so100large": { + 0: "Pick up roll of tape and put it in the bin." + }, + "lirislab/sweep_tissue_cube": { + 0: "Sweep the red cubes to the right with the tissue bag." + }, + "lirislab/lemon_into_bowl": { + 0: "Pick the yellow lemon and drop it in the red bowl" + }, + "lirislab/red_cube_into_green_lego_block": { + 0: "Put the red cube on top of the yellow cube." + }, + "lirislab/red_cube_into_blue_cube": { + 0: "Put the red cube on top of the blue cube." + }, + "00ri/so100_battery": { + 0: "Grasp a battery and put it in the bin." + }, + "frk2/so100largediffcam": { + 0: "Pick up roll of tape and put it in the bin" + }, + "FsqZ/so100_1": { + 0: "Put the yellow cube inside the purple box." + }, + "ZGGZZG/so100_drop0": { + 0: "Grasp a ball and put it in the hole." + }, + "Chojins/chess_game_000_white_red": { + 0: "Move the red chess pieces to the highlighted squares." + }, + "smanni/train_so100_fluffy_box": { + 0: "Grasp a small object and place it in the box." + }, + "ganker5/so100_push_20250328": { + 0: "Grasp a lego block and put it in the bin." + }, + "ganker5/so100_dataline_0328": { + 0: "Grasp a lego block and put it in the bin." + }, + "ganker5/so100_color_0328": { + 0: "Grasp a lego block and put it in the bin." + }, + "CrazyYhang/A1234-B-C_mvA2B": { + 0: "Move the top disk from the left column to the middle column." + }, + "RasmusP/so100_Orange2Green": { + 0: "Grasp the orange block and drop it in the box." + }, + "sixpigs1/so100_pick_cube_in_box": { + 0: "Pick up the red cube and put it in the box." + }, + "ganker5/so100_push_20250331": { + 0: "Grasp a lego block and put it in the bin." + }, + "ganker5/so100_dataline_20250331": { + 0: "Grasp a lego block and put it in the bin." + }, + "lirislab/put_caps_into_teabox": { + 0: "Pick the coffee capsule and put it into the top drawer of the teabox" + }, + "lirislab/close_top_drawer_teabox": { + 0: "Close the top drawer of the teabox" + }, + "lirislab/open_top_drawer_teabox": { + 0: "Open the top drawer of the teabox" + }, + "lirislab/unfold_bottom_right": { + 0: "Unfold the bag from bottom right corner" + }, + "lirislab/push_cup_target": { + 0: "Push the red cup to the pink target" + }, + "lirislab/put_banana_bowl": { + 0: "Put the banana into the red bowl" + }, + "Chojins/chess_game_001_blue_stereo": { + 0: "Move the blue chess pieces to the highlighted squares" + }, + "Chojins/chess_game_001_red_stereo": { + 0: "Move the red chess pieces to the highlighted squares" + }, + "ganker5/so100_toy_20250402": { + 0: "Grasp a lego block and put it in the bin." + }, + "Gano007/so100_medic": { + 0: "Grasp a medic box and put it in the bin." + }, + "00ri/so100_battery_bin_center": { + 0: "Grasp a battery and put it in the bin." + }, + "paszea/so100_whale_2": { + 0: "Grasp a whale and put it in the plate." + }, + "lirislab/fold_bottom_right": { + 0: "Fold the bag from the bottom right corner." + }, + "lirislab/put_coffee_cap_teabox": { + 0: "Put the coffee capsule into the top drawer of the teabox." + }, + "therarelab/so100_pick_place_2": { + 0: "Pick a plaster roll and place it to the blue sticker." + }, + "paszea/so100_whale_3": { + 0: "Grasp a whale and put it in the plate." + }, + "paszea/so100_whale_4": { + 0: "Grasp a whale and put it in the plate." + }, + "paszea/so100_lego": { + 0: "Grasp a lego and put it in the basket." + }, + "LemonadeDai/so100_coca": { + 0: "Grasp the Coca-Cola can and orient it upright with the top facing up." + }, + "zijian2022/backgrounda": { + 0: "Grasp a lego block and put it in the bin." + }, + "zijian2022/backgroundb": { + 0: "Grasp a lego block and put it in the bin." + }, + "356c/so100_nut_sort_1": { + 0: "Pick up the steel nuts and sort them by color." + }, + "Mwuqiu/so100_0408_muti": { + 0: "Grasp a yellow duck and put it in the box." + }, + "aimihat/so100_tape": { + 0: "Pick up the tape and put it in the bowl." + }, + "lirislab/so100_demo": { + 0: "Put the banana into the red bowl." + }, + "356c/so100_duck_reposition_1": { + 0: "Grasp the tool and use it to move the duck to the indicated position." + }, + "zijian2022/sort1": { + 0: "Grasp a box and sort it by color: place grey boxes on the left and black boxes on the right." + }, + "weiye11/so100_410_zwy": { + 0: "Pick up the cube and place it on the black circle." + }, + "VoicAndrei/so100_banana_to_plate_only": { + 0: "Pick up the banana and place it on the plate." + }, + "sixpigs1/so100_stack_cube_error": { + 0: "Pick up the red cube and stack it on the green cube with position offset when grasping.", + 1: "Pick up the red cube and stack it on the green cube with gripper error when grasping.", + 2: "Pick up the red cube and stack it on the green cube with position offset when stacking.", + 3: "Pick up the red cube and stack it on the green cube without errors", + }, + "isadev/bougies3": { + 0: "Grab the candle wick by the aluminium plate and place it in the box." + }, + "zijian2022/close3": { + 0: "Grasp a lego block and put it in the bin." + }, + "bensprenger/left_arm_yellow_brick_in_box_v0": { + 0: "Grasp the yellow lego block and put it in the box." + }, + "bensprenger/left_arm_yellow_brick_in_box_with_purple_noise_v0": { + 0: "Grasp a yellow lego block and put it in the bin." + }, + "roboticshack/team16-can-stacking": { + 0: "Grasp the flipped cup and stack it on top of the midpoint between the two other cups to create a tower" + }, + "bensprenger/right_arm_p_brick_in_box_with_y_noise_v0": { + 0: "Grasp the purple lego block and put it in the box." + }, + "pierfabre/pig2": { + 0: "Pick the pig and place it to the right." + }, + "zijian2022/insert2": { + 0: "Grasp a lego block and put it in the bin." + }, + "roboticshack/team-7-right-arm-grasp-tape": { + 0: "Grasp the tape and put it in the box." + }, + "pierfabre/pig3": { + 0: "Pick the pig and place it to the right." + }, + "Jiangeng/so100_413": { + 0: "Pick up the cube and place it on top of the black circle." + }, + "roboticshack/team9-pick_cube_place_static_plate": { + 0: "Pick up the green cube and place on orange plate." + }, + "AndrejOrsula/lerobot_double_ball_stacking_random": { + 0: "Stack the balls on top of each other." + }, + "roboticshack/left-arm-grasp-lego-brick": { + 0: "Grasp the lego brick and put it in the box." + }, + "roboticshack/team-7-left-arm-grasp-motor": { + 0: "Grasp the black motor and put it in the box." + }, + "pierfabre/cow2": { + 0: "Pick the cow and place it to the right." + }, + "pierfabre/sheep": { + 0: "Pick the sheep and place it to the right." + }, + "roboticshack/team9-pick_chicken_place_plate": { + 0: "Pick up the chicken and place on orange plate" + }, + "roboticshack/team13-two-balls-stacking": { + 0: "Stack the balls on top of each other." + }, + "tkc79/so100_lego_box_1": { + 0: "Grasp a lego block and put it in the box." + }, + "pierfabre/rabbit": { + 0: "Pick the rabbit and place it to the right.", + 1: "Pick the rabbit and put it to the right" + }, + "roboticshack/team13-three-balls-stacking": { + 0: "Stack the balls on top of each other." + }, + "pierfabre/horse": { + 0: "Pick the horse and place it to the right." + }, + "pierfabre/chicken": { + 0: "Pick the chicken and place it to the right." + }, + "roboticshack/team16-water-pouring": { + 0: "Pouring water from one cup to another cup" + }, + "ad330/cubePlace": { + 0: "Grasp white cube and place it in the bowl." + }, + "paszea/so100_lego_2cam": { + 0: "Grap lego blocks and put them in the plate." + }, + "bensprenger/chess_game_001_blue_stereo": { + 0: "Move the blue chess pieces to the highlighted squares." + }, + "Mohamedal/put_banana": { + 0: "Put the banana in the red bowl." + }, + "tkc79/so100_lego_box_2": { + 0: "Grasp a lego block and put it in the box." + }, + "samanthalhy/so100_herding_1": { + 0: "Grasp a green tool and herd all the particles to the grey bin." + }, + "jlesein/TestBoulon7": { + 0: "Pick up the bolt and put it on the plate." + }, + # V3 with VLM (Qwen-VL-2.5-instruct) annotation + "satvikahuja/mixer_on_off_new_1": {0: "Press the button on the blender."}, + "andy309/so100_0314_fold_cloths": {0: "fold the cloths, use two cameras, two arms."}, + "jchun/so100_pickplace_small_20250323_120056": {0: "Grasp items from white bowl and place in black tray"}, + "Ofiroz91/so_100_cube2bowl": {0: "placing cube inside a red bawl"}, + "ZCM5115/so100_1210": {0: "picks up the USB cable."}, + "francescocrivelli/carrot_eating": {0: "pick up carrot and bring to mouth"}, + "ZCM5115/so100_2Arm3cameras_movebox": {0: "Pick up the white box from the table."}, + "pranavsaroha/so100_carrot_1": {0: "pick a carrot and put it in the bin"}, + "pranavsaroha/so100_carrot_3": {0: "pick a carrot and put it in the bin"}, + "maximilienroberti/so100_lego_red_box": {0: "Placing the red Lego in the red box bin."}, "pranavsaroha/so100_squishy": {0: "pick a squishy and put it in the bin"}, "rabhishek100/so100_train_dataset": {0: "picks tape and places it in a cup."}, "pranavsaroha/so100_squishy100": {0: "pick a squishy and put it in the bin"}, "pandaRQ/pickmed": {0: "Place the green block on the table."}, "swarajgosavi/act_kikobot_pusht_real": {0: "picks up the red block."}, "pranavsaroha/so100_squishy2colors_1": {0: "pick the squishies and put them in the bins with their corresponding colors"}, "Chojins/chess_game_001_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "jmrog/so100_sweet_pick": {0: "Pick up the candy and place it in the bowl."}, "Chojins/chess_game_002_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "pranavsaroha/so100_squishy2colors_2_new": {0: "pick the squishies and put them in the bins with their corresponding colors"}, "Chojins/chess_game_003_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "Chojins/chess_game_004_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "Chojins/chess_game_005_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "Chojins/chess_game_006_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "Chojins/chess_game_007_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "koenvanwijk/blue52": {0: "places blue block on red LEGO piece."}, "jlitch/so100multicam7": {0: "pick up brick and put in bin"}, "vladfatu/so100_ds": {0: "Pick up the cube and place it in the box."}, "Chojins/chess_game_000_white": {0: "Move the blue chess pieces as indicated by the highlighted squares"}, "satvikahuja/orange_mixer_1": {0: "pick orange and place in mixer"}, "satvikahuja/mixer_on_off": {0: "switch the mixer on or off"}, "satvikahuja/orange_pick_place_new1": {0: "Pick up the orange and place it in the bowl."}, "satvikahuja/mixer_on_off_new": {0: "Adjust the s position."}, "FeiYjf/Makalu_push": {0: "Pick up the blue cube."}, "chmadran/so100_dataset04": {0: "picks the blue block and places it in the red cup."}, "FeiYjf/Maklu_dataset": {0: "Pick up the blue cube and place it on the paper."}, "FeiYjf/new_Dataset": {0: "Pick up the blue cube."}, "satvikahuja/mixer_on_off_new_4": {0: "Place the lid on the blender."}, "CSCSXX/pick_place_cube_1.17": {0: "Pick up the red block and place it in the box."}, "liyitenga/so100_pick_taffy3": {0: "Place the eraser in the container."}, "liyitenga/so100_pick_taffy6": {0: "Pick up the toy and place it in the purple cup."}, "yuz1wan/so100_pickplace": {0: "Pick the pink block and place it in the paper cup."}, "liyitenga/so100_pick_taffy7": {0: "Pick up the toy and place it in the box."}, "swarajgosavi/act_kikobot_block_real": {0: "Pick up the blue cube and place it in the box."}, "SeanLMH/so100_picknplace_v2": {0: "picks up blue cube and places it in yellow box."}, "DimiSch/so100_50ep_2": {0: "Place the yellow object in the bowl."}, "DimiSch/so100_50ep_3": {0: "Pick the yellow button from the table."}, "SeanLMH/so100_picknplace": {0: "Pick up the blue block and place it in the yellow box."}, "nbaron99/so100_pick_and_place2": {0: "picks up the white object."}, "chmadran/so100_dataset08": {0: "places blue block on paper."}, "Ityl/so100_recording1": {0: "Putting the red square onto the yellow piece"}, "ad330/so100_box_pickPlace": {0: "places jar in box."}, "carpit680/giraffe_task": {0: "Grasp a block and put it in the bin."}, "carpit680/giraffe_sock_demo_1": {0: "Grasp a sock off the floor."}, "DimiSch/so100_terra_50_2": {0: "Grasp a lego block and put it in the bin."}, "aractingi/push_cube_offline_data_cropped_resized": {0: "Push the green cube to the yellow sticker"}, "FeiYjf/Test_NNNN": {0: "Pick up the purple cube and move it to the right."}, "HITHY/so100_peach": {0: "Grasp a peach and put it in the bin."}, "zaringleb/so100_cube_4_binary": {0: "Grasp a lego block and put it in the bin."}, "FeiYjf/Grab_Pieces": {0: "places the black object on the table."}, "hegdearyandev/so100_eraser_cup_v1": {0: "picks up the red object."}, "jbraumann/so100_1902": {0: "picks up the yellow ball."}, "zaringleb/so100_cube_5_linear": {0: "Grasp a lego block and put it in the bin."}, "samsam0510/tape_insert_1": {0: "Grasp a red tape and put it on the box."}, "samsam0510/tape_insert_2": {0: "Grasp a red tape and put it in the yellow tape."}, "pengjunkun/so100_push_to_hole": {0: "Push the T into the hole."}, "Deason11/Random_Kitchen": {0: "Pick up the cup and place it on the table."}, "Loki0929/so100_100": {0: "Grasp a rubber duck and put it in the box."}, "speedyyoshi/so100_grasp_pink_block": {0: "Grasp a lego block and put it in the bin."}, "lirislab/green_lego_block_into_mug": {0: "pick the green block and place it in the red cup"}, "kevin510/lerobot-cat-toy-placement": {0: "Grasp the cat toy and put it in the cup."}, "NONHUMAN-RESEARCH/SOARM100_TASK_VENDA_BOX": {0: "Move the cube to the right side of the table."}, "zijian2022/noticehuman5": {0: "picks up the box."}, "zijian2022/noticehuman70": {0: "Stop movement when human encounter testbed.", 1: "Stop movement when human encounter testbed w/ trigger."}, "Bartm3/tape_to_bin": {0: "Grasp a tape and put it in the bin."}, "Pi-robot/barbecue_flip": {0: "Pick up the orange cone and place it on the table."}, "Pi-robot/barbecue_put": {0: "Pick up the stick and place it in the grill."}, "sshh11/so100_orange_50ep_1": {0: "Grasp an orange object and put it in the bin."}, "sshh11/so100_orange_50ep_2": {0: "Grasp an orange object and put it in the bin."}, "DorayakiLin/so100_pick_cube_in_box": {0: "Pick up the red cube and put it in the box."}, "Bartm3/tape_to_bin2": {0: "Grasp a tape and put it in the bin."}, "andy309/so100_0311_1152": {0: "Grasp and put it in the bin."}, "sihyun77/suho_so100": {0: "Grasp a lego block and put it in the bin."}, "sihyun77/si_so100": {0: "Grasp a lego block and put it in the bin."}, "shreyasgite/so100_base_left": {0: "Grasp a lego block and put it in the bin."}, "sihyun77/suho_red": {0: "Grasp a lego block and put it in the bin."}, "liuhuanjim013/so100_block": {0: "Grasp a lego block and put it in the bin."}, "joaoocruz00/so100_makeitD1": {0: "Grasp a lego block and put it in the bin."}, "sihyun77/suho_angel": {0: "Grasp a lego block and put it in the bin."}, "sihyun77/sihyun_king": {0: "Grasp a lego block and put it in the bin."}, "acrampette/third_arm_01": {0: "Pick up the circuit board from the table."}, "Winster/so100_cube": {0: "Grasp a lego block and put it in the bin."}, "1g0rrr/sam_openpi03": {0: "Grasp a blue cube and put it in the gray box."}, "thedevansh/mar16_1336": {0: "Grasp a lego block and put it in the bin."}, "hkphoooey/throw_stuffie": {0: "Grab stuffed animal and throw it on the dot."}, "acrampette/third_arm_02": {0: "Pick up the tie and place it in the box."}, "kumarhans/so100_tape_task": {0: "Grasp a roll of tape and put it over the candle case."}, "Odog16/so100_tea_towel_folding_v1": {0: "Fold tea towel into quarters"}, "pietroom/first_task_short": {0: "Pick up the marker from the box."}, "zijian2022/c0": {0: "Grasp a lego block and put it in the bin based on color.", 1: "Grasp a lego block and put it in the bin."}, "1g0rrr/sam_openpi_solder1": {0: "bring contact to the pad on the board."}, "1g0rrr/sam_openpi_solder2": {0: "bring contact to the pad on the board."}, "bnarin/so100_tic_tac_toe_we_do_it_live": {0: "move tic tac toe as player 2."}, "chmadran/so100_home_dataset": {0: "Grasp a lego block and put it in the bin."}, "baladhurgesh97/so100_final_picking_3": {0: "Grasp a carrot, plastic bottle and put it in respective bins."}, "zaringleb/so100_cube_6_2d": {0: "Grasp a lego block and put it in the bin."}, "ZGGZZG/so100_drop1": {0: "Grasp a cube and put it in the right place."}, "abhisb/so100_51_ep": {0: "Pick up the cube and place it in the box."}, "allenchienxxx/so100Test": {0: "Grasp a lego block and put it in the bin."}, "lizi178119985/so100_jia": {0: "Grasp a lego block and put it in the bin."}, "andrewcole712/so100_tape_bin_place": {0: "Place the tape in the wooden box."}, "Gano007/so100_doliprane": {0: "Grasp a medic box and put it in the bin."}, "XXRRSSRR/so100_v3_num_episodes_50": {0: "Grasp a box and put it in the side."}, "Gano007/so100_gano": {0: "Grasp a box and put it in the bin."}, "paszea/so100_whale_grab": {0: "Grasp a whale and put it in the plate."}, "Clementppr/lerobot_pick_and_place_dataset_world_model": {0: "Grasp a fruit and put it in the cup."}, "RasmusP/so100_dataset50ep": {0: "Grasp a square block and put it in the box."}, "Gano007/so100_second": {0: "Grasp a yellow box and put it in the bin."}, "zaringleb/so100_cude_linear_and_2d_comb": {0: "Grasp a lego block and put it in the bin."}, "zijian2022/digitalfix3": {0: "Grasp a lego block and put it in the bin."}, "sihyun77/mond_13": {0: "Grasp a lego block and put it in the bin."}, "356c/so100_rope_reposition_1": {0: "Grasp rope and reposition."}, "paszea/so100_lego_mix": {0: "Grasp lego blocks and put them in the plate."}, "jiajun001/eraser00_2": {0: "picks tissue paper from box."}, "VoicAndrei/so100_banana_to_plate_rebel_full": {0: "Pick up the banana and place it on the place"}, "isadev/bougies1": {0: "Put the candles in the box."}, "sixpigs1/so100_pick_cube_in_box_error": {0: "Pick up the red cube and put it in the box with position offset when grasping.", 1: "Pick up the red cube and put it in the box with gripper error when grasping.", 2: "Pick up the red cube and put it in the box with position offset when releasing.", 3: "Pick up the red cube and put it in the box without errors."}, "sixpigs1/so100_push_cube_error": {0: "Push the blue cube to the red and white target with position offset when reaching.", 1: "Push the blue cube to the red and white target with position offset when pushing.", 2: "Push the blue cube to the red and white target with gripper error when pushing.", 3: "Push the blue cube to the red and white target without errors."}, "sixpigs1/so100_pull_cube_error": {0: "Pull the yellow cube to the red and white target with position offset when reaching.", 1: "Pull the yellow cube to the red and white target with position offset when pulling.", 2: "Pull the yellow cube to the red and white target with gripper error when pulling.", 3: "Pull the yellow cube to the red and white target without errors."}, "isadev/bougies2": {0: "grab the candle wick and place it in the tray."}, "therarelab/med_dis_rare_6": {0: "places green object in box."}, "sixpigs1/so100_pull_cube_by_tool_error": {0: "Pick up the L-shaped tool and pull the purple cube by the tool with position offset when grasping.", 1: "Pick up the L-shaped tool and pull the purple cube by the tool with rotation offset when grasping.", 2: "Pick up the L-shaped tool and pull the purple cube by the tool with gripper error when grasping.", 3: "Pick up the L-shaped tool and pull the purple cube by the tool with position offset when lowering.", 4: "Pick up the L-shaped tool and pull the purple cube by the tool with rotation offset when lowering.", 5: "Pick up the L-shaped tool and pull the purple cube by the tool with position offset when pulling.", 6: "Pick up the L-shaped tool and pull the purple cube by the tool with gripper error when pulling.", 7: "Pick up the L-shaped tool and pull the purple cube by the tool without errors."}, "sixpigs1/so100_insert_cylinder_error": {0: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf with position offset when grasping.", 1: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf with gripper error when grasping.", 2: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf with rotation offset when uprighting.", 3: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf with position offset when inserting.", 4: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf with choice error when inserting.", 5: "Pick up the cylinder, upright it, and insert it into the middle hole of the shelf without errors."}, "lirislab/guess_who_no_cond": {0: "Place the card in the slot."}, "lirislab/guess_who_lighting": {0: "Pick up the card from the shelf."}, "nguyen-v/so100_press_red_button": {0: "The places the cube in the box."}, "nguyen-v/so100_bimanual_grab_lemon_put_in_box2": {0: "Grab the lemon with the black arm, then give it to the green arm, then place the lemon in the cardboard box with the green arm."}, "nguyen-v/press_red_button_new": {0: "Press the red button with the black arm"}, "nguyen-v/so100_rotate_red_button": {0: "Rotate the red button clockwise with the black arm"}, "roboticshack/team10-red-block": {0: "Pick a red lego block and move it to the right."}, "Cidoyi/so100_all_notes_1": {0: "Connect the cable to the device."}, "roboticshack/team11_pianobot": {0: "Point at the keyboard."}, "roboticshack/team2-guess_who_so100": {0: "Pick up the card from the shelf."}, "roboticshack/team2-guess_who_so100_light": {0: "Place the card in the slot."}, "roboticshack/team2-guess_who_less_ligth": {0: "Pick up the card and place it in the slot."}, "jiajun001/eraser00_3": {0: "Pick up the white object from the table."}, + "Setchii/so100_grab_ball": {0: "Grasp a ball and put it on a goblet."}, + # V4 with VLM annotation + 'ctbfl/sort_battery': {0: 'put the battery into battery_box'}, 'lerobot/aloha_static_screw_driver': {0: 'Pick up the screwdriver with the right arm, hand it over to the left arm then place it into the cup.'}, 'lerobot/aloha_static_candy': {0: 'Pick up the candy and unwrap it.'}, 'lerobot/aloha_mobile_wipe_wine': {0: 'Pick up the wet cloth on the faucet and use it to clean the spilled wine on the table and underneath the glass.'}, 'lerobot/aloha_static_coffee': {0: "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons."}, 'lerobot/aloha_static_towel': {0: 'Pick up a piece of paper towel and place it on the spilled liquid.'}, 'lerobot/aloha_static_vinh_cup': {0: 'Pick up the platic cup with the right arm, then pop its lid open with the left arm.'}, 'lerobot/aloha_static_vinh_cup_left': {0: 'Pick up the platic cup with the left arm, then pop its lid open with the right arm.'}, 'lerobot/aloha_static_ziploc_slide': {0: 'Slide open the ziploc bag.'}, 'lerobot/aloha_static_coffee_new': {0: 'Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray.'}, 'lerobot/aloha_static_cups_open': {0: 'Pick up the plastic cup and open its lid.'}, 'lerobot/aloha_static_pro_pencil': {0: 'Pick up the pencil with the right arm, hand it over to the left arm then place it back onto the table.'}, 'lerobot/aloha_mobile_wash_pan': {0: 'Pick up the pan, rinse it in the sink and then place it in the drying rack.'}, 'lerobot/aloha_mobile_cabinet': {0: 'Open the top cabinet, store the pot inside it then close the cabinet.'}, 'lerobot/aloha_mobile_chair': {0: 'Push the chairs in front of the desk to place them against it.'}, 'lerobot/aloha_mobile_elevator': {0: 'Take the elevator to the 1st floor.'}, 'aliberts/koch_tutorial': {0: 'Pick the Lego block and drop it in the box on the right.'}, 'underctrl/single-block_multi-color_pick-up_50': {0: 'Pick single block of multi-color and drop it in the box on the right.'}, 'underctrl/single-block_blue-color_pick-up_80': {0: 'Pick single block of multi-color and drop it in the box on the right.'}, 'underctrl/mutli-stacked-block_mutli-color_pick-up_80': {0: 'Pick single block of multi-color and drop it in the box on the right.'}, 'underctrl/single-stacked-block_two-color_pick-up_80': {0: 'Pick single block of multi-color and drop it in the box on the right.'}, 'underctrl/single-stacked-block_mutli-color_pick-up_80': {0: 'Pick single block of multi-color and drop it in the box on the right.'}, 'underctrl/handcamera_single_blue': {0: 'Pick the Lego block and drop it in the box on the right.'}, 'cmcgartoll/cube_color_organizer': {0: 'Organize blue cube', 1: 'Organize red cube', 2: 'Organize yellow cube'}, 'T-K-233/koch_k1_pour_shot': {0: 'Place the glass on the table.'}, 'Beegbrain/stack_2_cubes': {0: 'picks up the red block.'}, 'seeingrain/pick_place_lego': {0: 'Place the cube in the basket.'}, 'seeingrain/pick_place_lego_wider_range_richard': {0: 'Pick up the blue cube and place it in the basket.'}, 'seeingrain/pick_place_lego_wider_range_dang': {0: 'Pick up the cube and place it in the basket.'}, 'seeingrain/pick_place_lego_wider_range_dong': {0: 'Pick up the blue object and place it in the basket.'}, 'seeingrain/pick_lego_to_hand': {0: 'places blue object on table.'}, 'seeingrain/pick_place_pink_lego': {0: 'Pick up the red cube and place it in the basket.'}, 'seeingrain/pick_place_pink_lego_few_samples': {0: 'pick_place_pink_lego_few_samples '}, 'seeingrain/one_shot_learning_18episodes': {0: 'Pick up the red block and place it in the basket.'}, 'helper2424/hil-serl-push-circle-classifier': {0: 'Push small circle object to the correct position'}, 'seeingrain/lego_3cameras': {0: 'Pick up the red block and place it in the basket.'}, 'Lugenbott/koch_1225': {0: 'Pick up the blue block and place it in the red box.'}, 'twerdster/koch_training_red': {0: 'Pick up the red block.'}, 'dboemer/koch_50-samples': {0: 'Pick up the red block and place it on top of the yellow box.'}, 'seeingrain/241228_pick_place_2cams': {0: 'Place the cube in the basket.'}, 'Eyas/grab_pink_lighter_10_per_loc': {0: 'Pick up the pink object from the table.'}, 'Eyas/grab_bouillon': {0: 'picks up the box and places it in the box.'}, 'twerdster/koch_new_training_red': {0: 'move red cube into cellotape circle'}, 'andabi/shoes_easy': {0: 'picks up the shoe.'}, 'andabi/D2': {0: 'Pick up the shoe and place it on the table.'}, 'Beegbrain/oc_stack_cubes': {0: 'stack the red cube on the blue cube'}, 'abougdira/cube_target': {0: 'put_the cube on the yellow target'}, 'andabi/D3': {0: 'picks up the shoe.'}, 'andabi/D4': {0: 'picks up the shoe.'}, 'andabi/D5': {0: 'places shoe on table.'}, 'andabi/D6': {0: 'picks up the shoe.'}, 'andabi/D7': {0: 'picks up the shoe.'}, 'jainamit/koch_realcube3': {0: 'pick up the cube real with keyboard input'}, 'jainamit/koch_pickcube': {0: 'Pick up the blue cube and place it in the box.'}, 'andabi/D8': {0: 'picks up the shoe.'}, 'andabi/D9': {0: 'The picks up the paper and places it on the table.'}, 'andabi/D10': {0: 'picks up the shoe.'}, 'andabi/D11': {0: 'The picks up the paper and places it on the table.'}, 'andabi/D12': {0: 'The places the cube in the box.'}, 'andabi/D13': {0: 'picks up shoes.'}, 'andabi/D14': {0: 'picks up shoes.'}, 'andabi/D15': {0: 'picks up the shoe.'}, 'rgarreta/koch_pick_place_lego': {0: 'Pick the Lego block and drop it in the box on the right.'}, 'shin1107/koch_train_block': {0: 'Grasp a block and put it in the hole.'}, 'andabi/D16': {0: 'picks up the shoe.'}, 'TrossenRoboticsCommunity/aloha_fold_tshirt': {0: 'Fold the t-shirt.'}, 'rgarreta/koch_pick_place_lego_v2': {0: 'Pick the Lego block and drop it in the box on the right.'}, '1g0rrr/screw1': {0: 'Grasp a lego block and put it in the bin.'}, 'ncavallo/moss_train_grasp': {0: 'Grasp a lego block and put it in the bin.'}, 'andabi/D17': {0: 'picks up the shoe.'}, 'ma3oun/rpi_squares_1': {0: 'Raspberry Pi 5 squares recording 1'}, 'shin1107/koch_move_block_with_some_shapes': {0: 'Grasp a block and put it in the hole with some shapes.'}, 'jannick-st/classifier': {0: 'Move the blue object to the right side of the table.'}, 'rgarreta/koch_pick_place_lego_v3': {0: 'Pick the Lego block and drop it in the box on the right. Top and wrist cameras.'}, 'TrossenRoboticsCommunity/aloha_stationary_logo_assembly': {0: 'Assemble the Trossen Robotics Logo.'}, 'rgarreta/koch_pick_place_lego_v6': {0: 'Grasp a lego block and put it in the bin.'}, 'Beegbrain/put_green_into_blue_bin': {0: 'Put the green cube into the blue bin'}, 'Beegbrain/put_screwdriver_box': {0: 'Put the screwdriver into the box'}, 'Beegbrain/align_three_pens': {0: 'picks up a pen.'}, 'Beegbrain/stack_green_on_blue_cube': {0: 'Stack the blue cube on top of the green cube'}, 'Beegbrain/align_cubes_green_blue': {0: 'Put the green cubes on the left and the blue cube on the right'}, 'IPEC-COMMUNITY/ucsd_kitchen_dataset_lerobot': {0: 'Turn on the faucet', 1: 'Put the bowl inside the kitchen cabinet', 2: 'Open the oven door', 3: 'Place the teapot on the stove', 4: 'Put the white box into the sink', 5: 'Open the carbinet door', 6: 'Put the green box into the sink', 7: 'Put the canned spam into the sink'}, 'dkdltu1111/omx-bottle1': {0: 'Grasp a lego block and put it in the bin.'}, 'rgarreta/koch_pick_place_lego_v7': {0: 'Grasp a lego block and put it in the bin.'}, 'rgarreta/koch_pick_place_lego_v8': {0: 'Grasp a lego block and put it in the bin.'}, 'IPEC-COMMUNITY/berkeley_mvp_lerobot': {0: 'push wooden cube', 1: 'pick detergent from the sink', 2: 'reach red block', 3: 'pick yellow cube', 4: 'close fridge door', 5: 'pick fruit'}, 'BlobDieKatze/GrabBlocks': {0: 'Grasp a lego block and put it in the bin.'}, 'Yuanzhu/koch_bimanual_grasp_0': {0: 'places the yellow block on the mousepad.'}, 'Yuanzhu/koch_bimanual_grasp_3': {0: 'picks up a yellow block.'}, 'takuzennn/aloha-pick100': {0: 'arm pick pen and put it into the cup'}, 'abbyoneill/pusht': {0: 'Grasp a lego block and put it in the bin.'}, 'ncavallo/moss_train_grasp_new': {0: 'Grasp a lego block and put it in the bin.'}, 'mlfu7/pi0_conversion_no_pad_video': {0: 'pickplace deer greybowl', 1: 'stack red on green', 2: 'stack orange cup to yellow cup', 3: 'put orange cup into yellow cup', 4: 'put push red pen to blue pen', 5: 'put tiger to black bowl', 6: 'put potato in bot to black bowl', 7: 'pick up green triangle', 8: 'put push blue pen to red pen', 9: 'close drawer', 10: 'put push green block to red', 11: 'pickup potato', 12: 'open drawer', 13: 'put closing tongs', 14: 'poke block', 15: 'put push blue cube', 16: 'poke tiger', 17: 'pick red cube into black bowl', 18: 'pick blue cube stack on wood block', 19: 'pick blue cube into grey bowl', 20: 'put red ball in black bowl', 21: 'pick green triangle into pink bowl', 22: 'pick red ball into pink bowl', 23: 'poke green triangle', 24: 'poke grey bowl', 25: 'put blue cube pink bowl', 26: 'put red cube into black bowl', 27: 'put deer into gray bowl', 28: 'put red ball into pink bowl', 29: 'put green triangle into pink bowl', 30: 'put pour from yellow cup into black bowl', 31: 'put pour blue cup into pink bowl', 32: 'put brown cube into gray bowl'}, 'pepijn223/lekiwi_pen': {0: 'Fold the jeans.'}, 'TrossenRoboticsCommunity/aloha_baseline_dataset': {0: 'Pick up a blue block and put it in a green bowl. Baseline dataset for testing'}, 'KeWangRobotics/piper_rl_1': {0: 'Pick up the cube and place it in the box.'}, 'nduque/cam_setup2': {0: 'Grasp a green block and put it in the bin.'}, 'KeWangRobotics/piper_rl_1_cropped_resized': {0: 'picks up the cube.'}, 'abbyoneill/new_dataset_pick_place': {0: 'Grasp a lego block and put it in the bin.'}, 'abbyoneill/thurs1120pickplace': {0: 'Grasp a lego block and put it in the bin.'}, 'abbyoneill/data_w_mug': {0: 'Grasp a lego block and put it in the bin.'}, 'pepijn223/lekiwi_drive_in_circle': {0: 'Pick up the red object and place it on the table.'}, 'pepijn223/lekiwi_block_cleanup2': {0: 'Put red block in black box'}, 'KeWangRobotics/piper_rl_2': {0: 'Move the cube to the right side of the table.'}, 'KeWangRobotics/piper_rl_2_cropped_resized': {0: 'Move the block to the right side of the table.'}, 'hannesill/koch_pnp_simple_50': {0: 'Grasp a small block with a specific orientation and put it in the bin with a specific position and orientation.'}, 'KeWangRobotics/piper_rl_3': {0: 'Pick up the wooden block.'}, 'KeWangRobotics/piper_rl_3_cropped_resized': {0: 'Place the block in the box.'}, 'hannesill/koch_pnp_2_blocks_2_bins_200': {0: 'Grasp the blue block first and put it in the first bin that has a specific position and orientation. Then grasp the white block and put it in the second bin that has a specific position and orientation.'}, 'ellen2imagine/pusht_green1': {0: 'Place the green block in the box.'}, 'imatrixlee/koch_place': {0: 'Pick up the white object and place it on the table.'}, 'ellen2imagine/pusht_green_same_init2': {0: 'Place the green block in the correct position.'}, 'KeWangRobotics/piper_rl_4': {0: 'Move the block slightly.'}, 'KeWangRobotics/piper_rl_4_cropped_resized': {0: 'picks the wooden block.'}, 'lalalala0620/koch_blue_paper_tape': {0: 'Grasp a blue paper tape and put it in the bin.'}, 'nduque/act_50_ep': {0: 'Grasp a green block and put it in the bin.'}, 'nduque/act_50_ep2': {0: 'Grasp a green block and put it in the bin.'}, 'HWJ658970/fat_fish': {0: 'Grasp a fat fish toy and put it in the bin.'}, 'Beegbrain/put_red_triangle_green_rect': {0: 'Put the red triangle on top of the green rectangle'}, 'ncavallo/moss_train_gc_block': {0: 'Grasp a lego block and put it in the bin.'}, 'takuzennn/square3': {0: 'Pick the cube from the table.'}, 'HWJ658970/lego_50': {0: 'Grasp a yellow lego block and put it in the bin.'}, 'Deason11/mobile_manipulator_0319': {0: 'Grasp a lego block and put it in the bin.'}, 'Gongsta/grasp_duck_in_cup': {0: 'Grasp the rubber duck and put it in the cup.'}, 'nduque/robustness_e2': {0: 'Grasp a green dice and put it in the bin.'}, 'ibru/bobo_trash_collector': {0: 'Bobo Trash collect and place it in a bin'}, 'Beegbrain/moss_open_drawer_teabox': {0: 'Open the top drawer of the teabox'}, 'Beegbrain/moss_put_cube_teabox': {0: 'Put the green cube in the top drawer of the teabox'}, 'Beegbrain/moss_close_drawer_teabox': {0: 'Close the top drawer of the teabox'}, 'Beegbrain/moss_stack_cubes': {0: 'Stack the green cube on top of the blue cube'}, 'HWJ658970/lego_50_camera_change': {0: 'Grasp a yellow lego block and put it in the bin.'}, 'HWJ658970/lego_100_class': {0: 'Separate yellow and white Lego blocks and place them into the bin.'}, 'nduque/robustness_e3': {0: 'Grasp a green dice and put it in the bin.'}, 'nimitvasavat/Gr00t_lerobot': {0: 'Place the cereal box on the shelf.'}, 'Deason11/mobile_manipulator_0326': {0: 'Grasp a lego block and put it in the bin.', 1: 'mobile_lekiwi.'}, 'KeWangRobotics/piper_push_cube_gamepad_1': {0: 'push the cube to the black area'}, 'KeWangRobotics/piper_push_cube_gamepad_1_cropped_resized': {0: 'push the cube to the black area'}, 'jannick-st/push-cube-classifier_cropped_resized': {0: 'Close the cabinet door.'}, 'nimitvasavat/Gr00t_lerobotV2': {0: 'Place the chocolate chip cookie dough box on the table.'}, 'Zhaoting123/koch_cleanDesk_': {0: 'Grasp a card and use it to clean the desk'}, 'arclabmit/koch_gear_and_bin': {0: 'Pick the gear and place it in the bin.'}, 'Allen-488/koch_dataset_50': {0: 'Grasp a block and put it in the bin.'}, 'dop0/koch_pick_terminal': {0: 'Pick up the terminal and place on the cover.'}, 'nduque/robustness_e4': {0: 'Grasp a green dice and put it in the bin.'}, 'arclabmit/Koch_twoarms': {0: 'official two arms recordings10'}, 'nimitvasavat/Gr00t_lerobot_state_action': {0: 'Place the chocolate chip cookie dough box on the table.'}, 'zliu157/i3r': {0: 'Grasp a lego block and put it in the bin.'}, 'hangwu/koch_pick_terminal': {0: 'Pick up the terminal and place on the cover.'}, 'nduque/robustness_e5': {0: 'Grasp a green dice and put it in the bin.'}, 'zliu157/i3r2': {0: 'Grasp a i3r logo and put it in the bin.'}, 'HuaihaiLyu/groceries': {0: 'Pick the brown long bread and Egg yolk pasry into package'}, 'zliu157/i3r3': {0: 'Grasp a i3r logo and put it in the bin.'}, 'hangwu/piper_pick_terminal_and_place': {0: 'Grasp a terminal and put it on the black box.'}, 'hangwu/piper_pick_terminal_2': {0: 'Grasp the white terminal and put it on the green lid.'}, 'engineer0002/pepper': {0: 'Place the bottle on the table.'}, 'theo-michel/lekiwi_v2': {0: 'Pick up the can on the ground'}, 'theo-michel/lekiwi_v5': {0: 'Pick up the can on the ground'}, 'roboticshack/sandee-kiwiv10': {0: 'Place the bottle on the table.'}, 'ibru/bob_jetson': {0: 'Drive forward pickup the object and put it in the red box and drive back.'}, 'ibru/bobo_jetson': {0: 'Drive forward pickup the object and put it in the red box and drive back.', 1: 'Driver forward'}, 'zliu157/i3r5': {0: 'Grasp a i3r logo and put it in the bin.'}, 'Dongkkka/cable_pick_and_place2': {0: 'Put a black charging cable in a black bowl and put a red charging cable in a green bowl'}, + + +} diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index b4fca77c1..4bd0baccb 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -396,35 +396,57 @@ def test_factory(env_name, repo_id, policy_name): # TODO(alexander-soare): If you're hunting for savings on testing time, this takes about 5 seconds. @pytest.mark.skip("TODO after fix multidataset") def test_multidataset_frames(): - """Check that all dataset frames are incorporated.""" - # Note: use the image variants of the dataset to make the test approx 3x faster. - # Note: We really do need three repo_ids here as at some point this caught an issue with the chaining - # logic that wouldn't be caught with two repo IDs. + """Check that all dataset frames are incorporated and aligned correctly.""" repo_ids = [ "lerobot/aloha_sim_insertion_human_image", "lerobot/aloha_sim_transfer_cube_human_image", "lerobot/aloha_sim_insertion_scripted_image", ] + + # dummy padding dimensions (simulate training setup) + MAX_ACTION_DIM = 14 + MAX_STATE_DIM = 30 + MAX_NUM_IMAGES = 3 + MAX_IMAGE_DIM = 224 + sub_datasets = [LeRobotDataset(repo_id) for repo_id in repo_ids] - dataset = MultiLeRobotDataset(repo_ids) + dataset = MultiLeRobotDataset( + repo_ids, + max_action_dim=MAX_ACTION_DIM, + max_state_dim=MAX_STATE_DIM, + max_num_images=MAX_NUM_IMAGES, + max_image_dim=MAX_IMAGE_DIM, + ) + assert len(dataset) == sum(len(d) for d in sub_datasets) assert dataset.num_frames == sum(d.num_frames for d in sub_datasets) assert dataset.num_episodes == sum(d.num_episodes for d in sub_datasets) - # Run through all items of the LeRobotDatasets in parallel with the items of the MultiLerobotDataset and - # check they match. expected_dataset_indices = [] for i, sub_dataset in enumerate(sub_datasets): expected_dataset_indices.extend([i] * len(sub_dataset)) - for expected_dataset_index, sub_dataset_item, dataset_item in zip( + for expected_dataset_index, sub_item, multi_item in zip( expected_dataset_indices, chain(*sub_datasets), dataset, strict=True ): - dataset_index = dataset_item.pop("dataset_index") + dataset_index = multi_item.pop("dataset_index") assert dataset_index == expected_dataset_index - assert sub_dataset_item.keys() == dataset_item.keys() - for k in sub_dataset_item: - assert torch.equal(sub_dataset_item[k], dataset_item[k]) + + # we ignore padding_mask and dataset_index keys in multi_item + extra_keys = {k for k in multi_item if "padding_mask" in k} + filtered_multi_keys = set(multi_item.keys()) - extra_keys + assert set(sub_item.keys()) == filtered_multi_keys, f"mismatch in keys" + + for k in sub_item: + if k not in multi_item: + continue + v1, v2 = sub_item[k], multi_item[k] + if isinstance(v1, torch.Tensor) and isinstance(v2, torch.Tensor): + assert torch.equal(v1, v2), f"tensor mismatch on key: {k}" + else: + assert v1 == v2, f"value mismatch on key: {k}" + + # TODO(aliberts): Move to more appropriate location