Merge branch 'main' into feat/add-multidataset-training

* refactor(scripts): update system info script

* chore(scripts): rename info script

* feat(scripts): add entrypoint for info

* chore(ci): update issue report template

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -25,7 +25,7 @@ body:
     id: system-info
     attributes:
       label: System Info
-      description: If needed, you can share your lerobot configuration with us by running `python -m lerobot.scripts.display_sys_info` and copy-pasting its outputs below
+      description: Please share your LeRobot configuration by running `lerobot-info` (if installed) or `python -m lerobot.scripts.display_sys_info` (if not installed) and pasting the output below.
       render: Shell
       placeholder: lerobot version, OS, python version, numpy version, torch version, and lerobot's configuration
     validations:

docs/source/hilserl.mdx

@@ -62,7 +62,7 @@ pip install -e ".[hilserl]"
 
 ### Understanding Configuration
 
-The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/scripts/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
+The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
 
 <!-- prettier-ignore-start -->
 ```python
@@ -518,7 +518,7 @@ During the online training, press `space` to take over the policy and `space` ag
 Start the recording process, an example of the config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json):
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
+python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
 ```
 
 During recording:
@@ -549,7 +549,7 @@ Note: If you already know the crop parameters, you can skip this step and just s
 Use the `crop_dataset_roi.py` script to interactively select regions of interest in your camera images:
 
 ```bash
-python -m lerobot.scripts.rl.crop_dataset_roi --repo-id username/pick_lift_cube
+python -m lerobot.rl.crop_dataset_roi --repo-id username/pick_lift_cube
 ```
 
 1. For each camera view, the script will display the first frame
@@ -618,7 +618,7 @@ Before training, you need to collect a dataset with labeled examples. The `recor
 To collect a dataset, you need to modify some parameters in the environment configuration based on HILSerlRobotEnvConfig.
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
+python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
 ```
 
 **Key Parameters for Data Collection**
@@ -764,7 +764,7 @@ or set the argument in the json config file.
 Run `gym_manipulator.py` to test the model.
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config.json
+python -m lerobot.rl.gym_manipulator --config_path path/to/env_config.json
 ```
 
 The reward classifier will automatically provide rewards based on the visual input from the robot's cameras.
@@ -777,7 +777,7 @@ The reward classifier will automatically provide rewards based on the visual inp
 2. **Collect a dataset**:
 
    ```bash
-   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
+   python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 3. **Train the classifier**:
@@ -788,7 +788,7 @@ The reward classifier will automatically provide rewards based on the visual inp
 
 4. **Test the classifier**:
    ```bash
-   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
+   python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 ### Training with Actor-Learner
@@ -810,7 +810,7 @@ Create a training configuration file (example available [here](https://huggingfa
 First, start the learner server process:
 
 ```bash
-python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The learner:
@@ -825,7 +825,7 @@ The learner:
 In a separate terminal, start the actor process with the same configuration:
 
 ```bash
-python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The actor:

docs/source/hilserl_sim.mdx

@@ -91,7 +91,7 @@ Important parameters:
 To run the environment, set mode to null:
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
+python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Recording a Dataset
@@ -118,7 +118,7 @@ To collect a dataset, set the mode to `record` whilst defining the repo_id and n
 ```
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
+python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Training a Policy
@@ -126,13 +126,13 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.j
 To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/train_config.json) and run the actor and learner servers:
 
 ```bash
-python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
+python -m lerobot.rl.actor --config_path path/to/train_gym_hil_env.json
 ```
 
 In a different terminal, run the learner server:
 
 ```bash
-python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
+python -m lerobot.rl.learner --config_path path/to/train_gym_hil_env.json
 ```
 
 The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.

docs/source/il_sim.mdx

@@ -61,14 +61,14 @@ Then we can run this command to start:
 <hfoption id="Linux">
 
 ```bash
-python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
+python -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
+mjpython -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
@@ -198,14 +198,14 @@ Then you can run this command to visualize your trained policy
 <hfoption id="Linux">
 
 ```bash
-python -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
+python -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
+mjpython -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>

examples/tester.py

@@ -0,0 +1,66 @@
+from lerobot.datasets.lerobot_dataset import MultiLeRobotDataset
+
+REPO_A = "lerobot/pusht"
+REPO_B = "lerobot/aloha_mobile_cabinet"  # replace with the actual repo id
+
+feature_keys_mapping = {
+    REPO_A: {  # pusht (1 camera, 2-dim)
+        "action": "actions",
+        "observation.state": "obs_state",
+        "observation.image": "obs_image.cam_high",
+    },
+    REPO_B: {  # dual arm (3 cameras, 14-dim)
+        "action": "actions",
+        "observation.state": "obs_state",
+        "observation.images.cam_high": "obs_image.cam_high",
+        "observation.images.cam_left_wrist": "obs_image.cam_left_wrist",
+        "observation.images.cam_right_wrist": "obs_image.cam_right_wrist",
+    },
+}
+
+from torchvision.transforms.v2 import Compose, ToImage, Resize
+image_tf = Compose([
+    ToImage(),          # converts to tensor if needed
+    Resize((224, 224)), # unify sizes across datasets (96x96 vs 480x640)
+])
+
+from torch.utils.data import DataLoader
+
+dataset = MultiLeRobotDataset(
+    repo_ids=[REPO_A, REPO_B],
+    image_transforms=image_tf,              # ensures same HxW
+    feature_keys_mapping=feature_keys_mapping,
+    train_on_all_features=True,             # keep union of cameras; zero-fill missing
+    # optional: override if you want fixed maxima; else inferred:
+    # max_action_dim=14,
+    # max_state_dim=14,
+    max_action_dim=14,
+    max_state_dim=14,
+    max_image_dim=224,
+    ignore_keys=[
+        "next.*", # drop reward/done/success
+        "index",
+        "timestamp",
+        "videos/*",               # drop all video metadata
+        "observation.effort",     # 👈 drop effort everywhere
+    ],
+)
+breakpoint()
+loader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=0, pin_memory=True)
+for _ in range(100):
+    batch = next(iter(loader))
+
+breakpoint()
+# vectors padded to maxima (pusht:2 -> 14; dual-arm:14 -> 14)
+assert batch["actions"].shape[-1] == 14
+assert batch["obs_state"].shape[-1] == 14
+assert batch["actions_padding_mask"].shape[-1] == 14
+assert batch["obs_state_padding_mask"].shape[-1] == 14
+
+# cameras: all canonical keys exist; pusht will have wrists zero-filled
+for cam in ["obs_image.cam_high", "obs_image.cam_left_wrist", "obs_image.cam_right_wrist"]:
+    assert cam in batch
+    assert f"{cam}_is_pad" in batch
+    # images should all be 3x224x224 (or your transform’s size)
+    img = batch[cam]
+    assert img.ndim in (4, 5)  # (B,C,H,W) or (B,T,C,H,W) depending on your loader

examples/tester.sh

@@ -0,0 +1,16 @@
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+
+python examples/tester.py

pyproject.toml

@@ -171,6 +171,7 @@ lerobot-setup-motors="lerobot.setup_motors:main"
 lerobot-teleoperate="lerobot.teleoperate:main"
 lerobot-eval="lerobot.scripts.eval:main"
 lerobot-train="lerobot.scripts.train:main"
+lerobot-info="lerobot.scripts.lerobot_info:main"
 
 # ---------------- Tool Configurations ----------------
 [tool.setuptools.packages.find]

src/lerobot/configs/policies.py

@@ -196,11 +196,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
             config = json.load(f)
 
         config.pop("type")
-        with tempfile.NamedTemporaryFile("w+") as f:
+        with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
             json.dump(config, f)
             config_file = f.name
-            f.flush()
 
-            cli_overrides = policy_kwargs.pop("cli_overrides", [])
-            with draccus.config_type("json"):
-                return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
+        cli_overrides = policy_kwargs.pop("cli_overrides", [])
+        with draccus.config_type("json"):
+            return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)

src/lerobot/datasets/compute_stats.py

@@ -174,3 +174,79 @@ def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np
         aggregated_stats[key] = aggregate_feature_stats(stats_with_key)
 
     return aggregated_stats
+
+import numpy as np
+
+def aggregate_stats_multi(
+    stats_list: list[dict[str, dict]],
+    max_action_dim: int | None = None,
+    max_state_dim: int | None = None,
+) -> dict[str, dict[str, np.ndarray]]:
+    """Aggregate stats from multiple compute_stats outputs into a single set of stats.
+    
+    Supports heterogeneous robots by padding action/state stats to the max dim.
+    The final stats will have the union of all data keys from each of the stats dicts.
+
+    - new_min  = elementwise min across datasets
+    - new_max  = elementwise max across datasets
+    - new_mean = weighted mean (by count)
+    - new_std  = recomputed from total variance
+    """
+
+    data_keys = {key for stats in stats_list for key in stats}
+    aggregated_stats = {key: {} for key in data_keys}
+
+    def _pad(arr: np.ndarray, target: int) -> np.ndarray:
+        if arr.ndim == 0:  # scalar
+            return arr
+        if target is None or target <= 0 or arr.shape[-1] == target:
+            return arr
+        pad_width = [(0, 0)] * arr.ndim
+        pad_width[-1] = (0, target - arr.shape[-1])
+        return np.pad(arr, pad_width, mode="constant")
+
+    for key in data_keys:
+        stats_with_key = [stats[key] for stats in stats_list if key in stats]
+
+        # decide if this key should be padded
+        target_dim = None
+        if "action" in key and max_action_dim:
+            target_dim = max_action_dim
+        elif "state" in key and max_state_dim:
+            target_dim = max_state_dim
+
+        padded = []
+        counts = []
+        for s in stats_with_key:
+            mean = _pad(np.array(s["mean"]), target_dim)
+            std = _pad(np.array(s["std"]), target_dim)
+            min_ = _pad(np.array(s["min"]), target_dim)
+            max_ = _pad(np.array(s["max"]), target_dim)
+            count = s.get("count", 1)
+
+            padded.append(dict(mean=mean, std=std, min=min_, max=max_, count=count))
+            counts.append(count)
+
+        counts = np.array(counts, dtype=np.float64)
+        total_count = counts.sum()
+
+        means = np.stack([p["mean"] for p in padded])
+        stds  = np.stack([p["std"]  for p in padded])
+        mins  = np.stack([p["min"]  for p in padded])
+        maxs  = np.stack([p["max"]  for p in padded])
+
+        # weighted mean (broadcast weights properly)
+        new_mean = np.average(means, axis=0, weights=counts)
+        new_var  = np.average(stds**2 + (means - new_mean)**2, axis=0, weights=counts)
+
+        new_std = np.sqrt(new_var)
+
+        aggregated_stats[key] = {
+            "min": mins.min(axis=0),
+            "max": maxs.max(axis=0),
+            "mean": new_mean,
+            "std": new_std,
+            "count": int(total_count),
+        }
+
+    return aggregated_stats

src/lerobot/datasets/lerobot_dataset.py

@@ -31,6 +31,7 @@ import torch.utils
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.errors import RevisionNotFoundError
 
+from collections import defaultdict
 from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats
 from lerobot.datasets.image_writer import AsyncImageWriter, write_image
@@ -81,7 +82,12 @@ from lerobot.datasets.video_utils import (
 )
 
 CODEBASE_VERSION = "v3.0"
-
+OBS_IMAGE = "observation.image"
+OBS_IMAGE_2 = "observation.image_2"
+OBS_IMAGE_3 = "observation.image_3"
+OBS_STATE = "observation.state"
+OBS_ENV_STATE = "observation.env_state"
+ACTION = "action"
 
 class LeRobotDatasetMetadata:
     def __init__(
@@ -1322,13 +1328,139 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 
+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",
+}
+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]
 
+        # 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, strict=False):
+                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, strict=False)}
+        self.tasks = {repo_id: ds.meta.tasks for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+        self.info = {repo_id: ds.meta.info for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+
+
+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
+
+        # step 2: keep datasets with same features per robot type
+        consistent_datasets, keep_mask = keep_datasets_with_the_same_features_per_robot_type(
+            datasets
+        )
+
+        self.cleaned_datasets = consistent_datasets
+        self.keep_mask = keep_mask
+        self.cleaned_weights = [sampling_weights[i] for i in range(len(datasets)) if keep_mask[i]]
+        self.cleaned_repo_ids = [repo_ids[i] for i in range(len(datasets)) if keep_mask[i]]
+        self.cleaned_datasets_repo_ids = [
+            datasets_repo_ids[i] for i in range(len(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)
+
+# --- at the top of the file (same imports as before) ---
+from collections import defaultdict
+from typing import Callable
+import copy
+import numpy as np
+import torch
+import datasets
+from pathlib import Path
+
+# If you already have these in your codebase, reuse them
+try:
+    from lerobot.common.constants import (
+        ACTION, OBS_ENV_STATE, OBS_STATE, OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3
+    )
+except Exception:
+    # Fallbacks if constants are already strings elsewhere
+    ACTION = "action"
+    OBS_ENV_STATE = "observation.env_state"
+    OBS_STATE = "observation.state"
+    OBS_IMAGE = "observation.image"
+    OBS_IMAGE_2 = "observation.image_2"
+    OBS_IMAGE_3 = "observation.image_3"
+
+IGNORED_KEYS = ["observation.effort"]
 class MultiLeRobotDataset(torch.utils.data.Dataset):
-    """A dataset consisting of multiple underlying `LeRobotDataset`s.
-
-    The underlying `LeRobotDataset`s are effectively concatenated, and this class adopts much of the API
-    structure of `LeRobotDataset`.
-    """
+    # ... keep your existing docstring ...
 
     def __init__(
         self,
@@ -1336,99 +1468,253 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         root: str | Path | None = None,
         episodes: dict | None = None,
         image_transforms: Callable | None = None,
-        delta_timestamps: dict[str, list[float]] | None = None,
+        delta_timestamps: dict[list[float]] | None = None,
         tolerances_s: dict | None = None,
         download_videos: bool = True,
         video_backend: str | None = None,
+        # --- NEW: simple add-ons ---
+        sampling_weights: list[float] | None = None,
+        feature_keys_mapping: dict[str, dict[str, str]] | None = None,
+        max_action_dim: int | None = None,
+        max_state_dim: int | None = None,
+        max_num_images: int | None = None,
+        max_image_dim: int | None = None,
+        train_on_all_features: bool = False,
+        min_fps: int = 1,
+        max_fps: int = 100,
+        ignore_keys: list[str] | None = None,  # exact or glob patterns
     ):
         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)
-        # 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
-        ]
 
-        # 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)
+        # --- NEW: store mapping and simple knobs ---
+        self.feature_keys_mapping: dict[str, dict[str, str]] = feature_keys_mapping or {}
+        self.train_on_all_features = train_on_all_features
+        self.max_action_dim = max_action_dim
+        self.max_state_dim = max_state_dim
+        self.max_image_dim = max_image_dim
+        self.max_num_images = max_num_images  # (optional, we don’t enforce count, we enforce names)
+        self._ignore_patterns = list(ignore_keys or [])
+        # Build underlying single datasets
+        _datasets = []
+        datasets_repo_ids = []
+        self.sampling_weights = []
 
+        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:
+                _datasets.append(
+                    LeRobotDataset(
+                        repo_id,
+                        root=self.root / repo_id,
+                        episodes=episodes.get(repo_id, None) if episodes else None,
+                        image_transforms=image_transforms,  # transforms applied inside single ds
+                        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,
+                    )
+                )
+                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: "
+            f"{self.sampling_weights} corresponding to: {datasets_repo_ids}"
+        )
+
+        # Bookkeeping for mapping & canonical image inventory
         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
+        self._datasets = _datasets
+        self.datasets_repo_ids = datasets_repo_ids
+
+        # --- NEW: compute “canonical image keys” (targets across all mappings) ---
+        self._canonical_image_keys: set[str] = set()
+        self._source_keys_per_repo: dict[str, set[str]] = {}
+        self._target_keys_per_repo: dict[str, set[str]] = {}
+        for rid, mapping in self.feature_keys_mapping.items():
+            src_keys = set(mapping.keys())
+            tgt_keys = set(mapping.values())
+            self._source_keys_per_repo[rid] = src_keys
+            self._target_keys_per_repo[rid] = tgt_keys
+            # union of target names (we will ensure these exist at __getitem__)
+            self._canonical_image_keys |= {
+                k for k in tgt_keys if self._is_image_key_like(k)
+            }
+
+        # If user didn’t give any mapping, fall back to native keys (no-ops)
+        if not self._canonical_image_keys and self.train_on_all_features:
+            # discover all image-like keys from raw features
+            for ds in self._datasets:
+                for k, v in ds.hf_features.items():
+                    if isinstance(v, (datasets.Image, VideoFrame)):
+                        self._canonical_image_keys.add(k)
+
+        # Cleaner: keep fps & consistent feature sets per robot type (unchanged)
+        cleaner = MultiLeRobotDatasetCleaner(
+            datasets=self._datasets,
+            repo_ids=repo_ids,
+            sampling_weights=self.sampling_weights,
+            datasets_repo_ids=self.datasets_repo_ids,
+            min_fps=min_fps,
+            max_fps=max_fps,
+        )
+        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
+
+        # Meta (unchanged): we give it dim maxima; it will reshape/pad vectors
+        self.meta = MultiLeRobotDatasetMeta(
+            datasets=self._datasets,
+            repo_ids=self.repo_ids,
+            keys_to_max_dim={
+                ACTION: self.max_action_dim if self.max_action_dim is not None else -1,
+                OBS_ENV_STATE: self.max_state_dim if self.max_state_dim is not None else -1,
+                OBS_STATE: self.max_state_dim if self.max_state_dim is not None else -1,
+                OBS_IMAGE: self.max_image_dim if self.max_image_dim is not None else -1,
+                OBS_IMAGE_2: self.max_image_dim if self.max_image_dim is not None else -1,
+                OBS_IMAGE_3: self.max_image_dim if self.max_image_dim is not None else -1,
+            },
+            train_on_all_features=train_on_all_features,
+        )
+
+        # --- NEW: track dropped (source) keys so collate won’t expect them
+        # Anything that we *rename away* should be considered disabled,
+        # otherwise downstream may expect them to exist.
+        self._dropped_keys = set()
+        for rid, mapping in self.feature_keys_mapping.items():
+            self._dropped_keys |= set(mapping.keys())
+
+        # Merge with meta’s disabled features
+        self.disabled_features = set(self.meta.disabled_features) | self._dropped_keys
+
+        self.stats = self.meta.stats
+
+        # --- NEW: cache an example image shape per canonical key (lazy, filled on first use)
+        self._cached_img_shape: dict[str, torch.Size] = {}
+
+    # ---------------------- NEW small helpers ----------------------
+
+    def _is_image_key_like(self, key: str) -> bool:
+        # A loose heuristic: rely on name OR on features later
+        return ("image" in key) or ("cam_" in key) or ("images." in key)
+    
+    def _should_ignore(self, key: str) -> bool:
+        # exact or glob-style match
+        for pat in self._ignore_patterns:
+            if key == pat or fnmatch.fnmatch(key, pat):
+                return True
+        return False
+    def _apply_feature_mapping(self, item: dict, repo_id: str) -> dict:
+        """
+        Rename features according to feature_keys_mapping[repo_id].
+        - Moves tensor/image under target key.
+        - Drops source key if moved.
+        - Adds *_is_pad=False for image targets we fill/keep.
+        """
+        mapping = self.feature_keys_mapping.get(repo_id, {}) or {}
+        if not mapping:
+            return item
+
+        for src, tgt in mapping.items():
+            if src in item:
+                # Move value
+                item[tgt] = item[src]
+                # Drop the source to avoid duplication
+                del item[src]
+        return item
+
+    def _ensure_union_image_keys(self, item: dict) -> dict:
+        """
+        Ensure that every canonical image key exists.
+        When missing, create a zero tensor matching (B,C,H,W) or (C,H,W) of an available image.
+        Also add boolean mask at f"{key}_is_pad".
+        """
+        if not self.train_on_all_features or not self._canonical_image_keys:
+            return item
+
+        # find any existing image tensor in item to copy shape/dtype
+        exemplar = None
+        for k in list(item.keys()):
+            v = item[k]
+            if torch.is_tensor(v) and v.ndim in (3, 4, 5):  # (C,H,W) or (B,C,H,W) or (B,T,C,H,W)
+                exemplar = v
+                break
+
+        # fallback to a safe 3x224x224 if nothing found
+        def _fallback_image():
+            return torch.zeros(3, 224, 224, dtype=torch.uint8)
+
+        for key in self._canonical_image_keys:
+            if key not in item:
+                img = torch.zeros_like(exemplar) if exemplar is not None else _fallback_image()
+                item[key] = img
+                item[f"{key}_is_pad"] = torch.tensor(True, dtype=torch.bool)
+            else:
+                # Add a mask saying it’s *not* padded
+                if f"{key}_is_pad" not in item:
+                    item[f"{key}_is_pad"] = torch.tensor(False, dtype=torch.bool)
+        return item
+
+    # ---------------------- existing API below (mostly unchanged) ----------------------
 
     @property
     def repo_id_to_index(self):
-        """Return a mapping from dataset repo_id to a dataset index automatically created by this class.
-
-        This index is incorporated as a data key in the dictionary returned by `__getitem__`.
-        """
         return {repo_id: i for i, repo_id in enumerate(self.repo_ids)}
 
     @property
     def repo_index_to_id(self):
-        """Return the inverse mapping if repo_id_to_index."""
         return {v: k for k, v in self.repo_id_to_index}
 
     @property
     def fps(self) -> int:
-        """Frames per second used during data collection.
-
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
         return self._datasets[0].meta.info["fps"]
 
     @property
     def video(self) -> bool:
-        """Returns True if this dataset loads video frames from mp4 files.
-
-        Returns False if it only loads images from png files.
-
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
         return self._datasets[0].meta.info.get("video", False)
 
     @property
     def features(self) -> datasets.Features:
-        features = {}
+        """
+        Extend native HF features with any *target* keys introduced by mapping.
+        We copy the source spec for targets that didn’t exist in any raw dataset.
+        """
+        features: dict[str, datasets.features.Feature] = {}
         for dataset in self._datasets:
-            features.update({k: v for k, v in dataset.hf_features.items() if k not in self.disabled_features})
+            for k, v in dataset.hf_features.items():
+                if k not in self.disabled_features:
+                    features[k] = v
+
+        # Add mapped target image specs if not present yet
+        for rid, mapping in self.feature_keys_mapping.items():
+            ds = None
+            # find the dataset object to read feature spec for source
+            for _ds, _rid in zip(self._datasets, self.repo_ids, strict=False):
+                if _rid == rid:
+                    ds = _ds
+                    break
+            if ds is None:
+                continue
+            for src, tgt in mapping.items():
+                if tgt not in features and src in ds.hf_features:
+                    features[tgt] = ds.hf_features[src]
+
         return features
 
     @property
     def camera_keys(self) -> list[str]:
-        """Keys to access image and video stream from cameras."""
         keys = []
         for key, feats in self.features.items():
             if isinstance(feats, (datasets.Image, VideoFrame)):
@@ -1437,12 +1723,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
 
     @property
     def video_frame_keys(self) -> list[str]:
-        """Keys to access video frames that requires to be decoded into images.
-
-        Note: It is empty if the dataset contains images only,
-        or equal to `self.cameras` if the dataset contains videos only,
-        or can even be a subset of `self.cameras` in a case of a mixed image/video dataset.
-        """
         video_frame_keys = []
         for key, feats in self.features.items():
             if isinstance(feats, VideoFrame):
@@ -1451,21 +1731,14 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
 
     @property
     def num_frames(self) -> int:
-        """Number of samples/frames."""
         return sum(d.num_frames for d in self._datasets)
 
     @property
     def num_episodes(self) -> int:
-        """Number of episodes."""
         return sum(d.num_episodes for d in self._datasets)
 
     @property
     def tolerance_s(self) -> float:
-        """Tolerance in seconds used to discard loaded frames when their timestamps
-        are not close enough from the requested frames. It is only used when `delta_timestamps`
-        is provided or when loading video frames from mp4 files.
-        """
-        # 1e-4 to account for possible numerical error
         return 1 / self.fps - 1e-4
 
     def __len__(self):
@@ -1474,22 +1747,83 @@ 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]
+
+        # Identify which repo this sample came from
+        repo_id = self.datasets_repo_ids[dataset_idx]
+
+        # --- NEW: apply mapping and ensure union of image keys ---
+        item = self._apply_feature_mapping(item, repo_id)
+        item = self._ensure_union_image_keys(item)
+
+        # annotate dataset index for downstream
         item["dataset_index"] = torch.tensor(dataset_idx)
+
+        # Pad vector features to max dims using meta (unchanged)
+        item = create_padded_features(item, self.meta.features)
+
+        # Drop any disabled (including original source keys we remapped away)
         for data_key in self.disabled_features:
             if data_key in item:
                 del item[data_key]
+        for k in IGNORED_KEYS:
+            if k in item:
+                item.pop(k)
+        # Convert any datasets.Image still present to tensor
+        if self.image_transforms is not None:
+            for cam in [k for k in item.keys() if self._is_image_key_like(k)]:
+                val = item[cam]
+                if not torch.is_tensor(val):
+                    item[cam] = self.image_transforms(val)
+        # 🔑 Pad actions if too short
+        if "actions" in item and self.max_action_dim is not None:
+            act = item["actions"]
+            if act.shape[-1] < self.max_action_dim:
+                pad_len = self.max_action_dim - act.shape[-1]
+                item["actions"] = torch.cat([act, torch.zeros(pad_len, dtype=act.dtype)], dim=-1)
+                item["actions_padding_mask"] = torch.cat(
+                    [torch.zeros_like(act, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+
+        # pad obs_state if too short
+        if "obs_state" in item and self.max_state_dim is not None:
+            st = item["obs_state"]
+            if st.shape[-1] < self.max_state_dim:
+                pad_len = self.max_state_dim - st.shape[-1]
+                item["obs_state"] = torch.cat([st, torch.zeros(pad_len, dtype=st.dtype)], dim=-1)
+                item["obs_state_padding_mask"] = torch.cat(
+                    [torch.zeros_like(st, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+        # actions
+        if "actions" in item and self.max_action_dim is not None:
+            act = item["actions"]
+            if act.shape[-1] < self.max_action_dim:
+                pad_len = self.max_action_dim - act.shape[-1]
+                item["actions"] = torch.cat([act, torch.zeros(pad_len, dtype=act.dtype)], dim=-1)
+                mask = torch.cat(
+                    [torch.zeros_like(act, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+            else:
+                mask = torch.zeros(self.max_action_dim, dtype=torch.bool)  # 👈 all False if no padding
+            item["actions_padding_mask"] = mask
+        # obs state
+        if "obs_state" in item and self.max_state_dim is not None:
+            st = item["obs_state"]
+            if st.shape[-1] < self.max_state_dim:
+                pad_len = self.max_state_dim - st.shape[-1]
+                item["obs_state"] = torch.cat([st, torch.zeros(pad_len, dtype=st.dtype)], dim=-1)
+                mask = torch.cat(
+                    [torch.zeros_like(st, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+            else:
+                mask = torch.zeros(self.max_state_dim, dtype=torch.bool)  # 👈 always add mask
+            item["obs_state_padding_mask"] = mask
 
         return item
 
@@ -1506,3 +1840,149 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
             f"  Transformations: {self.image_transforms},\n"
             f")"
         )
+
+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 episode_stats_values(ds.meta)
+        ]
+        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(episode_stats_values(ds.meta)), 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(episode_stats_values(ds.meta)))
+        # 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 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 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
+
+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 episode_stats_values(meta):
+    episodes = meta.episodes.to_pandas().to_dict(orient="records")
+    return [
+        {k: v for k, v in ep.items() if isinstance(v, dict) and "mean" in v}
+        for ep in episodes
+    ]

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -404,7 +404,7 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
     info["video_files_size_in_mb"] = video_file_size_in_mb
     info["data_path"] = DEFAULT_DATA_PATH
     info["video_path"] = DEFAULT_VIDEO_PATH
-    info["fps"] = float(info["fps"])
+    info["fps"] = int(info["fps"])
     for key in info["features"]:
         if info["features"][key]["dtype"] == "video":
             # already has fps in video_info

src/lerobot/policies/pretrained.py

@@ -246,7 +246,9 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             base_model=base_model,
         )
 
-        template_card = files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text()
+        template_card = (
+            files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
+        )
         card = ModelCard.from_template(card_data, template_str=template_card)
         card.validate()
         return card

src/lerobot/rl/actor.py

@@ -24,7 +24,7 @@ Examples of usage:
 
 - Start an actor server for real robot training with human-in-the-loop intervention:
 ```bash
-python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 **NOTE**: The actor server requires a running learner server to connect to. Ensure the learner
@@ -64,12 +64,6 @@ from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.processor import TransitionKey
 from lerobot.robots import so100_follower  # noqa: F401
-from lerobot.scripts.rl.gym_manipulator import (
-    create_transition,
-    make_processors,
-    make_robot_env,
-    step_env_and_process_transition,
-)
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.transport import services_pb2, services_pb2_grpc
@@ -96,6 +90,13 @@ from lerobot.utils.utils import (
     init_logging,
 )
 
+from .gym_manipulator import (
+    create_transition,
+    make_processors,
+    make_robot_env,
+    step_env_and_process_transition,
+)
+
 ACTOR_SHUTDOWN_TIMEOUT = 30
 
 # Main entry point

src/lerobot/rl/eval_policy.py

@@ -25,12 +25,13 @@ from lerobot.robots import (  # noqa: F401
     make_robot_from_config,
     so100_follower,
 )
-from lerobot.scripts.rl.gym_manipulator import make_robot_env
 from lerobot.teleoperators import (
     gamepad,  # noqa: F401
     so101_leader,  # noqa: F401
 )
 
+from .gym_manipulator import make_robot_env
+
 logging.basicConfig(level=logging.INFO)
 
 

src/lerobot/rl/learner.py

@@ -25,7 +25,7 @@ Examples of usage:
 
 - Start a learner server for training:
 ```bash
-python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 **NOTE**: Start the learner server before launching the actor server. The learner opens a gRPC server
@@ -73,7 +73,6 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.robots import so100_follower  # noqa: F401
-from lerobot.scripts.rl import learner_service
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.transport import services_pb2_grpc
@@ -100,6 +99,8 @@ from lerobot.utils.utils import (
 )
 from lerobot.utils.wandb_utils import WandBLogger
 
+from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
+
 LOG_PREFIX = "[LEARNER]"
 
 
@@ -639,7 +640,7 @@ def start_learner(
         # TODO: Check if its useful
         _ = ProcessSignalHandler(False, display_pid=True)
 
-    service = learner_service.LearnerService(
+    service = LearnerService(
         shutdown_event=shutdown_event,
         parameters_queue=parameters_queue,
         seconds_between_pushes=cfg.policy.actor_learner_config.policy_parameters_push_frequency,
@@ -649,7 +650,7 @@ def start_learner(
     )
 
     server = grpc.server(
-        ThreadPoolExecutor(max_workers=learner_service.MAX_WORKERS),
+        ThreadPoolExecutor(max_workers=MAX_WORKERS),
         options=[
             ("grpc.max_receive_message_length", MAX_MESSAGE_SIZE),
             ("grpc.max_send_message_length", MAX_MESSAGE_SIZE),
@@ -670,7 +671,7 @@ def start_learner(
 
     shutdown_event.wait()
     logging.info("[LEARNER] Stopping gRPC server...")
-    server.stop(learner_service.SHUTDOWN_TIMEOUT)
+    server.stop(SHUTDOWN_TIMEOUT)
     logging.info("[LEARNER] gRPC server stopped")
 
 

src/lerobot/scripts/display_sys_info.py

@@ -1,90 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Use this script to get a quick summary of your system config.
-It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
-"""
-
-import platform
-
-HAS_HF_HUB = True
-HAS_HF_DATASETS = True
-HAS_NP = True
-HAS_TORCH = True
-HAS_LEROBOT = True
-
-try:
-    import huggingface_hub
-except ImportError:
-    HAS_HF_HUB = False
-
-try:
-    import datasets
-except ImportError:
-    HAS_HF_DATASETS = False
-
-try:
-    import numpy as np
-except ImportError:
-    HAS_NP = False
-
-try:
-    import torch
-except ImportError:
-    HAS_TORCH = False
-
-try:
-    import lerobot
-except ImportError:
-    HAS_LEROBOT = False
-
-
-lerobot_version = lerobot.__version__ if HAS_LEROBOT else "N/A"
-hf_hub_version = huggingface_hub.__version__ if HAS_HF_HUB else "N/A"
-hf_datasets_version = datasets.__version__ if HAS_HF_DATASETS else "N/A"
-np_version = np.__version__ if HAS_NP else "N/A"
-
-torch_version = torch.__version__ if HAS_TORCH else "N/A"
-torch_cuda_available = torch.cuda.is_available() if HAS_TORCH else "N/A"
-cuda_version = torch._C._cuda_getCompiledVersion() if HAS_TORCH and torch.version.cuda is not None else "N/A"
-
-
-# TODO(aliberts): refactor into an actual command `lerobot env`
-def display_sys_info() -> dict:
-    """Run this to get basic system info to help for tracking issues & bugs."""
-    info = {
-        "`lerobot` version": lerobot_version,
-        "Platform": platform.platform(),
-        "Python version": platform.python_version(),
-        "Huggingface_hub version": hf_hub_version,
-        "Dataset version": hf_datasets_version,
-        "Numpy version": np_version,
-        "PyTorch version (GPU?)": f"{torch_version} ({torch_cuda_available})",
-        "Cuda version": cuda_version,
-        "Using GPU in script?": "<fill in>",
-        # "Using distributed or parallel set-up in script?": "<fill in>",
-    }
-    print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
-    print(format_dict(info))
-    return info
-
-
-def format_dict(d: dict) -> str:
-    return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
-
-
-if __name__ == "__main__":
-    display_sys_info()

src/lerobot/scripts/lerobot_info.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Use this script to get a quick summary of your system config.
+It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
+
+Example:
+
+```shell
+lerobot-info
+```
+"""
+
+import importlib
+import platform
+
+
+def get_package_version(package_name: str) -> str:
+    """Get the version of a package if it exists, otherwise return 'N/A'."""
+    try:
+        module = importlib.import_module(package_name)
+        return getattr(module, "__version__", "Installed (version not found)")
+    except ImportError:
+        return "N/A"
+
+
+def get_sys_info() -> dict:
+    """Run this to get basic system info to help for tracking issues & bugs."""
+    # General package versions
+    info = {
+        "lerobot version": get_package_version("lerobot"),
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Huggingface Hub version": get_package_version("huggingface_hub"),
+        "Datasets version": get_package_version("datasets"),
+        "Numpy version": get_package_version("numpy"),
+    }
+
+    # PyTorch and GPU specific information
+    torch_version = "N/A"
+    torch_cuda_available = "N/A"
+    cuda_version = "N/A"
+    gpu_model = "N/A"
+    try:
+        import torch
+
+        torch_version = torch.__version__
+        torch_cuda_available = torch.cuda.is_available()
+        if torch_cuda_available:
+            cuda_version = torch.version.cuda
+            # Gets the name of the first available GPU
+            gpu_model = torch.cuda.get_device_name(0)
+    except ImportError:
+        # If torch is not installed, the default "N/A" values will be used.
+        pass
+
+    info.update(
+        {
+            "PyTorch version": torch_version,
+            "Is PyTorch built with CUDA support?": torch_cuda_available,
+            "Cuda version": cuda_version,
+            "GPU model": gpu_model,
+            "Using GPU in script?": "<fill in>",
+        }
+    )
+
+    return info
+
+
+def format_dict_for_markdown(d: dict) -> str:
+    """Formats a dictionary into a markdown-friendly bulleted list."""
+    return "\n".join([f"- {prop}: {val}" for prop, val in d.items()])
+
+
+def main():
+    system_info = get_sys_info()
+    print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
+    print(format_dict_for_markdown(system_info))
+
+
+if __name__ == "__main__":
+    main()

tests/rl/test_actor.py

@@ -65,7 +65,7 @@ def close_service_stub(channel, server):
 
 @require_package("grpc")
 def test_establish_learner_connection_success():
-    from lerobot.scripts.rl.actor import establish_learner_connection
+    from lerobot.rl.actor import establish_learner_connection
 
     """Test successful connection establishment."""
     stub, _servicer, channel, server = create_learner_service_stub()
@@ -82,7 +82,7 @@ def test_establish_learner_connection_success():
 
 @require_package("grpc")
 def test_establish_learner_connection_failure():
-    from lerobot.scripts.rl.actor import establish_learner_connection
+    from lerobot.rl.actor import establish_learner_connection
 
     """Test connection failure."""
     stub, servicer, channel, server = create_learner_service_stub()
@@ -101,7 +101,7 @@ def test_establish_learner_connection_failure():
 
 @require_package("grpc")
 def test_push_transitions_to_transport_queue():
-    from lerobot.scripts.rl.actor import push_transitions_to_transport_queue
+    from lerobot.rl.actor import push_transitions_to_transport_queue
     from lerobot.transport.utils import bytes_to_transitions
     from tests.transport.test_transport_utils import assert_transitions_equal
 
@@ -137,7 +137,7 @@ def test_push_transitions_to_transport_queue():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_transitions_stream():
-    from lerobot.scripts.rl.actor import transitions_stream
+    from lerobot.rl.actor import transitions_stream
 
     """Test transitions stream functionality."""
     shutdown_event = Event()
@@ -169,7 +169,7 @@ def test_transitions_stream():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_interactions_stream():
-    from lerobot.scripts.rl.actor import interactions_stream
+    from lerobot.rl.actor import interactions_stream
     from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test interactions stream functionality."""

tests/rl/test_actor_learner.py

@@ -90,13 +90,13 @@ def cfg():
 @require_package("grpc")
 @pytest.mark.timeout(10)  # force cross-platform watchdog
 def test_end_to_end_transitions_flow(cfg):
-    from lerobot.scripts.rl.actor import (
+    from lerobot.rl.actor import (
         establish_learner_connection,
         learner_service_client,
         push_transitions_to_transport_queue,
         send_transitions,
     )
-    from lerobot.scripts.rl.learner import start_learner
+    from lerobot.rl.learner import start_learner
     from lerobot.transport.utils import bytes_to_transitions
     from tests.transport.test_transport_utils import assert_transitions_equal
 
@@ -152,12 +152,12 @@ def test_end_to_end_transitions_flow(cfg):
 @require_package("grpc")
 @pytest.mark.timeout(10)
 def test_end_to_end_interactions_flow(cfg):
-    from lerobot.scripts.rl.actor import (
+    from lerobot.rl.actor import (
         establish_learner_connection,
         learner_service_client,
         send_interactions,
     )
-    from lerobot.scripts.rl.learner import start_learner
+    from lerobot.rl.learner import start_learner
     from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test complete interactions flow from actor to learner."""
@@ -226,8 +226,8 @@ def test_end_to_end_interactions_flow(cfg):
 @pytest.mark.parametrize("data_size", ["small", "large"])
 @pytest.mark.timeout(10)
 def test_end_to_end_parameters_flow(cfg, data_size):
-    from lerobot.scripts.rl.actor import establish_learner_connection, learner_service_client, receive_policy
-    from lerobot.scripts.rl.learner import start_learner
+    from lerobot.rl.actor import establish_learner_connection, learner_service_client, receive_policy
+    from lerobot.rl.learner import start_learner
     from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test complete parameter flow from learner to actor, with small and large data."""

tests/rl/test_learner_service.py

@@ -50,7 +50,7 @@ def create_learner_service_stub(
 ):
     import grpc
 
-    from lerobot.scripts.rl.learner_service import LearnerService
+    from lerobot.rl.learner_service import LearnerService
     from lerobot.transport import services_pb2_grpc  # generated from .proto
 
     """Fixture to start a LearnerService gRPC server and provide a connected stub."""