make training work

src/lerobot/datasets/factory.py

@@ -66,58 +66,6 @@ def resolve_delta_timestamps(
 
     return delta_timestamps
 
-
-def make_dataset1(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 isinstance(cfg.dataset.repo_id, str):
-        ds_meta = LeRobotDatasetMetadata(
-            cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
-        )
-        delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
-        dataset = LeRobotDataset(
-            cfg.dataset.repo_id,
-            root=cfg.dataset.root,
-            episodes=cfg.dataset.episodes,
-            delta_timestamps=delta_timestamps,
-            image_transforms=image_transforms,
-            revision=cfg.dataset.revision,
-            video_backend=cfg.dataset.video_backend,
-        )
-    else:
-        raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
-        dataset = MultiLeRobotDataset(
-            cfg.dataset.repo_id,
-            # TODO(aliberts): add proper support for multi dataset
-            # delta_timestamps=delta_timestamps,
-            image_transforms=image_transforms,
-            video_backend=cfg.dataset.video_backend,
-        )
-        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
-
 def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
@@ -144,7 +92,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         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,
         )
@@ -157,7 +104,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             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,
@@ -170,12 +116,13 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         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)
+        # training_features = TRAINING_FEATURES.get(cfg.dataset.features_version, None)
+        #FIXME: (jadechoghari): check support for training features
+        training_features = None
         dataset = MultiLeRobotDataset(
             repo_id,
             # TODO(aliberts): add proper support for multi dataset
@@ -183,11 +130,10 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             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_action_dim=cfg.policy.max_action_dim,
+            max_state_dim=cfg.policy.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,

src/lerobot/datasets/lerobot_dataset.py

@@ -82,7 +82,7 @@ from lerobot.datasets.video_utils import (
 )
 
 # mustafa stuff here
-from lerobot.common.datasets.utils_must import (
+from lerobot.datasets.utils_must import (
     reshape_features_to_max_dim,
     keep_datasets_with_valid_fps,
     keep_datasets_with_the_same_features_per_robot_type,
@@ -97,7 +97,7 @@ from lerobot.common.datasets.utils_must import (
     OBS_IMAGE_3,
     TASKS_KEYS_MAPPING,
 )
-from lerobot.common.constants import (
+from lerobot.constants import (
     ACTION,
     OBS_ENV_STATE,
     OBS_STATE,
@@ -124,7 +124,6 @@ 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

src/lerobot/datasets/utils_must.py

@@ -3,10 +3,19 @@ Utils function by Mustafa to refactor
 """
 import torch
 import numpy as np
-from lerobot.common.datasets.compute_stats import (
+from lerobot.datasets.compute_stats import (
     aggregate_stats
 )
 from collections import defaultdict
+from torch.utils.data.dataloader import default_collate
+import torch.nn.functional as F
+
+import torch
+from typing import Dict, List
+
+
+
+from typing import Dict, List
 OBS_IMAGE = "observation.image"
 OBS_IMAGE_2 = "observation.image2"
 OBS_IMAGE_3 = "observation.image3"
@@ -170,6 +179,9 @@ def pad_tensor(
     is_numpy = isinstance(tensor, np.ndarray)
     if is_numpy:
         tensor = torch.tensor(tensor)
+    if tensor.ndim == 0:
+        # Scalar — return as-is, no padding needed
+        return tensor
     pad = max_size - tensor.shape[pad_dim]
     if pad > 0:
         pad_sizes = (0, pad)  # pad right
@@ -189,6 +201,8 @@ def map_dict_keys(item: dict, feature_keys_mapping: dict, training_features: lis
         else:
             if training_features is None or key in training_features or pad_key in key:
                 features[key] = item[key]
+
+    # breakpoint()
     return features
 
 def find_start_of_motion(velocities, window_size, threshold, motion_buffer):
@@ -228,3 +242,48 @@ TRAINING_FEATURES = {
     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],
 }
+
+def is_batch_need_padding(values: list[torch.Tensor], pad_dim: int = -1) -> int:
+    return len(values[0].shape) > 0  # and len(set([v.shape[pad_dim] for v in values])) > 1
+
+
+def pad_tensor_to_shape(tensor: torch.Tensor, target_shape: tuple, pad_value: float = 0.0) -> torch.Tensor:
+    """Pads a tensor to the target shape (right/bottom only)."""
+    pad = []
+    for actual, target in zip(reversed(tensor.shape), reversed(target_shape)):
+        pad.extend([0, max(target - actual, 0)])
+    return F.pad(tensor, pad, value=pad_value)
+
+
+def multidataset_collate_fn(
+    batch: List[Dict[str, torch.Tensor]],
+    keys_to_max_dim: Dict[str, tuple] = {},
+    pad_value: float = 0.0,
+) -> Dict[str, torch.Tensor]:
+    """
+    Pads tensors to given target shape (if provided), otherwise uses per-batch max.
+    Supports 1D (e.g. action), 3D (e.g. [C,H,W] images).
+    """
+    collated_batch = [{} for _ in range(len(batch))]
+    batch_keys = batch[0].keys()
+
+    for key in batch_keys:
+        values = [sample[key] for sample in batch]
+        sample = values[0]
+
+        if not isinstance(sample, torch.Tensor):
+            for i in range(len(batch)):
+                collated_batch[i][key] = values[i]
+            continue
+
+        # use user-specified shape if available
+        if key in keys_to_max_dim and keys_to_max_dim[key] is not None:
+            target_shape = keys_to_max_dim[key]
+        else:
+            # compute per-batch max shape
+            target_shape = tuple(max(v.shape[i] for v in values) for i in range(sample.ndim))
+
+        for i in range(len(batch)):
+            collated_batch[i][key] = pad_tensor_to_shape(values[i], target_shape, pad_value=pad_value)
+
+    return default_collate(collated_batch)

src/lerobot/scripts/train.py

@@ -51,7 +51,8 @@ from lerobot.utils.utils import (
     init_logging,
 )
 from lerobot.utils.wandb_utils import WandBLogger
-
+from lerobot.datasets.utils_must import multidataset_collate_fn
+from functools import partial
 
 def update_policy(
     train_metrics: MetricsTracker,
@@ -173,7 +174,9 @@ def train(cfg: TrainPipelineConfig):
     else:
         shuffle = True
         sampler = None
-
+    
+    keys_to_max_dim = getattr(dataset.meta, "keys_to_max_dim", {})
+    collate_fn = partial(multidataset_collate_fn, keys_to_max_dim=keys_to_max_dim)
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=cfg.num_workers,