Merge branch 'pr-1451' into danaaubakirova/25_06_2025

src/lerobot/datasets/factory.py

@@ -32,6 +32,7 @@ IMAGENET_STATS = {
     "std": [[[0.229]], [[0.224]], [[0.225]]],  # (c,1,1)
 }
 
+from lerobot.common.datasets.utils_must import (EPISODES_DATASET_MAPPING, TRAINING_FEATURES, FEATURE_KEYS_MAPPING)
 
 def resolve_delta_timestamps(
     cfg: PreTrainedConfig, ds_meta: LeRobotDatasetMetadata
@@ -66,7 +67,7 @@ def resolve_delta_timestamps(
     return delta_timestamps
 
 
-def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
+def make_dataset1(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
     Args:
@@ -116,3 +117,96 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
                 dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
 
     return dataset
+
+def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
+    """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
+
+    Args:
+        cfg (TrainPipelineConfig): A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.
+
+    Raises:
+        NotImplementedError: The MultiLeRobotDataset is currently deactivated.
+
+    Returns:
+        LeRobotDataset | MultiLeRobotDataset
+    """
+    image_transforms = (
+        ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
+    )
+    if "," in cfg.dataset.repo_id:
+        repo_id = cfg.dataset.repo_id.split(",")
+        repo_id = [r for r in repo_id if r]
+    else:
+        repo_id = cfg.dataset.repo_id
+    sampling_weights = cfg.dataset.sampling_weights.split(",") if cfg.dataset.sampling_weights else None
+    feature_keys_mapping = FEATURE_KEYS_MAPPING
+    if isinstance(repo_id, str):
+        revision = getattr(cfg.dataset, "revision", None)
+        ds_meta = LeRobotDatasetMetadata(
+            cfg.dataset.repo_id,
+            local_files_only=cfg.dataset.local_files_only,
+            feature_keys_mapping=feature_keys_mapping,
+            revision=revision,
+        )
+        delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
+        dataset = LeRobotDataset(
+            cfg.dataset.repo_id,
+            root=getattr(cfg.dataset, "root", None),
+            episodes=cfg.dataset.episodes,
+            delta_timestamps=delta_timestamps,
+            image_transforms=image_transforms,
+            revision=revision,
+            video_backend=cfg.dataset.video_backend,
+            local_files_only=cfg.dataset.local_files_only,
+            feature_keys_mapping=feature_keys_mapping,
+            max_action_dim=cfg.dataset.max_action_dim,
+            max_state_dim=cfg.dataset.max_state_dim,
+            max_num_images=cfg.dataset.max_num_images,
+            max_image_dim=cfg.dataset.max_image_dim,
+        )
+    else:
+        delta_timestamps = {}
+        episodes = {}
+        for i in range(len(repo_id)):
+            ds_meta = LeRobotDatasetMetadata(
+                repo_id[i],
+                local_files_only=cfg.dataset.local_files_only,
+                feature_keys_mapping=feature_keys_mapping,
+            )  # FIXME(mshukor): ?
+            delta_timestamps[repo_id[i]] = resolve_delta_timestamps(cfg.policy, ds_meta)
+            episodes[repo_id[i]] =  EPISODES_DATASET_MAPPING.get(repo_id[i], cfg.dataset.episodes)
+        training_features = TRAINING_FEATURES.get(cfg.dataset.features_version, None)
+        dataset = MultiLeRobotDataset(
+            repo_id,
+            # TODO(aliberts): add proper support for multi dataset
+            episodes=episodes,
+            delta_timestamps=delta_timestamps,
+            image_transforms=image_transforms,
+            video_backend=cfg.dataset.video_backend,
+            local_files_only=cfg.dataset.local_files_only,
+            sampling_weights=sampling_weights,
+            feature_keys_mapping=feature_keys_mapping,
+            max_action_dim=cfg.dataset.max_action_dim,
+            max_state_dim=cfg.dataset.max_state_dim,
+            max_num_images=cfg.dataset.max_num_images,
+            max_image_dim=cfg.dataset.max_image_dim,
+            train_on_all_features=cfg.dataset.train_on_all_features,
+            training_features=training_features,
+            discard_first_n_frames=cfg.dataset.discard_first_n_frames,
+            min_fps=cfg.dataset.min_fps,
+            max_fps=cfg.dataset.max_fps,
+            discard_first_idle_frames=cfg.dataset.discard_first_idle_frames,
+            motion_threshold=cfg.dataset.motion_threshold,
+            motion_window_size=cfg.dataset.motion_window_size,
+            motion_buffer=cfg.dataset.motion_buffer,
+        )
+        logging.info(
+            "Multiple datasets were provided. Applied the following index mapping to the provided datasets: "
+            f"{pformat(dataset.repo_id_to_index , indent=2)}"
+        )
+    if cfg.dataset.use_imagenet_stats:
+        for key in dataset.meta.camera_keys:
+            for stats_type, stats in IMAGENET_STATS.items():
+                dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
+
+    return dataset

src/lerobot/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
 import os
@@ -80,10 +81,28 @@ from lerobot.datasets.video_utils import (
     get_video_info,
 )
 
-# from lerobot.robots.utils import Robot
-from lerobot.configs.datasets import ROBOT_TYPE_KEYS_MAPPING, TASKS_KEYS_MAPPING
-#FIXME: remove this import
-from lerobot.datasets.collators import pad_tensor
+# 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,
+    find_start_of_motion,
+    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"
 LEROBOT_HOME = Path(os.getenv("LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser()
@@ -105,6 +124,7 @@ class LeRobotDatasetMetadata:
         feature_keys_mapping: dict[str, str] | 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.local_files_only = local_files_only
@@ -122,6 +142,7 @@ 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 {}
@@ -129,6 +150,7 @@ class LeRobotDatasetMetadata:
         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)
         check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
@@ -206,7 +228,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]:
@@ -371,6 +401,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:
@@ -484,15 +527,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]
@@ -511,17 +570,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_STATE]).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,
@@ -676,6 +780,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]
@@ -699,12 +804,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
@@ -728,8 +839,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
@@ -746,15 +861,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):
@@ -1035,51 +1160,105 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 
-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
+class MultiLeRobotDatasetMeta:
+    def __init__(
+        self,
+        datasets: list[LeRobotDataset],
+        repo_ids: list[str],
+        keys_to_max_dim: dict[str, int],
+        train_on_all_features: bool = False,
+    ):
+        self.repo_ids = repo_ids
+        self.keys_to_max_dim = keys_to_max_dim
+        self.train_on_all_features = train_on_all_features
+        self.robot_types = [ds.meta.info["robot_type"] for ds in datasets]
 
-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)  
-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
+        # assign robot_type if missing
+        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"]
+
+        # step 1: compute disabled features
+        self.disabled_features = set()
+        if not self.train_on_all_features:
+            intersection = set(datasets[0].features)
+            for ds in datasets:
+                intersection.intersection_update(ds.features)
+            if not intersection:
+                raise RuntimeError("No common features across datasets.")
+            for repo_id, ds in zip(repo_ids, datasets):
+                extra = set(ds.features) - intersection
+                logging.warning(f"Disabling {extra} for repo {repo_id}")
+                self.disabled_features.update(extra)
+
+        # step 2: build union_features excluding disabled
+        self.union_features = {}
+        for ds in datasets:
+            for k, v in ds.features.items():
+                if k not in self.disabled_features:
+                    self.union_features[k] = v
+
+        # step 3: reshape feature schema
+        self.features = reshape_features_to_max_dim(
+            self.union_features, reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim
+        )
+
+        # step 4: aggregate stats
+        self.stats = aggregate_stats_per_robot_type(datasets)
+        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():
+                        pad_value = 0 if k in ["min", "mean"] else 1
+                        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
+                        )
+
+        # step 5: episodes & tasks
+        self.episodes = {
+            repo_id: ds.meta.episodes for repo_id, ds in zip(repo_ids, datasets)
+        }
+        self.tasks = {
+            repo_id: ds.meta.tasks for repo_id, ds in zip(repo_ids, datasets)
+        }
+        self.info = {
+            repo_id: ds.meta.info for repo_id, ds in zip(repo_ids, datasets)
+        }
+
+class MultiLeRobotDatasetCleaner:
+    def __init__(
+        self,
+        datasets: list[LeRobotDataset],
+        repo_ids: list[str],
+        sampling_weights: list[float],
+        datasets_repo_ids: list[str],
+        min_fps: int = 1,
+        max_fps: int = 100,
+    ):
+        self.original_datasets = datasets
+        self.original_repo_ids = repo_ids
+        self.original_weights = sampling_weights
+        self.original_datasets_repo_ids = datasets_repo_ids
+
+        # step 1: remove datasets with invalid fps
+        valid_fps_datasets = keep_datasets_with_valid_fps(datasets, min_fps=min_fps, max_fps=max_fps)
+
+        # step 2: keep datasets with same features per robot type
+        consistent_datasets, keep_mask = keep_datasets_with_the_same_features_per_robot_type(valid_fps_datasets)
+
+        self.cleaned_datasets = consistent_datasets
+        self.keep_mask = keep_mask
+        self.cleaned_weights = [sampling_weights[i] for i in range(len(valid_fps_datasets)) if keep_mask[i]]
+        self.cleaned_repo_ids = [repo_ids[i] for i in range(len(valid_fps_datasets)) if keep_mask[i]]
+        self.cleaned_datasets_repo_ids = [
+            datasets_repo_ids[i] for i in range(len(valid_fps_datasets)) if keep_mask[i]
+        ]
+
+        self.cumulative_sizes = np.array(
+            [0] + list(torch.cumsum(torch.tensor([len(d) for d in consistent_datasets]), dim=0))
+        )
+        self.cleaned_weights = np.array(self.cleaned_weights, dtype=np.float32)
+      
 
 class MultiLeRobotDataset(torch.utils.data.Dataset):
     """A dataset consisting of multiple underlying `LeRobotDataset`s.
@@ -1099,6 +1278,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         download_videos: bool = True,
         local_files_only: bool = False,
         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,
@@ -1118,105 +1299,93 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         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.
-        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)
-
-        self.image_transforms = image_transforms
-        self.delta_timestamps = self.delta_timestamps = delta_timestamps.get(
-            repo_id, None
-        )
-        # 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.
-        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"]
-        #self.stats = aggregate_stats([dataset.meta.stats for dataset in self._datasets])
-        _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): apply mapping to unify used keys
         # 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: max_state_dim,
-            OBS_ROBOT: 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
+        self.image_transforms = image_transforms
+        self.delta_timestamps = (
+            delta_timestamps.get(repo_id, None) if delta_timestamps else None
+        )  # delta_timestamps # FIXME(mshukor): last repo?
+        # In case datasets with the same robot_type have different features
+        cleaner = MultiLeRobotDatasetCleaner(
+            datasets=_datasets,
+            repo_ids=repo_ids,
+            sampling_weights=self.sampling_weights,
+            datasets_repo_ids=datasets_repo_ids,
+            min_fps=min_fps,
+            max_fps=max_fps,
         )
-        # 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, OBS_ROBOT]:
-                    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._datasets = cleaner.cleaned_datasets
+        self.sampling_weights = cleaner.cleaned_weights
+        self.repo_ids = cleaner.cleaned_repo_ids
+        self.datasets_repo_ids = cleaner.cleaned_datasets_repo_ids
+        self.cumulative_sizes = cleaner.cumulative_sizes
+        # 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
+        self.meta = MultiLeRobotDatasetMeta(
+            datasets=self._datasets,
+            repo_ids=self.repo_ids,
+            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,
+            },
+            train_on_all_features=train_on_all_features,
+        )
+        self.disabled_features = self.meta.disabled_features
+        self.stats = self.meta.stats
 
-        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):
@@ -1304,18 +1473,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.
         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)
-        item = create_padded_features(item, self.meta.info["features"])
+        item = create_padded_features(item, self.meta.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):

src/lerobot/datasets/utils_must.py

@@ -0,0 +1,230 @@
+"""
+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
+
+def find_start_of_motion(velocities, window_size, threshold, motion_buffer):
+    for t in range(len(velocities) - window_size):
+        window_mean = velocities[t:t+window_size].mean()
+        if window_mean > threshold:
+            return max(0, t - motion_buffer)  # include slight context before motion
+    return 0
+
+import yaml
+import requests
+def load_yaml_mapping(name: str) -> dict:
+    """
+    Loads a YAML mapping from a Hugging Face repo.
+    Example: name='features' → https://huggingface.co/jadechoghari/smolvla-keys/resolve/main/features.yaml
+    """
+    url = f"https://huggingface.co/jadechoghari/smolvla-keys/resolve/main/{name}.yaml"
+    response = requests.get(url)
+    response.raise_for_status()  # raise if the download fails
+
+    return yaml.safe_load(response.text)
+
+# Example usage
+TASKS_KEYS_MAPPING = load_yaml_mapping("tasks")
+FEATURE_KEYS_MAPPING = load_yaml_mapping("features")
+EPISODES_DATASET_MAPPING = {
+    "cadene/droid_1.0.1": list(range(50)),
+    "danaaubakirova/svla_so100_task5_v3": [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51],
+    "danaaubakirova/svla_so100_task4_v3": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53],
+}
+ACTION = "action"
+OBS_STATE = "observation.state"
+TASK = "task"
+ROBOT = "robot_type"
+TRAINING_FEATURES = {
+    0: [ACTION, OBS_STATE, TASK, ROBOT, OBS_IMAGE],
+    1: [ACTION, OBS_STATE, TASK, ROBOT, OBS_IMAGE, OBS_IMAGE_2],
+    2: [ACTION, OBS_STATE, TASK, ROBOT, OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3],
+}