Refactor SACObservationEncoder to improve modularity and readability. Split initialization into dedicated methods for image and state layers, and enhance caching logic for image features. Update forward method to streamline feature encoding and ensure proper normalization handling.

Co-authored-by: s1lent4gnt <kmeftah.khalil@gmail.com>

.gitignore

@@ -26,6 +26,7 @@ outputs
 
 # VS Code
 .vscode
+.devcontainer
 
 # HPC
 nautilus/*.yaml

.pre-commit-config.yaml

@@ -46,7 +46,7 @@ repos:
     rev: v3.19.1
     hooks:
     -   id: pyupgrade
-
+        exclude: '^(.*_pb2_grpc\.py|.*_pb2\.py$)'
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.11.5
     hooks:

benchmarks/video/run_video_benchmark.py

@@ -32,7 +32,11 @@ import numpy as np
 import pandas as pd
 import PIL
 import torch
-from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
+from skimage.metrics import (
+    mean_squared_error,
+    peak_signal_noise_ratio,
+    structural_similarity,
+)
 from tqdm import tqdm
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
@@ -94,7 +98,11 @@ def load_original_frames(imgs_dir: Path, timestamps: list[float], fps: int) -> t
 
 
 def save_decoded_frames(
-    imgs_dir: Path, save_dir: Path, frames: torch.Tensor, timestamps: list[float], fps: int
+    imgs_dir: Path,
+    save_dir: Path,
+    frames: torch.Tensor,
+    timestamps: list[float],
+    fps: int,
 ) -> None:
     if save_dir.exists() and len(list(save_dir.glob("frame_*.png"))) == len(timestamps):
         return
@@ -104,7 +112,10 @@ def save_decoded_frames(
         idx = int(ts * fps)
         frame_hwc = (frames[i].permute((1, 2, 0)) * 255).type(torch.uint8).cpu().numpy()
         PIL.Image.fromarray(frame_hwc).save(save_dir / f"frame_{idx:06d}_decoded.png")
-        shutil.copyfile(imgs_dir / f"frame_{idx:06d}.png", save_dir / f"frame_{idx:06d}_original.png")
+        shutil.copyfile(
+            imgs_dir / f"frame_{idx:06d}.png",
+            save_dir / f"frame_{idx:06d}_original.png",
+        )
 
 
 def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
@@ -120,7 +131,11 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
     imgs_dataset = hf_dataset.select_columns(img_keys[0])
 
     for i, item in enumerate(
-        tqdm(imgs_dataset, desc=f"saving {dataset.repo_id} first episode images", leave=False)
+        tqdm(
+            imgs_dataset,
+            desc=f"saving {dataset.repo_id} first episode images",
+            leave=False,
+        )
     ):
         img = item[img_keys[0]]
         img.save(str(imgs_dir / f"frame_{i:06d}.png"), quality=100)
@@ -275,7 +290,9 @@ def benchmark_encoding_decoding(
     random.seed(seed)
     benchmark_table = []
     for timestamps_mode in tqdm(
-        decoding_cfg["timestamps_modes"], desc="decodings (timestamps_modes)", leave=False
+        decoding_cfg["timestamps_modes"],
+        desc="decodings (timestamps_modes)",
+        leave=False,
     ):
         for backend in tqdm(decoding_cfg["backends"], desc="decodings (backends)", leave=False):
             benchmark_row = benchmark_decoding(

examples/12_train_hilserl_classifier.md

@@ -0,0 +1,94 @@
+# Training a HIL-SERL Reward Classifier with LeRobot
+
+This tutorial provides step-by-step instructions for training a reward classifier using LeRobot.
+
+---
+
+## Training Script Overview
+
+LeRobot includes a ready-to-use training script located at [`lerobot/scripts/train_hilserl_classifier.py`](../../lerobot/scripts/train_hilserl_classifier.py). Here's an outline of its workflow:
+
+1. **Configuration Loading**
+   The script uses Hydra to load a configuration file for subsequent steps. (Details on Hydra follow below.)
+
+2. **Dataset Initialization**
+   It loads a `LeRobotDataset` containing images and rewards. To optimize performance, a weighted random sampler is used to balance class sampling.
+
+3. **Classifier Initialization**
+   A lightweight classification head is built on top of a frozen, pretrained image encoder from HuggingFace. The classifier outputs either:
+   - A single probability (binary classification), or
+   - Logits (multi-class classification).
+
+4. **Training Loop Execution**
+   The script performs:
+   - Forward and backward passes,
+   - Optimization steps,
+   - Periodic logging, evaluation, and checkpoint saving.
+
+---
+
+## Configuring with Hydra
+
+For detailed information about Hydra usage, refer to [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md). However, note that training the reward classifier differs slightly and requires a separate configuration file.
+
+### Config File Setup
+
+The default `default.yaml` cannot launch the reward classifier training directly. Instead, you need a configuration file like [`lerobot/configs/policy/hilserl_classifier.yaml`](../../lerobot/configs/policy/hilserl_classifier.yaml), with the following adjustment:
+
+Replace the `dataset_repo_id` field with the identifier for your dataset, which contains images and sparse rewards:
+
+```yaml
+# Example: lerobot/configs/policy/reward_classifier.yaml
+dataset_repo_id: "my_dataset_repo_id"
+## Typical logs and metrics
+```
+When you start the training process, you will first see your full configuration being printed in the terminal. You can check it to make sure that you config it correctly and your config is not overrided by other files. The final configuration will also be saved with the checkpoint.
+
+After that, you will see training log like this one:
+
+```
+[2024-11-29 18:26:36,999][root][INFO] -
+Epoch 5/5
+Training:  82%|██████████████████████████████████████████████████████████████████████████████▋                 | 91/111 [00:50<00:09,  2.04it/s, loss=0.2999, acc=69.99%]
+```
+
+or evaluation log like:
+
+```
+Validation: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 28/28 [00:20<00:00,  1.37it/s]
+```
+
+### Metrics Tracking with Weights & Biases (WandB)
+
+If `wandb.enable` is set to `true`, the training and evaluation logs will also be saved in WandB. This allows you to track key metrics in real-time, including:
+
+- **Training Metrics**:
+  - `train/accuracy`
+  - `train/loss`
+  - `train/dataloading_s`
+- **Evaluation Metrics**:
+  - `eval/accuracy`
+  - `eval/loss`
+  - `eval/eval_s`
+
+#### Additional Features
+
+You can also log sample predictions during evaluation. Each logged sample will include:
+
+- The **input image**.
+- The **predicted label**.
+- The **true label**.
+- The **classifier's "confidence" (logits/probability)**.
+
+These logs can be useful for diagnosing and debugging performance issues.
+
+
+#### Generate protobuf files
+
+```bash
+python -m grpc_tools.protoc \
+    -I lerobot/scripts/server \
+    --python_out=lerobot/scripts/server \
+    --grpc_python_out=lerobot/scripts/server \
+    lerobot/scripts/server/hilserl.proto
+```

examples/1_load_lerobot_dataset.py

@@ -32,7 +32,10 @@ import torch
 from huggingface_hub import HfApi
 
 import lerobot
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import (
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+)
 
 # We ported a number of existing datasets ourselves, use this to see the list:
 print("List of available datasets:")

examples/3_train_policy.py

@@ -22,7 +22,10 @@ from pathlib import Path
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import (
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+)
 from lerobot.common.datasets.utils import dataset_to_policy_features
 from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
@@ -77,7 +80,24 @@ def main():
         # Load the previous action (-0.1), the next action to be executed (0.0),
         # and 14 future actions with a 0.1 seconds spacing. All these actions will be
         # used to supervise the policy.
-        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
+        "action": [
+            -0.1,
+            0.0,
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+            1.0,
+            1.1,
+            1.2,
+            1.3,
+            1.4,
+        ],
     }
 
     # We can then instantiate the dataset with these delta_timestamps configuration.

examples/advanced/2_calculate_validation_loss.py

@@ -26,7 +26,10 @@ import math
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import (
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+)
 from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
 
 
@@ -51,7 +54,24 @@ def main():
         # Load the previous action (-0.1), the next action to be executed (0.0),
         # and 14 future actions with a 0.1 seconds spacing. All these actions will be
         # used to calculate the loss.
-        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
+        "action": [
+            -0.1,
+            0.0,
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+            1.0,
+            1.1,
+            1.2,
+            1.3,
+            1.4,
+        ],
     }
 
     # Load the last 10% of episodes of the dataset as a validation set.

lerobot/common/datasets/compute_stats.py

@@ -19,7 +19,10 @@ from lerobot.common.datasets.utils import load_image_as_numpy
 
 
 def estimate_num_samples(
-    dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10_000, power: float = 0.75
+    dataset_len: int,
+    min_num_samples: int = 100,
+    max_num_samples: int = 10_000,
+    power: float = 0.75,
 ) -> int:
     """Heuristic to estimate the number of samples based on dataset size.
     The power controls the sample growth relative to dataset size.
@@ -123,7 +126,9 @@ def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
                     raise ValueError(f"Shape of '{k}' must be (3,1,1), but is {v.shape} instead.")
 
 
-def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
+def aggregate_feature_stats(
+    stats_ft_list: list[dict[str, dict]],
+) -> dict[str, dict[str, np.ndarray]]:
     """Aggregates stats for a single feature."""
     means = np.stack([s["mean"] for s in stats_ft_list])
     variances = np.stack([s["std"] ** 2 for s in stats_ft_list])
@@ -152,7 +157,9 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
     }
 
 
-def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
+def aggregate_stats(
+    stats_list: list[dict[str, dict]],
+) -> dict[str, dict[str, np.ndarray]]:
     """Aggregate stats from multiple compute_stats outputs into a single set of stats.
 
     The final stats will have the union of all data keys from each of the stats dicts.

lerobot/common/datasets/lerobot_dataset.py

@@ -318,7 +318,7 @@ class LeRobotDatasetMetadata:
         obj.root.mkdir(parents=True, exist_ok=False)
 
         if robot is not None:
-            features = get_features_from_robot(robot, use_videos)
+            features = {**(features or {}), **get_features_from_robot(robot)}
             robot_type = robot.robot_type
             if not all(cam.fps == fps for cam in robot.cameras.values()):
                 logging.warning(
@@ -821,7 +821,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
             if self.features[key]["dtype"] in ["image", "video"]:
                 img_path = self._get_image_file_path(
-                    episode_index=self.episode_buffer["episode_index"], image_key=key, frame_index=frame_index
+                    episode_index=self.episode_buffer["episode_index"],
+                    image_key=key,
+                    frame_index=frame_index,
                 )
                 if frame_index == 0:
                     img_path.parent.mkdir(parents=True, exist_ok=True)
@@ -867,7 +869,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
         for key, ft in self.features.items():
             # index, episode_index, task_index are already processed above, and image and video
             # are processed separately by storing image path and frame info as meta data
-            if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
+            if key in ["index", "episode_index", "task_index"] or ft["dtype"] in [
+                "image",
+                "video",
+            ]:
                 continue
             episode_buffer[key] = np.stack(episode_buffer[key])
 

lerobot/common/datasets/online_buffer.py

@@ -154,14 +154,32 @@ class OnlineBuffer(torch.utils.data.Dataset):
             OnlineBuffer.NEXT_INDEX_KEY: {"dtype": np.dtype("int64"), "shape": ()},
             # Since the memmap is initialized with all-zeros, this keeps track of which indices are occupied
             # with real data rather than the dummy initialization.
-            OnlineBuffer.OCCUPANCY_MASK_KEY: {"dtype": np.dtype("?"), "shape": (buffer_capacity,)},
-            OnlineBuffer.INDEX_KEY: {"dtype": np.dtype("int64"), "shape": (buffer_capacity,)},
-            OnlineBuffer.FRAME_INDEX_KEY: {"dtype": np.dtype("int64"), "shape": (buffer_capacity,)},
-            OnlineBuffer.EPISODE_INDEX_KEY: {"dtype": np.dtype("int64"), "shape": (buffer_capacity,)},
-            OnlineBuffer.TIMESTAMP_KEY: {"dtype": np.dtype("float64"), "shape": (buffer_capacity,)},
+            OnlineBuffer.OCCUPANCY_MASK_KEY: {
+                "dtype": np.dtype("?"),
+                "shape": (buffer_capacity,),
+            },
+            OnlineBuffer.INDEX_KEY: {
+                "dtype": np.dtype("int64"),
+                "shape": (buffer_capacity,),
+            },
+            OnlineBuffer.FRAME_INDEX_KEY: {
+                "dtype": np.dtype("int64"),
+                "shape": (buffer_capacity,),
+            },
+            OnlineBuffer.EPISODE_INDEX_KEY: {
+                "dtype": np.dtype("int64"),
+                "shape": (buffer_capacity,),
+            },
+            OnlineBuffer.TIMESTAMP_KEY: {
+                "dtype": np.dtype("float64"),
+                "shape": (buffer_capacity,),
+            },
         }
         for k, v in data_spec.items():
-            complete_data_spec[k] = {"dtype": v["dtype"], "shape": (buffer_capacity, *v["shape"])}
+            complete_data_spec[k] = {
+                "dtype": v["dtype"],
+                "shape": (buffer_capacity, *v["shape"]),
+            }
         return complete_data_spec
 
     def add_data(self, data: dict[str, np.ndarray]):

lerobot/common/datasets/push_dataset_to_hub/utils.py

@@ -77,7 +77,9 @@ def check_repo_id(repo_id: str) -> None:
 
 
 # TODO(aliberts): remove
-def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> Dict[str, torch.Tensor]:
+def calculate_episode_data_index(
+    hf_dataset: datasets.Dataset,
+) -> Dict[str, torch.Tensor]:
     """
     Calculate episode data index for the provided HuggingFace Dataset. Relies on episode_index column of hf_dataset.
 

lerobot/common/datasets/sampler.py

@@ -43,7 +43,10 @@ class EpisodeAwareSampler:
         ):
             if episode_indices_to_use is None or episode_idx in episode_indices_to_use:
                 indices.extend(
-                    range(start_index.item() + drop_n_first_frames, end_index.item() - drop_n_last_frames)
+                    range(
+                        start_index.item() + drop_n_first_frames,
+                        end_index.item() - drop_n_last_frames,
+                    )
                 )
 
         self.indices = indices

lerobot/common/datasets/utils.py

@@ -225,7 +225,10 @@ def load_episodes(local_dir: Path) -> dict:
 def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
     # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
     # is a dictionary of stats and not an integer.
-    episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
+    episode_stats = {
+        "episode_index": episode_index,
+        "stats": serialize_dict(episode_stats),
+    }
     append_jsonlines(episode_stats, local_dir / EPISODES_STATS_PATH)
 
 
@@ -409,7 +412,7 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
 
             names = ft["names"]
             # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
-            if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
+            if names is not None and names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
                 shape = (shape[2], shape[0], shape[1])
         elif key == "observation.environment_state":
             type = FeatureType.ENV
@@ -540,7 +543,10 @@ def check_timestamps_sync(
 
 
 def check_delta_timestamps(
-    delta_timestamps: dict[str, list[float]], fps: int, tolerance_s: float, raise_value_error: bool = True
+    delta_timestamps: dict[str, list[float]],
+    fps: int,
+    tolerance_s: float,
+    raise_value_error: bool = True,
 ) -> bool:
     """This will check if all the values in delta_timestamps are multiples of 1/fps +/- tolerance.
     This is to ensure that these delta_timestamps added to any timestamp from a dataset will themselves be

lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py

@@ -118,7 +118,10 @@ DATASETS = {
         "single_task": "Place the battery into the slot of the remote controller.",
         **ALOHA_STATIC_INFO,
     },
-    "aloha_static_candy": {"single_task": "Pick up the candy and unwrap it.", **ALOHA_STATIC_INFO},
+    "aloha_static_candy": {
+        "single_task": "Pick up the candy and unwrap it.",
+        **ALOHA_STATIC_INFO,
+    },
     "aloha_static_coffee": {
         "single_task": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons.",
         **ALOHA_STATIC_INFO,
@@ -167,13 +170,22 @@ DATASETS = {
         "single_task": "Pick up the plastic cup with the left arm, then pop its lid open with the right arm.",
         **ALOHA_STATIC_INFO,
     },
-    "aloha_static_ziploc_slide": {"single_task": "Slide open the ziploc bag.", **ALOHA_STATIC_INFO},
-    "aloha_sim_insertion_scripted": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
+    "aloha_static_ziploc_slide": {
+        "single_task": "Slide open the ziploc bag.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_insertion_scripted": {
+        "single_task": "Insert the peg into the socket.",
+        **ALOHA_STATIC_INFO,
+    },
     "aloha_sim_insertion_scripted_image": {
         "single_task": "Insert the peg into the socket.",
         **ALOHA_STATIC_INFO,
     },
-    "aloha_sim_insertion_human": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
+    "aloha_sim_insertion_human": {
+        "single_task": "Insert the peg into the socket.",
+        **ALOHA_STATIC_INFO,
+    },
     "aloha_sim_insertion_human_image": {
         "single_task": "Insert the peg into the socket.",
         **ALOHA_STATIC_INFO,
@@ -194,10 +206,19 @@ DATASETS = {
         "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
         **ALOHA_STATIC_INFO,
     },
-    "pusht": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
-    "pusht_image": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
+    "pusht": {
+        "single_task": "Push the T-shaped block onto the T-shaped target.",
+        **PUSHT_INFO,
+    },
+    "pusht_image": {
+        "single_task": "Push the T-shaped block onto the T-shaped target.",
+        **PUSHT_INFO,
+    },
     "unitreeh1_fold_clothes": {"single_task": "Fold the sweatshirt.", **UNITREEH_INFO},
-    "unitreeh1_rearrange_objects": {"single_task": "Put the object into the bin.", **UNITREEH_INFO},
+    "unitreeh1_rearrange_objects": {
+        "single_task": "Put the object into the bin.",
+        **UNITREEH_INFO,
+    },
     "unitreeh1_two_robot_greeting": {
         "single_task": "Greet the other robot with a high five.",
         **UNITREEH_INFO,
@@ -207,13 +228,31 @@ DATASETS = {
         **UNITREEH_INFO,
     },
     "xarm_lift_medium": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_replay": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
-    "xarm_lift_medium_replay_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
+    "xarm_lift_medium_image": {
+        "single_task": "Pick up the cube and lift it.",
+        **XARM_INFO,
+    },
+    "xarm_lift_medium_replay": {
+        "single_task": "Pick up the cube and lift it.",
+        **XARM_INFO,
+    },
+    "xarm_lift_medium_replay_image": {
+        "single_task": "Pick up the cube and lift it.",
+        **XARM_INFO,
+    },
     "xarm_push_medium": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_replay": {"single_task": "Push the cube onto the target.", **XARM_INFO},
-    "xarm_push_medium_replay_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
+    "xarm_push_medium_image": {
+        "single_task": "Push the cube onto the target.",
+        **XARM_INFO,
+    },
+    "xarm_push_medium_replay": {
+        "single_task": "Push the cube onto the target.",
+        **XARM_INFO,
+    },
+    "xarm_push_medium_replay_image": {
+        "single_task": "Push the cube onto the target.",
+        **XARM_INFO,
+    },
     "umi_cup_in_the_wild": {
         "single_task": "Put the cup on the plate.",
         "license": "apache-2.0",

lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py

@@ -379,7 +379,12 @@ def fix_lfs_video_files_tracking(work_dir: Path, lfs_untracked_videos: list[str]
     for i in range(0, len(lfs_untracked_videos), 100):
         files = lfs_untracked_videos[i : i + 100]
         try:
-            subprocess.run(["git", "rm", "--cached", *files], cwd=work_dir, capture_output=True, check=True)
+            subprocess.run(
+                ["git", "rm", "--cached", *files],
+                cwd=work_dir,
+                capture_output=True,
+                check=True,
+            )
         except subprocess.CalledProcessError as e:
             print("git rm --cached ERROR:")
             print(e.stderr)
@@ -402,7 +407,17 @@ def _lfs_clone(repo_id: str, work_dir: Path, branch: str) -> None:
     repo_url = f"https://huggingface.co/datasets/{repo_id}"
     env = {"GIT_LFS_SKIP_SMUDGE": "1"}  # Prevent downloading LFS files
     subprocess.run(
-        ["git", "clone", "--branch", branch, "--single-branch", "--depth", "1", repo_url, str(work_dir)],
+        [
+            "git",
+            "clone",
+            "--branch",
+            branch,
+            "--single-branch",
+            "--depth",
+            "1",
+            repo_url,
+            str(work_dir),
+        ],
         check=True,
         env=env,
     )
@@ -410,7 +425,11 @@ def _lfs_clone(repo_id: str, work_dir: Path, branch: str) -> None:
 
 def _get_lfs_untracked_videos(work_dir: Path, video_files: list[str]) -> list[str]:
     lfs_tracked_files = subprocess.run(
-        ["git", "lfs", "ls-files", "-n"], cwd=work_dir, capture_output=True, text=True, check=True
+        ["git", "lfs", "ls-files", "-n"],
+        cwd=work_dir,
+        capture_output=True,
+        text=True,
+        check=True,
     )
     lfs_tracked_files = set(lfs_tracked_files.stdout.splitlines())
     return [f for f in video_files if f not in lfs_tracked_files]
@@ -424,7 +443,11 @@ def get_videos_info(repo_id: str, local_dir: Path, video_keys: list[str], branch
     ]
     hub_api = HfApi()
     hub_api.snapshot_download(
-        repo_id=repo_id, repo_type="dataset", local_dir=local_dir, revision=branch, allow_patterns=video_files
+        repo_id=repo_id,
+        repo_type="dataset",
+        local_dir=local_dir,
+        revision=branch,
+        allow_patterns=video_files,
     )
     videos_info_dict = {}
     for vid_key, vid_path in zip(video_keys, video_files, strict=True):
@@ -451,7 +474,11 @@ def convert_dataset(
 
     hub_api = HfApi()
     hub_api.snapshot_download(
-        repo_id=repo_id, repo_type="dataset", revision=v1, local_dir=v1x_dir, ignore_patterns="videos*/"
+        repo_id=repo_id,
+        repo_type="dataset",
+        revision=v1,
+        local_dir=v1x_dir,
+        ignore_patterns="videos*/",
     )
     branch = "main"
     if test_branch:
@@ -509,12 +536,21 @@ def convert_dataset(
         dataset = dataset.remove_columns(video_keys)
         clean_gitattr = Path(
             hub_api.hf_hub_download(
-                repo_id=GITATTRIBUTES_REF, repo_type="dataset", local_dir=local_dir, filename=".gitattributes"
+                repo_id=GITATTRIBUTES_REF,
+                repo_type="dataset",
+                local_dir=local_dir,
+                filename=".gitattributes",
             )
         ).absolute()
         with tempfile.TemporaryDirectory() as tmp_video_dir:
             move_videos(
-                repo_id, video_keys, total_episodes, total_chunks, Path(tmp_video_dir), clean_gitattr, branch
+                repo_id,
+                video_keys,
+                total_episodes,
+                total_chunks,
+                Path(tmp_video_dir),
+                clean_gitattr,
+                branch,
             )
         videos_info = get_videos_info(repo_id, v1x_dir, video_keys=video_keys, branch=branch)
         for key in video_keys:
@@ -543,7 +579,11 @@ def convert_dataset(
 
     # Episodes
     episodes = [
-        {"episode_index": ep_idx, "tasks": tasks_by_episodes[ep_idx], "length": episode_lengths[ep_idx]}
+        {
+            "episode_index": ep_idx,
+            "tasks": tasks_by_episodes[ep_idx],
+            "length": episode_lengths[ep_idx],
+        }
         for ep_idx in episode_indices
     ]
     write_jsonlines(episodes, v20_dir / EPISODES_PATH)
@@ -572,7 +612,12 @@ def convert_dataset(
         hub_api.delete_folder(repo_id=repo_id, path_in_repo="data", repo_type="dataset", revision=branch)
 
     with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
-        hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta_data", repo_type="dataset", revision=branch)
+        hub_api.delete_folder(
+            repo_id=repo_id,
+            path_in_repo="meta_data",
+            repo_type="dataset",
+            revision=branch,
+        )
 
     with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
         hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta", repo_type="dataset", revision=branch)

lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py

@@ -37,8 +37,16 @@ import logging
 from huggingface_hub import HfApi
 
 from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.common.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
-from lerobot.common.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
+from lerobot.common.datasets.utils import (
+    EPISODES_STATS_PATH,
+    STATS_PATH,
+    load_stats,
+    write_info,
+)
+from lerobot.common.datasets.v21.convert_stats import (
+    check_aggregate_stats,
+    convert_stats,
+)
 
 V20 = "v2.0"
 V21 = "v2.1"
@@ -79,10 +87,16 @@ def convert_dataset(
 
     hub_api = HfApi()
     if hub_api.file_exists(
-        repo_id=dataset.repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"
+        repo_id=dataset.repo_id,
+        filename=STATS_PATH,
+        revision=branch,
+        repo_type="dataset",
     ):
         hub_api.delete_file(
-            path_in_repo=STATS_PATH, repo_id=dataset.repo_id, revision=branch, repo_type="dataset"
+            path_in_repo=STATS_PATH,
+            repo_id=dataset.repo_id,
+            revision=branch,
+            repo_type="dataset",
         )
 
     hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")

lerobot/common/datasets/v21/convert_stats.py

@@ -17,7 +17,11 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import numpy as np
 from tqdm import tqdm
 
-from lerobot.common.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
+from lerobot.common.datasets.compute_stats import (
+    aggregate_stats,
+    get_feature_stats,
+    sample_indices,
+)
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.common.datasets.utils import write_episode_stats
 
@@ -95,5 +99,9 @@ def check_aggregate_stats(
             if key in reference_stats and stat in reference_stats[key]:
                 err_msg = f"feature='{key}' stats='{stat}'"
                 np.testing.assert_allclose(
-                    val, reference_stats[key][stat], rtol=rtol, atol=atol, err_msg=err_msg
+                    val,
+                    reference_stats[key][stat],
+                    rtol=rtol,
+                    atol=atol,
+                    err_msg=err_msg,
                 )

lerobot/common/envs/configs.py

@@ -14,10 +14,12 @@
 
 import abc
 from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple
 
 import draccus
 
 from lerobot.common.constants import ACTION, OBS_ENV, OBS_IMAGE, OBS_IMAGES, OBS_ROBOT
+from lerobot.common.robot_devices.robots.configs import RobotConfig
 from lerobot.configs.types import FeatureType, PolicyFeature
 
 
@@ -154,3 +156,135 @@ class XarmEnv(EnvConfig):
             "visualization_height": self.visualization_height,
             "max_episode_steps": self.episode_length,
         }
+
+
+@dataclass
+class VideoRecordConfig:
+    """Configuration for video recording in ManiSkill environments."""
+
+    enabled: bool = False
+    record_dir: str = "videos"
+    trajectory_name: str = "trajectory"
+
+
+@dataclass
+class WrapperConfig:
+    """Configuration for environment wrappers."""
+
+    joint_masking_action_space: list[bool] | None = None
+
+
+@dataclass
+class EEActionSpaceConfig:
+    """Configuration parameters for end-effector action space."""
+
+    x_step_size: float
+    y_step_size: float
+    z_step_size: float
+    bounds: Dict[str, Any]  # Contains 'min' and 'max' keys with position bounds
+    use_gamepad: bool = False
+
+
+@dataclass
+class EnvWrapperConfig:
+    """Configuration for environment wrappers."""
+
+    display_cameras: bool = False
+    use_relative_joint_positions: bool = True
+    add_joint_velocity_to_observation: bool = False
+    add_ee_pose_to_observation: bool = False
+    crop_params_dict: Optional[Dict[str, Tuple[int, int, int, int]]] = None
+    resize_size: Optional[Tuple[int, int]] = None
+    control_time_s: float = 20.0
+    fixed_reset_joint_positions: Optional[Any] = None
+    reset_time_s: float = 5.0
+    joint_masking_action_space: Optional[Any] = None
+    ee_action_space_params: Optional[EEActionSpaceConfig] = None
+    use_gripper: bool = False
+    gripper_quantization_threshold: float | None = 0.8
+    gripper_penalty: float = 0.0
+    gripper_penalty_in_reward: bool = False
+    open_gripper_on_reset: bool = False
+
+
+@EnvConfig.register_subclass(name="gym_manipulator")
+@dataclass
+class HILSerlRobotEnvConfig(EnvConfig):
+    """Configuration for the HILSerlRobotEnv environment."""
+
+    robot: Optional[RobotConfig] = None
+    wrapper: Optional[EnvWrapperConfig] = None
+    fps: int = 10
+    name: str = "real_robot"
+    mode: str = None  # Either "record", "replay", None
+    repo_id: Optional[str] = None
+    dataset_root: Optional[str] = None
+    task: str = ""
+    num_episodes: int = 10  # only for record mode
+    episode: int = 0
+    device: str = "cuda"
+    push_to_hub: bool = True
+    pretrained_policy_name_or_path: Optional[str] = None
+    reward_classifier: dict[str, str | None] = field(
+        default_factory=lambda: {
+            "pretrained_path": None,
+            "config_path": None,
+        }
+    )
+
+    def gym_kwargs(self) -> dict:
+        return {}
+
+
+@EnvConfig.register_subclass("maniskill_push")
+@dataclass
+class ManiskillEnvConfig(EnvConfig):
+    """Configuration for the ManiSkill environment."""
+
+    name: str = "maniskill/pushcube"
+    task: str = "PushCube-v1"
+    image_size: int = 64
+    control_mode: str = "pd_ee_delta_pose"
+    state_dim: int = 25
+    action_dim: int = 7
+    fps: int = 200
+    episode_length: int = 50
+    obs_type: str = "rgb"
+    render_mode: str = "rgb_array"
+    render_size: int = 64
+    device: str = "cuda"
+    robot: str = "so100"  # This is a hack to make the robot config work
+    video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
+    wrapper: WrapperConfig = field(default_factory=WrapperConfig)
+    mock_gripper: bool = False
+    features: dict[str, PolicyFeature] = field(
+        default_factory=lambda: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+            "observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 64, 64)),
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(25,)),
+        }
+    )
+    features_map: dict[str, str] = field(
+        default_factory=lambda: {
+            "action": ACTION,
+            "observation.image": OBS_IMAGE,
+            "observation.state": OBS_ROBOT,
+        }
+    )
+    reward_classifier: dict[str, str | None] = field(
+        default_factory=lambda: {
+            "pretrained_path": None,
+            "config_path": None,
+        }
+    )
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {
+            "obs_type": self.obs_type,
+            "render_mode": self.render_mode,
+            "max_episode_steps": self.episode_length,
+            "control_mode": self.control_mode,
+            "sensor_configs": {"width": self.image_size, "height": self.image_size},
+            "num_envs": 1,
+        }

lerobot/common/envs/utils.py

@@ -37,29 +37,35 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
     """
     # map to expected inputs for the policy
     return_observations = {}
-    if "pixels" in observations:
-        if isinstance(observations["pixels"], dict):
-            imgs = {f"observation.images.{key}": img for key, img in observations["pixels"].items()}
-        else:
-            imgs = {"observation.image": observations["pixels"]}
+    # TODO: You have to merge all tensors from agent key and extra key
+    # You don't keep sensor param key in the observation
+    # And you keep sensor data rgb
+    for key, img in observations.items():
+        if "images" not in key:
+            continue
 
-        for imgkey, img in imgs.items():
-            # TODO(aliberts, rcadene): use transforms.ToTensor()?
+        # TODO(aliberts, rcadene): use transforms.ToTensor()?
+        if not torch.is_tensor(img):
             img = torch.from_numpy(img)
 
-            # sanity check that images are channel last
-            _, h, w, c = img.shape
-            assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
+        if img.ndim == 3:
+            img = img.unsqueeze(0)
 
-            # sanity check that images are uint8
-            assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"
+        # sanity check that images are channel last
+        _, h, w, c = img.shape
+        assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
 
-            # convert to channel first of type float32 in range [0,1]
-            img = einops.rearrange(img, "b h w c -> b c h w").contiguous()
-            img = img.type(torch.float32)
-            img /= 255
+        # sanity check that images are uint8
+        assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"
 
-            return_observations[imgkey] = img
+        # convert to channel first of type float32 in range [0,1]
+        img = einops.rearrange(img, "b h w c -> b c h w").contiguous()
+        img = img.type(torch.float32)
+        img /= 255
+
+        return_observations[key] = img
+        # obs state agent qpos and qvel
+        # image
 
     if "environment_state" in observations:
         return_observations["observation.environment_state"] = torch.from_numpy(
@@ -68,7 +74,8 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
 
     # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
     # requirement for "agent_pos"
-    return_observations["observation.state"] = torch.from_numpy(observations["agent_pos"]).float()
+    # return_observations["observation.state"] = torch.from_numpy(observations["agent_pos"]).float()
+    return_observations["observation.state"] = observations["observation.state"].float()
     return return_observations
 
 
@@ -86,7 +93,7 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
         else:
             feature = ft
 
-        policy_key = env_cfg.features_map[key]
+        policy_key = env_cfg.features_map.get(key, key)
         policy_features[policy_key] = feature
 
     return policy_features
@@ -101,7 +108,9 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
     with warnings.catch_warnings():
         warnings.simplefilter("once", UserWarning)  # Apply filter only in this function
 
-        if not (hasattr(env.envs[0], "task_description") and hasattr(env.envs[0], "task")):
+        if not (
+            hasattr(env.envs[0], "task_description") and hasattr(env.envs[0], "task")
+        ):
             warnings.warn(
                 "The environment does not have 'task_description' and 'task'. Some policies require these features.",
                 UserWarning,
@@ -115,7 +124,9 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
             )
 
 
-def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
+def add_envs_task(
+    env: gym.vector.VectorEnv, observation: dict[str, Any]
+) -> dict[str, Any]:
     """Adds task feature to the observation dict with respect to the first environment attribute."""
     if hasattr(env.envs[0], "task_description"):
         observation["task"] = env.call("task_description")

lerobot/common/optim/optimizers.py

@@ -14,8 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
+from typing import Any
 
 import draccus
 import torch
@@ -44,7 +45,7 @@ class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC):
         return "adam"
 
     @abc.abstractmethod
-    def build(self) -> torch.optim.Optimizer:
+    def build(self) -> torch.optim.Optimizer | dict[str, torch.optim.Optimizer]:
         raise NotImplementedError
 
 
@@ -94,7 +95,76 @@ class SGDConfig(OptimizerConfig):
         return torch.optim.SGD(params, **kwargs)
 
 
-def save_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None:
+@OptimizerConfig.register_subclass("multi_adam")
+@dataclass
+class MultiAdamConfig(OptimizerConfig):
+    """Configuration for multiple Adam optimizers with different parameter groups.
+
+    This creates a dictionary of Adam optimizers, each with its own hyperparameters.
+
+    Args:
+        lr: Default learning rate (used if not specified for a group)
+        weight_decay: Default weight decay (used if not specified for a group)
+        optimizer_groups: Dictionary mapping parameter group names to their hyperparameters
+        grad_clip_norm: Gradient clipping norm
+    """
+
+    lr: float = 1e-3
+    weight_decay: float = 0.0
+    grad_clip_norm: float = 10.0
+    optimizer_groups: dict[str, dict[str, Any]] = field(default_factory=dict)
+
+    def build(self, params_dict: dict[str, list]) -> dict[str, torch.optim.Optimizer]:
+        """Build multiple Adam optimizers.
+
+        Args:
+            params_dict: Dictionary mapping parameter group names to lists of parameters
+                         The keys should match the keys in optimizer_groups
+
+        Returns:
+            Dictionary mapping parameter group names to their optimizers
+        """
+        optimizers = {}
+
+        for name, params in params_dict.items():
+            # Get group-specific hyperparameters or use defaults
+            group_config = self.optimizer_groups.get(name, {})
+
+            # Create optimizer with merged parameters (defaults + group-specific)
+            optimizer_kwargs = {
+                "lr": group_config.get("lr", self.lr),
+                "betas": group_config.get("betas", (0.9, 0.999)),
+                "eps": group_config.get("eps", 1e-5),
+                "weight_decay": group_config.get("weight_decay", self.weight_decay),
+            }
+
+            optimizers[name] = torch.optim.Adam(params, **optimizer_kwargs)
+
+        return optimizers
+
+
+def save_optimizer_state(
+    optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path
+) -> None:
+    """Save optimizer state to disk.
+
+    Args:
+        optimizer: Either a single optimizer or a dictionary of optimizers.
+        save_dir: Directory to save the optimizer state.
+    """
+    if isinstance(optimizer, dict):
+        # Handle dictionary of optimizers
+        for name, opt in optimizer.items():
+            optimizer_dir = save_dir / name
+            optimizer_dir.mkdir(exist_ok=True, parents=True)
+            _save_single_optimizer_state(opt, optimizer_dir)
+    else:
+        # Handle single optimizer
+        _save_single_optimizer_state(optimizer, save_dir)
+
+
+def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None:
+    """Save a single optimizer's state to disk."""
     state = optimizer.state_dict()
     param_groups = state.pop("param_groups")
     flat_state = flatten_dict(state)
@@ -102,11 +172,44 @@ def save_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> No
     write_json(param_groups, save_dir / OPTIMIZER_PARAM_GROUPS)
 
 
-def load_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> torch.optim.Optimizer:
+def load_optimizer_state(
+    optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path
+) -> torch.optim.Optimizer | dict[str, torch.optim.Optimizer]:
+    """Load optimizer state from disk.
+
+    Args:
+        optimizer: Either a single optimizer or a dictionary of optimizers.
+        save_dir: Directory to load the optimizer state from.
+
+    Returns:
+        The updated optimizer(s) with loaded state.
+    """
+    if isinstance(optimizer, dict):
+        # Handle dictionary of optimizers
+        loaded_optimizers = {}
+        for name, opt in optimizer.items():
+            optimizer_dir = save_dir / name
+            if optimizer_dir.exists():
+                loaded_optimizers[name] = _load_single_optimizer_state(opt, optimizer_dir)
+            else:
+                loaded_optimizers[name] = opt
+        return loaded_optimizers
+    else:
+        # Handle single optimizer
+        return _load_single_optimizer_state(optimizer, save_dir)
+
+
+def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> torch.optim.Optimizer:
+    """Load a single optimizer's state from disk."""
     current_state_dict = optimizer.state_dict()
     flat_state = load_file(save_dir / OPTIMIZER_STATE)
     state = unflatten_dict(flat_state)
-    loaded_state_dict = {"state": {int(k): v for k, v in state["state"].items()}}
+
+    # Handle case where 'state' key might not exist (for newly created optimizers)
+    if "state" in state:
+        loaded_state_dict = {"state": {int(k): v for k, v in state["state"].items()}}
+    else:
+        loaded_state_dict = {"state": {}}
 
     if "param_groups" in current_state_dict:
         param_groups = deserialize_json_into_object(

lerobot/common/optim/schedulers.py

@@ -49,7 +49,11 @@ class DiffuserSchedulerConfig(LRSchedulerConfig):
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
         from diffusers.optimization import get_scheduler
 
-        kwargs = {**asdict(self), "num_training_steps": num_training_steps, "optimizer": optimizer}
+        kwargs = {
+            **asdict(self),
+            "num_training_steps": num_training_steps,
+            "optimizer": optimizer,
+        }
         return get_scheduler(**kwargs)
 
 
@@ -71,7 +75,10 @@ class VQBeTSchedulerConfig(LRSchedulerConfig):
                 progress = float(adjusted_step - self.num_warmup_steps) / float(
                     max(1, num_training_steps - self.num_warmup_steps)
                 )
-                return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)))
+                return max(
+                    0.0,
+                    0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)),
+                )
 
         return LambdaLR(optimizer, lr_lambda, -1)
 

lerobot/common/policies/act/modeling_act.py

@@ -241,7 +241,9 @@ class ACTTemporalEnsembler:
             # Note: The last dimension is unsqueeze to make sure we can broadcast properly for tensor
             # operations later.
             self.ensembled_actions_count = torch.ones(
-                (self.chunk_size, 1), dtype=torch.long, device=self.ensembled_actions.device
+                (self.chunk_size, 1),
+                dtype=torch.long,
+                device=self.ensembled_actions.device,
             )
         else:
             # self.ensembled_actions will have shape (batch_size, chunk_size - 1, action_dim). Compute
@@ -253,7 +255,10 @@ class ACTTemporalEnsembler:
             # The last action, which has no prior online average, needs to get concatenated onto the end.
             self.ensembled_actions = torch.cat([self.ensembled_actions, actions[:, -1:]], dim=1)
             self.ensembled_actions_count = torch.cat(
-                [self.ensembled_actions_count, torch.ones_like(self.ensembled_actions_count[-1:])]
+                [
+                    self.ensembled_actions_count,
+                    torch.ones_like(self.ensembled_actions_count[-1:]),
+                ]
             )
         # "Consume" the first action.
         action, self.ensembled_actions, self.ensembled_actions_count = (
@@ -333,7 +338,11 @@ class ACT(nn.Module):
         # Backbone for image feature extraction.
         if self.config.image_features:
             backbone_model = getattr(torchvision.models, config.vision_backbone)(
-                replace_stride_with_dilation=[False, False, config.replace_final_stride_with_dilation],
+                replace_stride_with_dilation=[
+                    False,
+                    False,
+                    config.replace_final_stride_with_dilation,
+                ],
                 weights=config.pretrained_backbone_weights,
                 norm_layer=FrozenBatchNorm2d,
             )
@@ -427,7 +436,11 @@ class ACT(nn.Module):
             action_embed = self.vae_encoder_action_input_proj(batch["action"])  # (B, S, D)
 
             if self.config.robot_state_feature:
-                vae_encoder_input = [cls_embed, robot_state_embed, action_embed]  # (B, S+2, D)
+                vae_encoder_input = [
+                    cls_embed,
+                    robot_state_embed,
+                    action_embed,
+                ]  # (B, S+2, D)
             else:
                 vae_encoder_input = [cls_embed, action_embed]
             vae_encoder_input = torch.cat(vae_encoder_input, axis=1)
@@ -540,7 +553,10 @@ class ACTEncoder(nn.Module):
         self.norm = nn.LayerNorm(config.dim_model) if config.pre_norm else nn.Identity()
 
     def forward(
-        self, x: Tensor, pos_embed: Tensor | None = None, key_padding_mask: Tensor | None = None
+        self,
+        x: Tensor,
+        pos_embed: Tensor | None = None,
+        key_padding_mask: Tensor | None = None,
     ) -> Tensor:
         for layer in self.layers:
             x = layer(x, pos_embed=pos_embed, key_padding_mask=key_padding_mask)
@@ -603,7 +619,10 @@ class ACTDecoder(nn.Module):
     ) -> Tensor:
         for layer in self.layers:
             x = layer(
-                x, encoder_out, decoder_pos_embed=decoder_pos_embed, encoder_pos_embed=encoder_pos_embed
+                x,
+                encoder_out,
+                decoder_pos_embed=decoder_pos_embed,
+                encoder_pos_embed=encoder_pos_embed,
             )
         if self.norm is not None:
             x = self.norm(x)

lerobot/common/policies/diffusion/modeling_diffusion.py

@@ -209,7 +209,10 @@ class DiffusionModel(nn.Module):
 
     # ========= inference  ============
     def conditional_sample(
-        self, batch_size: int, global_cond: Tensor | None = None, generator: torch.Generator | None = None
+        self,
+        batch_size: int,
+        global_cond: Tensor | None = None,
+        generator: torch.Generator | None = None,
     ) -> Tensor:
         device = get_device_from_parameters(self)
         dtype = get_dtype_from_parameters(self)
@@ -254,7 +257,10 @@ class DiffusionModel(nn.Module):
                 # Separate batch and sequence dims back out. The camera index dim gets absorbed into the
                 # feature dim (effectively concatenating the camera features).
                 img_features = einops.rearrange(
-                    img_features_list, "(n b s) ... -> b s (n ...)", b=batch_size, s=n_obs_steps
+                    img_features_list,
+                    "(n b s) ... -> b s (n ...)",
+                    b=batch_size,
+                    s=n_obs_steps,
                 )
             else:
                 # Combine batch, sequence, and "which camera" dims before passing to shared encoder.
@@ -264,7 +270,10 @@ class DiffusionModel(nn.Module):
                 # Separate batch dim and sequence dim back out. The camera index dim gets absorbed into the
                 # feature dim (effectively concatenating the camera features).
                 img_features = einops.rearrange(
-                    img_features, "(b s n) ... -> b s (n ...)", b=batch_size, s=n_obs_steps
+                    img_features,
+                    "(b s n) ... -> b s (n ...)",
+                    b=batch_size,
+                    s=n_obs_steps,
                 )
             global_cond_feats.append(img_features)
 
@@ -515,7 +524,9 @@ class DiffusionRgbEncoder(nn.Module):
 
 
 def _replace_submodules(
-    root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]
+    root_module: nn.Module,
+    predicate: Callable[[nn.Module], bool],
+    func: Callable[[nn.Module], nn.Module],
 ) -> nn.Module:
     """
     Args:
@@ -633,10 +644,14 @@ class DiffusionConditionalUnet1d(nn.Module):
         self.mid_modules = nn.ModuleList(
             [
                 DiffusionConditionalResidualBlock1d(
-                    config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs
+                    config.down_dims[-1],
+                    config.down_dims[-1],
+                    **common_res_block_kwargs,
                 ),
                 DiffusionConditionalResidualBlock1d(
-                    config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs
+                    config.down_dims[-1],
+                    config.down_dims[-1],
+                    **common_res_block_kwargs,
                 ),
             ]
         )

lerobot/common/policies/factory.py

@@ -24,6 +24,7 @@ from lerobot.common.envs.configs import EnvConfig
 from lerobot.common.envs.utils import env_to_policy_features
 from lerobot.common.policies.act.configuration_act import ACTConfig
 from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.common.policies.hilserl.classifier.configuration_classifier import ClassifierConfig
 from lerobot.common.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.common.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.common.policies.pretrained import PreTrainedPolicy
@@ -59,6 +60,14 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
         from lerobot.common.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
 
         return PI0FASTPolicy
+    elif name == "sac":
+        from lerobot.common.policies.sac.modeling_sac import SACPolicy
+
+        return SACPolicy
+    elif name == "hilserl_classifier":
+        from lerobot.common.policies.hilserl.classifier.modeling_classifier import Classifier
+
+        return Classifier
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -76,6 +85,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return PI0Config(**kwargs)
     elif policy_type == "pi0fast":
         return PI0FASTConfig(**kwargs)
+    elif policy_type == "hilserl_classifier":
+        return ClassifierConfig(**kwargs)
     else:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 

lerobot/common/policies/hilserl/classifier/configuration_classifier.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+from typing import List
+
+from lerobot.common.optim.optimizers import AdamWConfig, OptimizerConfig
+from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.configs.policies import PreTrainedConfig
+
+
+@PreTrainedConfig.register_subclass(name="hilserl_classifier")
+@dataclass
+class ClassifierConfig(PreTrainedConfig):
+    """Configuration for the Classifier model."""
+
+    name: str = "hilserl_classifier"
+    num_classes: int = 2
+    hidden_dim: int = 256
+    dropout_rate: float = 0.1
+    model_name: str = "helper2424/resnet10"
+    device: str = "cpu"
+    model_type: str = "cnn"  # "transformer" or "cnn"
+    num_cameras: int = 2
+    learning_rate: float = 1e-4
+    normalization_mode = None
+    # output_features: Dict[str, PolicyFeature] = field(
+    #     default_factory=lambda: {"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(1,))}
+    # )
+
+    @property
+    def observation_delta_indices(self) -> List | None:
+        return None
+
+    @property
+    def action_delta_indices(self) -> List | None:
+        return None
+
+    @property
+    def reward_delta_indices(self) -> List | None:
+        return None
+
+    def get_optimizer_preset(self) -> OptimizerConfig:
+        return AdamWConfig(
+            lr=self.learning_rate,
+            weight_decay=0.01,
+            grad_clip_norm=1.0,
+        )
+
+    def get_scheduler_preset(self) -> LRSchedulerConfig | None:
+        return None
+
+    def validate_features(self) -> None:
+        """Validate feature configurations."""
+        # Classifier doesn't need specific feature validation
+        pass

lerobot/common/policies/hilserl/classifier/modeling_classifier.py

@@ -0,0 +1,237 @@
+import logging
+from typing import Dict, Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from lerobot.common.constants import OBS_IMAGE
+from lerobot.common.policies.hilserl.classifier.configuration_classifier import (
+    ClassifierConfig,
+)
+from lerobot.common.policies.normalize import Normalize, Unnormalize
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+class ClassifierOutput:
+    """Wrapper for classifier outputs with additional metadata."""
+
+    def __init__(
+        self,
+        logits: Tensor,
+        probabilities: Optional[Tensor] = None,
+        hidden_states: Optional[Tensor] = None,
+    ):
+        self.logits = logits
+        self.probabilities = probabilities
+        self.hidden_states = hidden_states
+
+    def __repr__(self):
+        return (
+            f"ClassifierOutput(logits={self.logits}, "
+            f"probabilities={self.probabilities}, "
+            f"hidden_states={self.hidden_states})"
+        )
+
+
+class Classifier(PreTrainedPolicy):
+    """Image classifier built on top of a pre-trained encoder."""
+
+    name = "hilserl_classifier"
+    config_class = ClassifierConfig
+
+    def __init__(
+        self,
+        config: ClassifierConfig,
+        dataset_stats: Dict[str, Dict[str, Tensor]] | None = None,
+    ):
+        from transformers import AutoModel
+
+        super().__init__(config)
+        self.config = config
+
+        # Initialize normalization (standardized with the policy framework)
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
+        # Set up encoder
+        encoder = AutoModel.from_pretrained(self.config.model_name, trust_remote_code=True)
+        # Extract vision model if we're given a multimodal model
+        if hasattr(encoder, "vision_model"):
+            logging.info("Multimodal model detected - using vision encoder only")
+            self.encoder = encoder.vision_model
+            self.vision_config = encoder.config.vision_config
+        else:
+            self.encoder = encoder
+            self.vision_config = getattr(encoder, "config", None)
+
+        # Model type from config
+        self.is_cnn = self.config.model_type == "cnn"
+
+        # For CNNs, initialize backbone
+        if self.is_cnn:
+            self._setup_cnn_backbone()
+
+        self._freeze_encoder()
+        self._build_classifier_head()
+
+    def _setup_cnn_backbone(self):
+        """Set up CNN encoder"""
+        if hasattr(self.encoder, "fc"):
+            self.feature_dim = self.encoder.fc.in_features
+            self.encoder = nn.Sequential(*list(self.encoder.children())[:-1])
+        elif hasattr(self.encoder.config, "hidden_sizes"):
+            self.feature_dim = self.encoder.config.hidden_sizes[-1]  # Last channel dimension
+        else:
+            raise ValueError("Unsupported CNN architecture")
+
+    def _freeze_encoder(self) -> None:
+        """Freeze the encoder parameters."""
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+
+    def _build_classifier_head(self) -> None:
+        """Initialize the classifier head architecture."""
+        # Get input dimension based on model type
+        if self.is_cnn:
+            input_dim = self.feature_dim
+        else:  # Transformer models
+            if hasattr(self.encoder.config, "hidden_size"):
+                input_dim = self.encoder.config.hidden_size
+            else:
+                raise ValueError("Unsupported transformer architecture since hidden_size is not found")
+
+        self.classifier_head = nn.Sequential(
+            nn.Linear(input_dim * self.config.num_cameras, self.config.hidden_dim),
+            nn.Dropout(self.config.dropout_rate),
+            nn.LayerNorm(self.config.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(
+                self.config.hidden_dim,
+                1 if self.config.num_classes == 2 else self.config.num_classes,
+            ),
+        )
+
+    def _get_encoder_output(self, x: torch.Tensor) -> torch.Tensor:
+        """Extract the appropriate output from the encoder."""
+        with torch.no_grad():
+            if self.is_cnn:
+                # The HF ResNet applies pooling internally
+                outputs = self.encoder(x)
+                # Get pooled output directly
+                features = outputs.pooler_output
+
+                if features.dim() > 2:
+                    features = features.squeeze(-1).squeeze(-1)
+                return features
+            else:  # Transformer models
+                outputs = self.encoder(x)
+                if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+                    return outputs.pooler_output
+                return outputs.last_hidden_state[:, 0, :]
+
+    def extract_images_and_labels(self, batch: Dict[str, Tensor]) -> Tuple[list, Tensor]:
+        """Extract image tensors and label tensors from batch."""
+        # Find image keys in input features
+        image_keys = [key for key in self.config.input_features if key.startswith(OBS_IMAGE)]
+
+        # Extract the images and labels
+        images = [batch[key] for key in image_keys]
+        labels = batch["next.reward"]
+
+        return images, labels
+
+    def predict(self, xs: list) -> ClassifierOutput:
+        """Forward pass of the classifier for inference."""
+        encoder_outputs = torch.hstack([self._get_encoder_output(x) for x in xs])
+        logits = self.classifier_head(encoder_outputs)
+
+        if self.config.num_classes == 2:
+            logits = logits.squeeze(-1)
+            probabilities = torch.sigmoid(logits)
+        else:
+            probabilities = torch.softmax(logits, dim=-1)
+
+        return ClassifierOutput(logits=logits, probabilities=probabilities, hidden_states=encoder_outputs)
+
+    def forward(self, batch: Dict[str, Tensor]) -> Tuple[Tensor, Dict[str, Tensor]]:
+        """Standard forward pass for training compatible with train.py."""
+        # Normalize inputs if needed
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+
+        # Extract images and labels
+        images, labels = self.extract_images_and_labels(batch)
+
+        # Get predictions
+        outputs = self.predict(images)
+
+        # Calculate loss
+        if self.config.num_classes == 2:
+            # Binary classification
+            loss = nn.functional.binary_cross_entropy_with_logits(outputs.logits, labels)
+            predictions = (torch.sigmoid(outputs.logits) > 0.5).float()
+        else:
+            # Multi-class classification
+            loss = nn.functional.cross_entropy(outputs.logits, labels.long())
+            predictions = torch.argmax(outputs.logits, dim=1)
+
+        # Calculate accuracy for logging
+        correct = (predictions == labels).sum().item()
+        total = labels.size(0)
+        accuracy = 100 * correct / total
+
+        # Return loss and metrics for logging
+        output_dict = {
+            "accuracy": accuracy,
+            "correct": correct,
+            "total": total,
+        }
+
+        return loss, output_dict
+
+    def predict_reward(self, batch, threshold=0.6):
+        """Legacy method for compatibility."""
+        images, _ = self.extract_images_and_labels(batch)
+        if self.config.num_classes == 2:
+            probs = self.predict(images).probabilities
+            logging.debug(f"Predicted reward images: {probs}")
+            return (probs > threshold).float()
+        else:
+            return torch.argmax(self.predict(images).probabilities, dim=1)
+
+    # Methods required by PreTrainedPolicy abstract class
+
+    def get_optim_params(self) -> dict:
+        """Return optimizer parameters for the policy."""
+        return {
+            "params": self.parameters(),
+            "lr": getattr(self.config, "learning_rate", 1e-4),
+            "weight_decay": getattr(self.config, "weight_decay", 0.01),
+        }
+
+    def reset(self):
+        """Reset any stateful components (required by PreTrainedPolicy)."""
+        # Classifier doesn't have stateful components that need resetting
+        pass
+
+    def select_action(self, batch: Dict[str, Tensor]) -> Tensor:
+        """Return action (class prediction) based on input observation."""
+        images, _ = self.extract_images_and_labels(batch)
+
+        with torch.no_grad():
+            outputs = self.predict(images)
+
+            if self.config.num_classes == 2:
+                # For binary classification return 0 or 1
+                return (outputs.probabilities > 0.5).float()
+            else:
+                # For multi-class return the predicted class
+                return torch.argmax(outputs.probabilities, dim=1)

lerobot/common/policies/hilserl/configuration_hilserl.py

@@ -0,0 +1,23 @@
+#!/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.
+
+from dataclasses import dataclass
+
+
+@dataclass
+class HILSerlConfig:
+    pass

lerobot/common/policies/hilserl/modeling_hilserl.py

@@ -0,0 +1,29 @@
+#!/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.
+
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class HILSerlPolicy(
+    nn.Module,
+    PyTorchModelHubMixin,
+    library_name="lerobot",
+    repo_url="https://github.com/huggingface/lerobot",
+    tags=["robotics", "hilserl"],
+):
+    pass

lerobot/common/policies/normalize.py

@@ -79,28 +79,46 @@ def create_stats_buffers(
             )
 
         # TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
-        if stats:
-            if isinstance(stats[key]["mean"], np.ndarray):
-                if norm_mode is NormalizationMode.MEAN_STD:
+        if stats and key in stats:
+            if norm_mode is NormalizationMode.MEAN_STD:
+                if "mean" not in stats[key] or "std" not in stats[key]:
+                    raise ValueError(
+                        f"Missing 'mean' or 'std' in stats for key {key} with MEAN_STD normalization"
+                    )
+
+                if isinstance(stats[key]["mean"], np.ndarray):
                     buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
                     buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
-                elif norm_mode is NormalizationMode.MIN_MAX:
-                    buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
-                    buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
-            elif isinstance(stats[key]["mean"], torch.Tensor):
-                # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
-                # tensors anywhere (for example, when we use the same stats for normalization and
-                # unnormalization). See the logic here
-                # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
-                if norm_mode is NormalizationMode.MEAN_STD:
+                elif isinstance(stats[key]["mean"], torch.Tensor):
+                    # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
+                    # tensors anywhere (for example, when we use the same stats for normalization and
+                    # unnormalization). See the logic here
+                    # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
                     buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
                     buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
-                elif norm_mode is NormalizationMode.MIN_MAX:
+                else:
+                    type_ = type(stats[key]["mean"])
+                    raise ValueError(
+                        f"np.ndarray or torch.Tensor expected for 'mean', but type is '{type_}' instead."
+                    )
+
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                if "min" not in stats[key] or "max" not in stats[key]:
+                    raise ValueError(
+                        f"Missing 'min' or 'max' in stats for key {key} with MIN_MAX normalization"
+                    )
+
+                if isinstance(stats[key]["min"], np.ndarray):
+                    buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
+                    buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
+                elif isinstance(stats[key]["min"], torch.Tensor):
                     buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
                     buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
-            else:
-                type_ = type(stats[key]["mean"])
-                raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
+                else:
+                    type_ = type(stats[key]["min"])
+                    raise ValueError(
+                        f"np.ndarray or torch.Tensor expected for 'min', but type is '{type_}' instead."
+                    )
 
         stats_buffers[key] = buffer
     return stats_buffers
@@ -149,12 +167,13 @@ class Normalize(nn.Module):
             setattr(self, "buffer_" + key.replace(".", "_"), buffer)
 
     # TODO(rcadene): should we remove torch.no_grad?
-    @torch.no_grad
+    # @torch.no_grad
     def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
         batch = dict(batch)  # shallow copy avoids mutating the input batch
         for key, ft in self.features.items():
             if key not in batch:
                 # FIXME(aliberts, rcadene): This might lead to silent fail!
+                # NOTE: (azouitine) This continues help us for instantiation SACPolicy
                 continue
 
             norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
@@ -223,7 +242,7 @@ class Unnormalize(nn.Module):
             setattr(self, "buffer_" + key.replace(".", "_"), buffer)
 
     # TODO(rcadene): should we remove torch.no_grad?
-    @torch.no_grad
+    # @torch.no_grad
     def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
         batch = dict(batch)  # shallow copy avoids mutating the input batch
         for key, ft in self.features.items():

lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py

@@ -61,7 +61,11 @@ from lerobot.common.policies.pi0.conversion_scripts.conversion_utils import (
 )
 from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
 
-PRECISIONS = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}
+PRECISIONS = {
+    "bfloat16": torch.bfloat16,
+    "float32": torch.float32,
+    "float16": torch.float16,
+}
 
 
 def slice_paligemma_state_dict(state_dict, config):

lerobot/common/policies/pi0/flex_attention.py

@@ -48,18 +48,32 @@ def flex_attention_forward(
 
     key_states = key_states[:, :, :, None, :]
     key_states = key_states.expand(
-        batch_size, key_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
+        batch_size,
+        key_states.shape[1],
+        num_key_value_heads,
+        num_key_value_groups,
+        head_dim,
     )
     key_states = key_states.reshape(
-        batch_size, key_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
+        batch_size,
+        key_states.shape[1],
+        num_key_value_heads * num_key_value_groups,
+        head_dim,
     )
 
     value_states = value_states[:, :, :, None, :]
     value_states = value_states.expand(
-        batch_size, value_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
+        batch_size,
+        value_states.shape[1],
+        num_key_value_heads,
+        num_key_value_groups,
+        head_dim,
     )
     value_states = value_states.reshape(
-        batch_size, value_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
+        batch_size,
+        value_states.shape[1],
+        num_key_value_heads * num_key_value_groups,
+        head_dim,
     )
 
     query_states = query_states.transpose(1, 2)

lerobot/common/policies/pi0/modeling_pi0.py

@@ -69,7 +69,11 @@ from lerobot.common.utils.utils import get_safe_dtype
 
 
 def create_sinusoidal_pos_embedding(
-    time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
+    time: torch.tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device="cpu",
 ) -> Tensor:
     """Computes sine-cosine positional embedding vectors for scalar positions."""
     if dimension % 2 != 0:
@@ -577,7 +581,11 @@ class PI0FlowMatching(nn.Module):
 
         # Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
         time_emb = create_sinusoidal_pos_embedding(
-            timestep, self.config.proj_width, min_period=4e-3, max_period=4.0, device=device
+            timestep,
+            self.config.proj_width,
+            min_period=4e-3,
+            max_period=4.0,
+            device=device,
         )
         time_emb = time_emb.type(dtype=dtype)
 
@@ -609,7 +617,15 @@ class PI0FlowMatching(nn.Module):
         return embs, pad_masks, att_masks
 
     def forward(
-        self, images, img_masks, lang_tokens, lang_masks, state, actions, noise=None, time=None
+        self,
+        images,
+        img_masks,
+        lang_tokens,
+        lang_masks,
+        state,
+        actions,
+        noise=None,
+        time=None,
     ) -> Tensor:
         """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)"""
         if noise is None:
@@ -655,7 +671,11 @@ class PI0FlowMatching(nn.Module):
         device = state.device
 
         if noise is None:
-            actions_shape = (bsize, self.config.n_action_steps, self.config.max_action_dim)
+            actions_shape = (
+                bsize,
+                self.config.n_action_steps,
+                self.config.max_action_dim,
+            )
             noise = self.sample_noise(actions_shape, device)
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(

lerobot/common/policies/pi0/paligemma_with_expert.py

@@ -293,12 +293,18 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
                     # in `transformers`. (molbap)
                     key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
                     value_states = torch.cat(
-                        [past_key_values[layer_idx]["value_states"], value_states], dim=1
+                        [past_key_values[layer_idx]["value_states"], value_states],
+                        dim=1,
                     )
 
             attention_interface = self.get_attention_interface()
             att_output = attention_interface(
-                attention_mask, batch_size, head_dim, query_states, key_states, value_states
+                attention_mask,
+                batch_size,
+                head_dim,
+                query_states,
+                key_states,
+                value_states,
             )
             att_output = att_output.to(dtype=torch.bfloat16)
 
@@ -358,12 +364,24 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
         return attention_interface
 
     def flash_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
+        self,
+        attention_mask,
+        batch_size,
+        head_dim,
+        query_states,
+        key_states,
+        value_states,
     ):
         raise NotImplementedError("FA2 is not implemented (yet)")
 
     def eager_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
+        self,
+        attention_mask,
+        batch_size,
+        head_dim,
+        query_states,
+        key_states,
+        value_states,
     ):
         num_att_heads = self.config.paligemma_config.text_config.num_attention_heads
         num_key_value_heads = self.config.paligemma_config.text_config.num_key_value_heads
@@ -375,17 +393,31 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
         sequence_length = key_states.shape[1]
 
         key_states = key_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+            batch_size,
+            sequence_length,
+            num_key_value_heads,
+            num_key_value_groups,
+            head_dim,
         )
         key_states = key_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+            batch_size,
+            sequence_length,
+            num_key_value_heads * num_key_value_groups,
+            head_dim,
         )
 
         value_states = value_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+            batch_size,
+            sequence_length,
+            num_key_value_heads,
+            num_key_value_groups,
+            head_dim,
         )
         value_states = value_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+            batch_size,
+            sequence_length,
+            num_key_value_heads * num_key_value_groups,
+            head_dim,
         )
 
         # Attention here is upcasted to float32 to match the original eager implementation.

lerobot/common/policies/sac/configuration_sac.py

@@ -0,0 +1,229 @@
+#!/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.
+
+from dataclasses import dataclass, field
+
+from lerobot.common.optim.optimizers import MultiAdamConfig
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+
+
+@dataclass
+class ConcurrencyConfig:
+    actor: str = "threads"
+    learner: str = "threads"
+
+
+@dataclass
+class ActorLearnerConfig:
+    learner_host: str = "127.0.0.1"
+    learner_port: int = 50051
+    policy_parameters_push_frequency: int = 4
+
+
+@dataclass
+class CriticNetworkConfig:
+    hidden_dims: list[int] = field(default_factory=lambda: [256, 256])
+    activate_final: bool = True
+    final_activation: str | None = None
+
+
+@dataclass
+class ActorNetworkConfig:
+    hidden_dims: list[int] = field(default_factory=lambda: [256, 256])
+    activate_final: bool = True
+
+
+@dataclass
+class PolicyConfig:
+    use_tanh_squash: bool = True
+    log_std_min: float = 1e-5
+    log_std_max: float = 10.0
+    init_final: float = 0.05
+
+
+@PreTrainedConfig.register_subclass("sac")
+@dataclass
+class SACConfig(PreTrainedConfig):
+    """Soft Actor-Critic (SAC) configuration.
+
+    SAC is an off-policy actor-critic deep RL algorithm based on the maximum entropy
+    reinforcement learning framework. It learns a policy and a Q-function simultaneously
+    using experience collected from the environment.
+
+    This configuration class contains all the parameters needed to define a SAC agent,
+    including network architectures, optimization settings, and algorithm-specific
+    hyperparameters.
+
+    Args:
+        actor_network: Configuration for the actor network architecture.
+        critic_network: Configuration for the critic network architecture.
+        policy: Configuration for the policy parameters.
+        n_obs_steps: Number of observation steps to consider.
+        normalization_mapping: Mapping of feature types to normalization modes.
+        dataset_stats: Statistics for normalizing different types of inputs.
+        input_features: Dictionary of input features with their types and shapes.
+        output_features: Dictionary of output features with their types and shapes.
+        camera_number: Number of cameras used for visual observations.
+        device: Device to run the model on (e.g., "cuda", "cpu").
+        storage_device: Device to store the model on.
+        vision_encoder_name: Name of the vision encoder model.
+        freeze_vision_encoder: Whether to freeze the vision encoder during training.
+        image_encoder_hidden_dim: Hidden dimension size for the image encoder.
+        shared_encoder: Whether to use a shared encoder for actor and critic.
+        num_discrete_actions: Number of discrete actions, eg for gripper actions.
+        image_embedding_pooling_dim: Dimension of the image embedding pooling.
+        concurrency: Configuration for concurrency settings.
+        actor_learner: Configuration for actor-learner architecture.
+        online_steps: Number of steps for online training.
+        online_env_seed: Seed for the online environment.
+        online_buffer_capacity: Capacity of the online replay buffer.
+        offline_buffer_capacity: Capacity of the offline replay buffer.
+        async_prefetch: Whether to use asynchronous prefetching for the buffers.
+        online_step_before_learning: Number of steps before learning starts.
+        policy_update_freq: Frequency of policy updates.
+        discount: Discount factor for the SAC algorithm.
+        temperature_init: Initial temperature value.
+        num_critics: Number of critics in the ensemble.
+        num_subsample_critics: Number of subsampled critics for training.
+        critic_lr: Learning rate for the critic network.
+        actor_lr: Learning rate for the actor network.
+        temperature_lr: Learning rate for the temperature parameter.
+        critic_target_update_weight: Weight for the critic target update.
+        utd_ratio: Update-to-data ratio for the UTD algorithm.
+        state_encoder_hidden_dim: Hidden dimension size for the state encoder.
+        latent_dim: Dimension of the latent space.
+        target_entropy: Target entropy for the SAC algorithm.
+        use_backup_entropy: Whether to use backup entropy for the SAC algorithm.
+        grad_clip_norm: Gradient clipping norm for the SAC algorithm.
+    """
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.MEAN_STD,
+            "STATE": NormalizationMode.MIN_MAX,
+            "ENV": NormalizationMode.MIN_MAX,
+            "ACTION": NormalizationMode.MIN_MAX,
+        }
+    )
+
+    dataset_stats: dict[str, dict[str, list[float]]] | None = field(
+        default_factory=lambda: {
+            "observation.image": {
+                "mean": [0.485, 0.456, 0.406],
+                "std": [0.229, 0.224, 0.225],
+            },
+            "observation.state": {
+                "min": [0.0, 0.0],
+                "max": [1.0, 1.0],
+            },
+            "action": {
+                "min": [0.0, 0.0, 0.0],
+                "max": [1.0, 1.0, 1.0],
+            },
+        }
+    )
+
+    # Architecture specifics
+    camera_number: int = 1
+    device: str = "cuda"
+    storage_device: str = "cpu"
+    # Set to "helper2424/resnet10" for hil serl
+    vision_encoder_name: str | None = None
+    freeze_vision_encoder: bool = True
+    image_encoder_hidden_dim: int = 32
+    shared_encoder: bool = True
+    num_discrete_actions: int | None = None
+    image_embedding_pooling_dim: int = 8
+
+    # Training parameter
+    online_steps: int = 1000000
+    online_env_seed: int = 10000
+    online_buffer_capacity: int = 100000
+    offline_buffer_capacity: int = 100000
+    async_prefetch: bool = False
+    online_step_before_learning: int = 100
+    policy_update_freq: int = 1
+
+    # SAC algorithm parameters
+    discount: float = 0.99
+    temperature_init: float = 1.0
+    num_critics: int = 2
+    num_subsample_critics: int | None = None
+    critic_lr: float = 3e-4
+    actor_lr: float = 3e-4
+    temperature_lr: float = 3e-4
+    critic_target_update_weight: float = 0.005
+    utd_ratio: int = 1  # If you want enable utd_ratio, you need to set it to >1
+    state_encoder_hidden_dim: int = 256
+    latent_dim: int = 256
+    target_entropy: float | None = None
+    use_backup_entropy: bool = True
+    grad_clip_norm: float = 40.0
+
+    # Network configuration
+    critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
+    actor_network_kwargs: ActorNetworkConfig = field(default_factory=ActorNetworkConfig)
+    policy_kwargs: PolicyConfig = field(default_factory=PolicyConfig)
+    grasp_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
+    actor_learner_config: ActorLearnerConfig = field(default_factory=ActorLearnerConfig)
+    concurrency: ConcurrencyConfig = field(default_factory=ConcurrencyConfig)
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Any validation specific to SAC configuration
+
+    def get_optimizer_preset(self) -> MultiAdamConfig:
+        return MultiAdamConfig(
+            weight_decay=0.0,
+            optimizer_groups={
+                "actor": {"lr": self.actor_lr},
+                "critic": {"lr": self.critic_lr},
+                "temperature": {"lr": self.temperature_lr},
+            },
+        )
+
+    def get_scheduler_preset(self) -> None:
+        return None
+
+    def validate_features(self) -> None:
+        has_image = any(key.startswith("observation.image") for key in self.input_features)
+        has_state = "observation.state" in self.input_features
+
+        if not (has_state or has_image):
+            raise ValueError(
+                "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
+            )
+
+        if "action" not in self.output_features:
+            raise ValueError("You must provide 'action' in the output features")
+
+    @property
+    def image_features(self) -> list[str]:
+        return [key for key in self.input_features if "image" in key]
+
+    @property
+    def observation_delta_indices(self) -> list:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list:
+        return None  # SAC typically predicts one action at a time
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

lerobot/common/policies/tdmpc/modeling_tdmpc.py

@@ -39,7 +39,11 @@ from lerobot.common.constants import OBS_ENV, OBS_ROBOT
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.pretrained import PreTrainedPolicy
 from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.common.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
+from lerobot.common.policies.utils import (
+    get_device_from_parameters,
+    get_output_shape,
+    populate_queues,
+)
 
 
 class TDMPCPolicy(PreTrainedPolicy):
@@ -63,7 +67,11 @@ class TDMPCPolicy(PreTrainedPolicy):
     config_class = TDMPCConfig
     name = "tdmpc"
 
-    def __init__(self, config: TDMPCConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
+    def __init__(
+        self,
+        config: TDMPCConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
@@ -189,13 +197,20 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         # In the CEM loop we will need this for a call to estimate_value with the gaussian sampled
         # trajectories.
-        z = einops.repeat(z, "b d -> n b d", n=self.config.n_gaussian_samples + self.config.n_pi_samples)
+        z = einops.repeat(
+            z,
+            "b d -> n b d",
+            n=self.config.n_gaussian_samples + self.config.n_pi_samples,
+        )
 
         # Model Predictive Path Integral (MPPI) with the cross-entropy method (CEM) as the optimization
         # algorithm.
         # The initial mean and standard deviation for the cross-entropy method (CEM).
         mean = torch.zeros(
-            self.config.horizon, batch_size, self.config.action_feature.shape[0], device=device
+            self.config.horizon,
+            batch_size,
+            self.config.action_feature.shape[0],
+            device=device,
         )
         # Maybe warm start CEM with the mean from the previous step.
         if self._prev_mean is not None:
@@ -291,9 +306,10 @@ class TDMPCPolicy(PreTrainedPolicy):
         if self.config.q_ensemble_size > 2:
             G += (
                 running_discount
-                * torch.min(terminal_values[torch.randint(0, self.config.q_ensemble_size, size=(2,))], dim=0)[
-                    0
-                ]
+                * torch.min(
+                    terminal_values[torch.randint(0, self.config.q_ensemble_size, size=(2,))],
+                    dim=0,
+                )[0]
             )
         else:
             G += running_discount * torch.min(terminal_values, dim=0)[0]
@@ -329,7 +345,10 @@ class TDMPCPolicy(PreTrainedPolicy):
         # Apply random image augmentations.
         if self.config.image_features and self.config.max_random_shift_ratio > 0:
             observations["observation.image"] = flatten_forward_unflatten(
-                partial(random_shifts_aug, max_random_shift_ratio=self.config.max_random_shift_ratio),
+                partial(
+                    random_shifts_aug,
+                    max_random_shift_ratio=self.config.max_random_shift_ratio,
+                ),
                 observations["observation.image"],
             )
 
@@ -553,7 +572,10 @@ class TDMPCTOLD(nn.Module):
         self._Qs = nn.ModuleList(
             [
                 nn.Sequential(
-                    nn.Linear(config.latent_dim + config.action_feature.shape[0], config.mlp_dim),
+                    nn.Linear(
+                        config.latent_dim + config.action_feature.shape[0],
+                        config.mlp_dim,
+                    ),
                     nn.LayerNorm(config.mlp_dim),
                     nn.Tanh(),
                     nn.Linear(config.mlp_dim, config.mlp_dim),
@@ -702,11 +724,26 @@ class TDMPCObservationEncoder(nn.Module):
                     stride=2,
                 ),
                 nn.ReLU(),
-                nn.Conv2d(config.image_encoder_hidden_dim, config.image_encoder_hidden_dim, 5, stride=2),
+                nn.Conv2d(
+                    config.image_encoder_hidden_dim,
+                    config.image_encoder_hidden_dim,
+                    5,
+                    stride=2,
+                ),
                 nn.ReLU(),
-                nn.Conv2d(config.image_encoder_hidden_dim, config.image_encoder_hidden_dim, 3, stride=2),
+                nn.Conv2d(
+                    config.image_encoder_hidden_dim,
+                    config.image_encoder_hidden_dim,
+                    3,
+                    stride=2,
+                ),
                 nn.ReLU(),
-                nn.Conv2d(config.image_encoder_hidden_dim, config.image_encoder_hidden_dim, 3, stride=2),
+                nn.Conv2d(
+                    config.image_encoder_hidden_dim,
+                    config.image_encoder_hidden_dim,
+                    3,
+                    stride=2,
+                ),
                 nn.ReLU(),
             )
             dummy_shape = (1, *next(iter(config.image_features.values())).shape)
@@ -749,7 +786,8 @@ class TDMPCObservationEncoder(nn.Module):
         if self.config.image_features:
             feat.append(
                 flatten_forward_unflatten(
-                    self.image_enc_layers, obs_dict[next(iter(self.config.image_features))]
+                    self.image_enc_layers,
+                    obs_dict[next(iter(self.config.image_features))],
                 )
             )
         if self.config.env_state_feature:
@@ -796,7 +834,9 @@ def update_ema_parameters(ema_net: nn.Module, net: nn.Module, alpha: float):
     """Update EMA parameters in place with ema_param <- alpha * ema_param + (1 - alpha) * param."""
     for ema_module, module in zip(ema_net.modules(), net.modules(), strict=True):
         for (n_p_ema, p_ema), (n_p, p) in zip(
-            ema_module.named_parameters(recurse=False), module.named_parameters(recurse=False), strict=True
+            ema_module.named_parameters(recurse=False),
+            module.named_parameters(recurse=False),
+            strict=True,
         ):
             assert n_p_ema == n_p, "Parameter names don't match for EMA model update"
             if isinstance(p, dict):

lerobot/common/policies/vqbet/configuration_vqbet.py

@@ -193,7 +193,12 @@ class VQBeTConfig(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        return list(range(1 - self.n_obs_steps, self.n_action_pred_token + self.action_chunk_size - 1))
+        return list(
+            range(
+                1 - self.n_obs_steps,
+                self.n_action_pred_token + self.action_chunk_size - 1,
+            )
+        )
 
     @property
     def reward_delta_indices(self) -> None:

lerobot/common/policies/vqbet/modeling_vqbet.py

@@ -29,7 +29,11 @@ from torch import Tensor, nn
 
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
+from lerobot.common.policies.utils import (
+    get_device_from_parameters,
+    get_output_shape,
+    populate_queues,
+)
 from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.common.policies.vqbet.vqbet_utils import GPT, ResidualVQ
 
@@ -324,7 +328,8 @@ class VQBeTModel(nn.Module):
 
         # To input state and observation features into GPT layers, we first project the features to fit the shape of input size of GPT.
         self.state_projector = MLP(
-            config.robot_state_feature.shape[0], hidden_channels=[self.config.gpt_input_dim]
+            config.robot_state_feature.shape[0],
+            hidden_channels=[self.config.gpt_input_dim],
         )
         self.rgb_feature_projector = MLP(
             self.rgb_encoder.feature_dim, hidden_channels=[self.config.gpt_input_dim]
@@ -354,7 +359,11 @@ class VQBeTModel(nn.Module):
         )
         # Separate batch and sequence dims.
         img_features = einops.rearrange(
-            img_features, "(b s n) ... -> b s n ...", b=batch_size, s=n_obs_steps, n=self.num_images
+            img_features,
+            "(b s n) ... -> b s n ...",
+            b=batch_size,
+            s=n_obs_steps,
+            n=self.num_images,
         )
 
         # Arrange prior and current observation step tokens as shown in the class docstring.
@@ -391,7 +400,11 @@ class VQBeTModel(nn.Module):
         # Thus, it predicts a historical action sequence, in addition to current and future actions (predicting future actions : optional).
         if len_additional_action_token > 0:
             features = torch.cat(
-                [features[:, historical_act_pred_index], features[:, -len_additional_action_token:]], dim=1
+                [
+                    features[:, historical_act_pred_index],
+                    features[:, -len_additional_action_token:],
+                ],
+                dim=1,
             )
         else:
             features = features[:, historical_act_pred_index]
@@ -514,7 +527,13 @@ class VQBeTHead(nn.Module):
 
             cbet_secondary_logits = self.map_to_cbet_preds_secondary_bin(
                 torch.cat(
-                    (x, F.one_hot(sampled_primary_centers, num_classes=self.config.vqvae_n_embed)),
+                    (
+                        x,
+                        F.one_hot(
+                            sampled_primary_centers,
+                            num_classes=self.config.vqvae_n_embed,
+                        ),
+                    ),
                     axis=1,
                 )
             )
@@ -532,7 +551,9 @@ class VQBeTHead(nn.Module):
         else:
             cbet_logits = self.map_to_cbet_preds_bin(x)
             cbet_logits = einops.rearrange(
-                cbet_logits, "(NT) (G C) -> (NT) G C", G=self.vqvae_model.vqvae_num_layers
+                cbet_logits,
+                "(NT) (G C) -> (NT) G C",
+                G=self.vqvae_model.vqvae_num_layers,
             )
             cbet_probs = torch.softmax(cbet_logits / self.config.bet_softmax_temperature, dim=-1)
             NT, G, choices = cbet_probs.shape
@@ -730,7 +751,9 @@ class VQBeTRgbEncoder(nn.Module):
 
 
 def _replace_submodules(
-    root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]
+    root_module: nn.Module,
+    predicate: Callable[[nn.Module], bool],
+    func: Callable[[nn.Module], nn.Module],
 ) -> nn.Module:
     """
     Args:

lerobot/common/policies/vqbet/vqbet_utils.py

@@ -377,7 +377,10 @@ class ResidualVQ(nn.Module):
         self.layers = nn.ModuleList(
             [
                 VectorQuantize(
-                    dim=codebook_dim, codebook_dim=codebook_dim, accept_image_fmap=accept_image_fmap, **kwargs
+                    dim=codebook_dim,
+                    codebook_dim=codebook_dim,
+                    accept_image_fmap=accept_image_fmap,
+                    **kwargs,
                 )
                 for _ in range(num_quantizers)
             ]

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -297,7 +297,11 @@ class IntelRealSenseCamera:
         if self.fps and self.capture_width and self.capture_height:
             # TODO(rcadene): can we set rgb8 directly?
             config.enable_stream(
-                rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
+                rs.stream.color,
+                self.capture_width,
+                self.capture_height,
+                rs.format.rgb8,
+                self.fps,
             )
         else:
             config.enable_stream(rs.stream.color)
@@ -305,7 +309,11 @@ class IntelRealSenseCamera:
         if self.use_depth:
             if self.fps and self.capture_width and self.capture_height:
                 config.enable_stream(
-                    rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
+                    rs.stream.depth,
+                    self.capture_width,
+                    self.capture_height,
+                    rs.format.z16,
+                    self.fps,
                 )
             else:
                 config.enable_stream(rs.stream.depth)

lerobot/common/robot_devices/cameras/utils.py

@@ -41,7 +41,9 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> list[C
             cameras[key] = OpenCVCamera(cfg)
 
         elif cfg.type == "intelrealsense":
-            from lerobot.common.robot_devices.cameras.intelrealsense import IntelRealSenseCamera
+            from lerobot.common.robot_devices.cameras.intelrealsense import (
+                IntelRealSenseCamera,
+            )
 
             cameras[key] = IntelRealSenseCamera(cfg)
         else:
@@ -58,7 +60,9 @@ def make_camera(camera_type, **kwargs) -> Camera:
         return OpenCVCamera(config)
 
     elif camera_type == "intelrealsense":
-        from lerobot.common.robot_devices.cameras.intelrealsense import IntelRealSenseCamera
+        from lerobot.common.robot_devices.cameras.intelrealsense import (
+            IntelRealSenseCamera,
+        )
 
         config = IntelRealSenseCameraConfig(**kwargs)
         return IntelRealSenseCamera(config)

lerobot/common/robot_devices/control_configs.py

@@ -87,6 +87,8 @@ class RecordControlConfig(ControlConfig):
     play_sounds: bool = True
     # Resume recording on an existing dataset.
     resume: bool = False
+    # Reset follower arms to an initial configuration.
+    reset_follower_arms: bool = True
 
     def __post_init__(self):
         # HACK: We parse again the cli args here to get the pretrained path if there was one.

lerobot/common/robot_devices/control_utils.py

@@ -25,6 +25,7 @@ from copy import copy
 from functools import cache
 
 import rerun as rr
+import numpy as np
 import torch
 from deepdiff import DeepDiff
 from termcolor import colored
@@ -128,14 +129,22 @@ def predict_action(observation, policy, device, use_amp):
     return action
 
 
-def init_keyboard_listener():
-    # Allow to exit early while recording an episode or resetting the environment,
-    # by tapping the right arrow key '->'. This might require a sudo permission
-    # to allow your terminal to monitor keyboard events.
+def init_keyboard_listener(assign_rewards=False):
+    """
+    Initializes a keyboard listener to enable early termination of an episode
+    or environment reset by pressing the right arrow key ('->'). This may require
+    sudo permissions to allow the terminal to monitor keyboard events.
+
+    Args:
+        assign_rewards (bool): If True, allows annotating the collected trajectory
+        with a binary reward at the end of the episode to indicate success.
+    """
     events = {}
     events["exit_early"] = False
     events["rerecord_episode"] = False
     events["stop_recording"] = False
+    if assign_rewards:
+        events["next.reward"] = 0
 
     if is_headless():
         logging.warning(
@@ -160,6 +169,13 @@ def init_keyboard_listener():
                 print("Escape key pressed. Stopping data recording...")
                 events["stop_recording"] = True
                 events["exit_early"] = True
+            elif assign_rewards and key == keyboard.Key.space:
+                events["next.reward"] = 1 if events["next.reward"] == 0 else 0
+                print(
+                    "Space key pressed. Assigning new reward to the subsequent frames. New reward:",
+                    events["next.reward"],
+                )
+
         except Exception as e:
             print(f"Error handling key press: {e}")
 
@@ -246,6 +262,8 @@ def control_loop(
     while timestamp < control_time_s:
         start_loop_t = time.perf_counter()
 
+        current_joint_positions = robot.follower_arms["main"].read("Present_Position")
+
         if teleoperate:
             observation, action = robot.teleop_step(record_data=True)
         else:
@@ -253,7 +271,10 @@ def control_loop(
 
             if policy is not None:
                 pred_action = predict_action(
-                    observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+                    observation,
+                    policy,
+                    get_safe_torch_device(policy.config.device),
+                    policy.config.use_amp,
                 )
                 # Action can eventually be clipped using `max_relative_target`,
                 # so action actually sent is saved in the dataset.
@@ -301,7 +322,17 @@ def reset_environment(robot, events, reset_time_s, fps):
     )
 
 
-def stop_recording(robot, listener, display_data):
+def reset_follower_position(robot: Robot, target_position):
+    current_position = robot.follower_arms["main"].read("Present_Position")
+    trajectory = torch.from_numpy(
+        np.linspace(current_position, target_position, 50)
+    )  # NOTE: 30 is just an aribtrary number
+    for pose in trajectory:
+        robot.send_action(pose)
+        busy_wait(0.015)
+
+
+def stop_recording(robot, listener, display_cameras):
     robot.disconnect()
 
     if not is_headless() and listener is not None:
@@ -327,12 +358,20 @@ def sanity_check_dataset_name(repo_id, policy_cfg):
 
 
 def sanity_check_dataset_robot_compatibility(
-    dataset: LeRobotDataset, robot: Robot, fps: int, use_videos: bool
+    dataset: LeRobotDataset,
+    robot: Robot,
+    fps: int,
+    use_videos: bool,
+    extra_features: dict = None,
 ) -> None:
+    features_from_robot = get_features_from_robot(robot, use_videos)
+    if extra_features is not None:
+        features_from_robot.update(extra_features)
+
     fields = [
         ("robot_type", dataset.meta.robot_type, robot.robot_type),
         ("fps", dataset.fps, fps),
-        ("features", dataset.features, get_features_from_robot(robot, use_videos)),
+        ("features", dataset.features, features_from_robot),
     ]
 
     mismatches = []

lerobot/common/robot_devices/motors/dynamixel.py

@@ -23,7 +23,10 @@ import numpy as np
 import tqdm
 
 from lerobot.common.robot_devices.motors.configs import DynamixelMotorsBusConfig
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
 from lerobot.common.utils.utils import capture_timestamp_utc
 
 PROTOCOL_VERSION = 2.0
@@ -784,7 +787,12 @@ class DynamixelMotorsBus:
                 f"{self.packet_handler.getTxRxResult(comm)}"
             )
 
-    def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
+    def write(
+        self,
+        data_name,
+        values: int | float | np.ndarray,
+        motor_names: str | list[str] | None = None,
+    ):
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
                 f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."

lerobot/common/robot_devices/motors/feetech.py

@@ -23,7 +23,10 @@ import numpy as np
 import tqdm
 
 from lerobot.common.robot_devices.motors.configs import FeetechMotorsBusConfig
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
 from lerobot.common.utils.utils import capture_timestamp_utc
 
 PROTOCOL_VERSION = 0
@@ -809,7 +812,12 @@ class FeetechMotorsBus:
                 f"{self.packet_handler.getTxRxResult(comm)}"
             )
 
-    def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
+    def write(
+        self,
+        data_name,
+        values: int | float | np.ndarray,
+        motor_names: str | list[str] | None = None,
+    ):
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
                 f"FeetechMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."

lerobot/common/robot_devices/motors/utils.py

@@ -30,7 +30,9 @@ class MotorsBus(Protocol):
     def write(self): ...
 
 
-def make_motors_buses_from_configs(motors_bus_configs: dict[str, MotorsBusConfig]) -> list[MotorsBus]:
+def make_motors_buses_from_configs(
+    motors_bus_configs: dict[str, MotorsBusConfig],
+) -> list[MotorsBus]:
     motors_buses = {}
 
     for key, cfg in motors_bus_configs.items():

lerobot/common/robot_devices/robots/configs.py

@@ -443,7 +443,7 @@ class So100RobotConfig(ManipulatorRobotConfig):
     leader_arms: dict[str, MotorsBusConfig] = field(
         default_factory=lambda: {
             "main": FeetechMotorsBusConfig(
-                port="/dev/tty.usbmodem58760431091",
+                port="/dev/tty.usbmodem58760433331",
                 motors={
                     # name: (index, model)
                     "shoulder_pan": [1, "sts3215"],
@@ -460,7 +460,7 @@ class So100RobotConfig(ManipulatorRobotConfig):
     follower_arms: dict[str, MotorsBusConfig] = field(
         default_factory=lambda: {
             "main": FeetechMotorsBusConfig(
-                port="/dev/tty.usbmodem585A0076891",
+                port="/dev/tty.usbmodem58760431631",
                 motors={
                     # name: (index, model)
                     "shoulder_pan": [1, "sts3215"],

lerobot/common/robot_devices/robots/feetech_calibration.py

@@ -207,7 +207,10 @@ def run_arm_auto_calibration_so100(arm: MotorsBus, robot_type: str, arm_name: st
 
     print("Calibrate elbow_flex")
     calib["elbow_flex"] = move_to_calibrate(
-        arm, "elbow_flex", positive_first=False, in_between_move_hook=in_between_move_hook
+        arm,
+        "elbow_flex",
+        positive_first=False,
+        in_between_move_hook=in_between_move_hook,
     )
     calib["elbow_flex"] = apply_offset(calib["elbow_flex"], offset=80 - 1024)
 
@@ -239,7 +242,11 @@ def run_arm_auto_calibration_so100(arm: MotorsBus, robot_type: str, arm_name: st
         }
         arm.write("Goal_Position", list(positions.values()), list(positions.keys()))
 
-    arm.write("Goal_Position", round(calib["shoulder_lift"]["zero_pos"] - 1600), "shoulder_lift")
+    arm.write(
+        "Goal_Position",
+        round(calib["shoulder_lift"]["zero_pos"] - 1600),
+        "shoulder_lift",
+    )
     time.sleep(2)
     arm.write("Goal_Position", round(calib["elbow_flex"]["zero_pos"] + 1700), "elbow_flex")
     time.sleep(2)
@@ -250,7 +257,11 @@ def run_arm_auto_calibration_so100(arm: MotorsBus, robot_type: str, arm_name: st
 
     print("Calibrate wrist_roll")
     calib["wrist_roll"] = move_to_calibrate(
-        arm, "wrist_roll", invert_drive_mode=True, positive_first=False, while_move_hook=while_move_hook
+        arm,
+        "wrist_roll",
+        invert_drive_mode=True,
+        positive_first=False,
+        while_move_hook=while_move_hook,
     )
 
     arm.write("Goal_Position", calib["wrist_roll"]["zero_pos"], "wrist_roll")

lerobot/common/robot_devices/robots/lekiwi_remote.py

@@ -61,7 +61,9 @@ def calibrate_follower_arm(motors_bus, calib_dir_str):
     calib_dir.mkdir(parents=True, exist_ok=True)
     calib_file = calib_dir / "main_follower.json"
     try:
-        from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+        from lerobot.common.robot_devices.robots.feetech_calibration import (
+            run_arm_manual_calibration,
+        )
     except ImportError:
         print("[WARNING] Calibration function not available. Skipping calibration.")
         return
@@ -116,7 +118,14 @@ def run_lekiwi(robot_config):
     robot = LeKiwi(motors_bus)
 
     # Define the expected arm motor IDs.
-    arm_motor_ids = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
+    arm_motor_ids = [
+        "shoulder_pan",
+        "shoulder_lift",
+        "elbow_flex",
+        "wrist_flex",
+        "wrist_roll",
+        "gripper",
+    ]
 
     # Disable torque for each arm motor.
     for motor in arm_motor_ids:
@@ -130,7 +139,9 @@ def run_lekiwi(robot_config):
     images_lock = threading.Lock()
     stop_event = threading.Event()
     cam_thread = threading.Thread(
-        target=run_camera_capture, args=(cameras, images_lock, latest_images_dict, stop_event), daemon=True
+        target=run_camera_capture,
+        args=(cameras, images_lock, latest_images_dict, stop_event),
+        daemon=True,
     )
     cam_thread.start()
 

lerobot/common/robot_devices/robots/manipulator.py

@@ -28,14 +28,22 @@ import numpy as np
 import torch
 
 from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
-from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
+from lerobot.common.robot_devices.motors.utils import (
+    MotorsBus,
+    make_motors_buses_from_configs,
+)
 from lerobot.common.robot_devices.robots.configs import ManipulatorRobotConfig
 from lerobot.common.robot_devices.robots.utils import get_arm_id
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
 
 
 def ensure_safe_goal_position(
-    goal_pos: torch.Tensor, present_pos: torch.Tensor, max_relative_target: float | list[float]
+    goal_pos: torch.Tensor,
+    present_pos: torch.Tensor,
+    max_relative_target: float | list[float],
 ):
     # Cap relative action target magnitude for safety.
     diff = goal_pos - present_pos
@@ -45,7 +53,7 @@ def ensure_safe_goal_position(
     safe_goal_pos = present_pos + safe_diff
 
     if not torch.allclose(goal_pos, safe_goal_pos):
-        logging.warning(
+        logging.debug(
             "Relative goal position magnitude had to be clamped to be safe.\n"
             f"  requested relative goal position target: {diff}\n"
             f"    clamped relative goal position target: {safe_diff}"
@@ -309,7 +317,9 @@ class ManipulatorRobot:
                 print(f"Missing calibration file '{arm_calib_path}'")
 
                 if self.robot_type in ["koch", "koch_bimanual", "aloha"]:
-                    from lerobot.common.robot_devices.robots.dynamixel_calibration import run_arm_calibration
+                    from lerobot.common.robot_devices.robots.dynamixel_calibration import (
+                        run_arm_calibration,
+                    )
 
                     calibration = run_arm_calibration(arm, self.robot_type, name, arm_type)
 
@@ -464,6 +474,14 @@ class ManipulatorRobot:
             before_fwrite_t = time.perf_counter()
             goal_pos = leader_pos[name]
 
+            # If specified, clip the goal positions within predefined bounds specified in the config of the robot
+            # if self.config.joint_position_relative_bounds is not None:
+            #     goal_pos = torch.clamp(
+            #         goal_pos,
+            #         self.config.joint_position_relative_bounds["min"],
+            #         self.config.joint_position_relative_bounds["max"],
+            #     )
+
             # Cap goal position when too far away from present position.
             # Slower fps expected due to reading from the follower.
             if self.config.max_relative_target is not None:
@@ -585,6 +603,14 @@ class ManipulatorRobot:
             goal_pos = action[from_idx:to_idx]
             from_idx = to_idx
 
+            # If specified, clip the goal positions within predefined bounds specified in the config of the robot
+            # if self.config.joint_position_relative_bounds is not None:
+            #     goal_pos = torch.clamp(
+            #         goal_pos,
+            #         self.config.joint_position_relative_bounds["min"],
+            #         self.config.joint_position_relative_bounds["max"],
+            #     )
+
             # Cap goal position when too far away from present position.
             # Slower fps expected due to reading from the follower.
             if self.config.max_relative_target is not None:

lerobot/common/robot_devices/robots/mobile_manipulator.py

@@ -25,9 +25,14 @@ import zmq
 
 from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
 from lerobot.common.robot_devices.motors.feetech import TorqueMode
-from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
+from lerobot.common.robot_devices.motors.utils import (
+    MotorsBus,
+    make_motors_buses_from_configs,
+)
 from lerobot.common.robot_devices.robots.configs import LeKiwiRobotConfig
-from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+from lerobot.common.robot_devices.robots.feetech_calibration import (
+    run_arm_manual_calibration,
+)
 from lerobot.common.robot_devices.robots.utils import get_arm_id
 from lerobot.common.robot_devices.utils import RobotDeviceNotConnectedError
 
@@ -324,7 +329,11 @@ class MobileManipulator:
         socks = dict(poller.poll(15))
         if self.video_socket not in socks or socks[self.video_socket] != zmq.POLLIN:
             # No new data arrived → reuse ALL old data
-            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+            return (
+                self.last_frames,
+                self.last_present_speed,
+                self.last_remote_arm_state,
+            )
 
         # Drain all messages, keep only the last
         last_msg = None
@@ -337,7 +346,11 @@ class MobileManipulator:
 
         if not last_msg:
             # No new message → also reuse old
-            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+            return (
+                self.last_frames,
+                self.last_present_speed,
+                self.last_remote_arm_state,
+            )
 
         # Decode only the final message
         try:
@@ -375,7 +388,11 @@ class MobileManipulator:
         except Exception as e:
             print(f"[DEBUG] Error decoding video message: {e}")
             # If decode fails, fall back to old data
-            return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+            return (
+                self.last_frames,
+                self.last_present_speed,
+                self.last_remote_arm_state,
+            )
 
         return frames, present_speed, remote_arm_state_tensor
 
@@ -461,7 +478,11 @@ class MobileManipulator:
 
         body_state = self.wheel_raw_to_body(present_speed)
 
-        body_state_mm = (body_state[0] * 1000.0, body_state[1] * 1000.0, body_state[2])  # Convert x,y to mm/s
+        body_state_mm = (
+            body_state[0] * 1000.0,
+            body_state[1] * 1000.0,
+            body_state[2],
+        )  # Convert x,y to mm/s
         wheel_state_tensor = torch.tensor(body_state_mm, dtype=torch.float32)
         combined_state_tensor = torch.cat((remote_arm_state_tensor, wheel_state_tensor), dim=0)
 
@@ -620,7 +641,11 @@ class MobileManipulator:
         # Convert each wheel’s angular speed (deg/s) to a raw integer.
         wheel_raw = [MobileManipulator.degps_to_raw(deg) for deg in wheel_degps]
 
-        return {"left_wheel": wheel_raw[0], "back_wheel": wheel_raw[1], "right_wheel": wheel_raw[2]}
+        return {
+            "left_wheel": wheel_raw[0],
+            "back_wheel": wheel_raw[1],
+            "right_wheel": wheel_raw[2],
+        }
 
     def wheel_raw_to_body(
         self, wheel_raw: dict, wheel_radius: float = 0.05, base_radius: float = 0.125

lerobot/common/robot_devices/robots/utils.py

@@ -72,7 +72,9 @@ def make_robot_from_config(config: RobotConfig):
 
         return ManipulatorRobot(config)
     elif isinstance(config, LeKiwiRobotConfig):
-        from lerobot.common.robot_devices.robots.mobile_manipulator import MobileManipulator
+        from lerobot.common.robot_devices.robots.mobile_manipulator import (
+            MobileManipulator,
+        )
 
         return MobileManipulator(config)
     else:

lerobot/common/robot_devices/utils.py

@@ -48,7 +48,8 @@ class RobotDeviceNotConnectedError(Exception):
     """Exception raised when the robot device is not connected."""
 
     def __init__(
-        self, message="This robot device is not connected. Try calling `robot_device.connect()` first."
+        self,
+        message="This robot device is not connected. Try calling `robot_device.connect()` first.",
     ):
         self.message = message
         super().__init__(self.message)

lerobot/common/utils/io_utils.py

@@ -28,7 +28,9 @@ def write_video(video_path, stacked_frames, fps):
     # Filter out DeprecationWarnings raised from pkg_resources
     with warnings.catch_warnings():
         warnings.filterwarnings(
-            "ignore", "pkg_resources is deprecated as an API", category=DeprecationWarning
+            "ignore",
+            "pkg_resources is deprecated as an API",
+            category=DeprecationWarning,
         )
         imageio.mimsave(video_path, stacked_frames, fps=fps)
 

lerobot/common/utils/random_utils.py

@@ -42,7 +42,11 @@ def deserialize_python_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> Non
     """
     Restores the rng state for `random` from a dictionary produced by `serialize_python_rng_state()`.
     """
-    py_state = (rng_state_dict["py_rng_version"].item(), tuple(rng_state_dict["py_rng_state"].tolist()), None)
+    py_state = (
+        rng_state_dict["py_rng_version"].item(),
+        tuple(rng_state_dict["py_rng_state"].tolist()),
+        None,
+    )
     random.setstate(py_state)
 
 

lerobot/common/utils/utils.py

@@ -18,6 +18,7 @@ import os
 import os.path as osp
 import platform
 import subprocess
+import time
 from copy import copy
 from datetime import datetime, timezone
 from pathlib import Path
@@ -107,11 +108,11 @@ def is_amp_available(device: str):
         raise ValueError(f"Unknown device '{device}.")
 
 
-def init_logging():
+def init_logging(log_file=None):
     def custom_format(record):
         dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         fnameline = f"{record.pathname}:{record.lineno}"
-        message = f"{record.levelname} {dt} {fnameline[-15:]:>15} {record.msg}"
+        message = f"{record.levelname} [PID: {os.getpid()}] {dt} {fnameline[-15:]:>15} {record.msg}"
         return message
 
     logging.basicConfig(level=logging.INFO)
@@ -125,6 +126,12 @@ def init_logging():
     console_handler.setFormatter(formatter)
     logging.getLogger().addHandler(console_handler)
 
+    if log_file is not None:
+        # File handler
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setFormatter(formatter)
+        logging.getLogger().addHandler(file_handler)
+
 
 def format_big_number(num, precision=0):
     suffixes = ["", "K", "M", "B", "T", "Q"]
@@ -228,3 +235,33 @@ def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
     except TypeError:
         # If a TypeError is raised, the string is not a valid dtype
         return False
+
+
+class TimerManager:
+    def __init__(
+        self,
+        elapsed_time_list: list[float] | None = None,
+        label="Elapsed time",
+        log=True,
+    ):
+        self.label = label
+        self.elapsed_time_list = elapsed_time_list
+        self.log = log
+        self.elapsed = 0.0
+
+    def __enter__(self):
+        self.start = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.elapsed: float = time.perf_counter() - self.start
+
+        if self.elapsed_time_list is not None:
+            self.elapsed_time_list.append(self.elapsed)
+
+        if self.log:
+            print(f"{self.label}: {self.elapsed:.6f} seconds")
+
+    @property
+    def elapsed_seconds(self):
+        return self.elapsed

lerobot/common/utils/wandb_utils.py

@@ -30,9 +30,10 @@ def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[st
     """Return a group name for logging. Optionally returns group name as list."""
     lst = [
         f"policy:{cfg.policy.type}",
-        f"dataset:{cfg.dataset.repo_id}",
         f"seed:{cfg.seed}",
     ]
+    if cfg.dataset is not None:
+        lst.append(f"dataset:{cfg.dataset.repo_id}")
     if cfg.env is not None:
         lst.append(f"env:{cfg.env.type}")
     return lst if return_list else "-".join(lst)
@@ -92,6 +93,12 @@ class WandBLogger:
             resume="must" if cfg.resume else None,
             mode=self.cfg.mode if self.cfg.mode in ["online", "offline", "disabled"] else "online",
         )
+        run_id = wandb.run.id
+        # NOTE: We will override the cfg.wandb.run_id with the wandb run id.
+        # This is because we want to be able to resume the run from the wandb run id.
+        cfg.wandb.run_id = run_id
+        # Handle custom step key for rl asynchronous training.
+        self._wandb_custom_step_key: set[str] | None = None
         print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
         logging.info(f"Track this run --> {colored(wandb.run.get_url(), 'yellow', attrs=['bold'])}")
         self._wandb = wandb
@@ -108,9 +115,26 @@ class WandBLogger:
         artifact.add_file(checkpoint_dir / PRETRAINED_MODEL_DIR / SAFETENSORS_SINGLE_FILE)
         self._wandb.log_artifact(artifact)
 
-    def log_dict(self, d: dict, step: int, mode: str = "train"):
+    def log_dict(
+        self, d: dict, step: int | None = None, mode: str = "train", custom_step_key: str | None = None
+    ):
         if mode not in {"train", "eval"}:
             raise ValueError(mode)
+        if step is None and custom_step_key is None:
+            raise ValueError("Either step or custom_step_key must be provided.")
+
+        # NOTE: This is not simple. Wandb step is it must always monotonically increase and it
+        # increases with each wandb.log call, but in the case of asynchronous RL for example,
+        # multiple time steps is possible for example, the interaction step with the environment,
+        # the training step, the evaluation step, etc. So we need to define a custom step key
+        # to log the correct step for each metric.
+        if custom_step_key is not None:
+            if self._wandb_custom_step_key is None:
+                self._wandb_custom_step_key = set()
+            new_custom_key = f"{mode}/{custom_step_key}"
+            if new_custom_key not in self._wandb_custom_step_key:
+                self._wandb_custom_step_key.add(new_custom_key)
+                self._wandb.define_metric(new_custom_key, hidden=True)
 
         for k, v in d.items():
             if not isinstance(v, (int, float, str)):
@@ -118,7 +142,22 @@ class WandBLogger:
                     f'WandB logging of key "{k}" was ignored as its type is not handled by this wrapper.'
                 )
                 continue
-            self._wandb.log({f"{mode}/{k}": v}, step=step)
+
+            # Do not log the custom step key itself.
+            if self._wandb_custom_step_key is not None and k in self._wandb_custom_step_key:
+                continue
+
+            if custom_step_key is not None:
+                value_custom_step = d[custom_step_key]
+                self._wandb.log(
+                    {
+                        f"{mode}/{k}": v,
+                        f"{mode}/{custom_step_key}": value_custom_step,
+                    }
+                )
+                continue
+
+            self._wandb.log(data={f"{mode}/{k}": v}, step=step)
 
     def log_video(self, video_path: str, step: int, mode: str = "train"):
         if mode not in {"train", "eval"}:

lerobot/configs/parser.py

@@ -42,7 +42,10 @@ def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> lis
         args = sys.argv[1:]
     attr_level_args = []
     detect_string = f"--{field_name}."
-    exclude_strings = (f"--{field_name}.{draccus.CHOICE_TYPE_KEY}=", f"--{field_name}.{PATH_KEY}=")
+    exclude_strings = (
+        f"--{field_name}.{draccus.CHOICE_TYPE_KEY}=",
+        f"--{field_name}.{PATH_KEY}=",
+    )
     for arg in args:
         if arg.startswith(detect_string) and not arg.startswith(exclude_strings):
             denested_arg = f"--{arg.removeprefix(detect_string)}"

lerobot/configs/policies.py

@@ -26,7 +26,11 @@ from huggingface_hub.errors import HfHubHTTPError
 from lerobot.common.optim.optimizers import OptimizerConfig
 from lerobot.common.optim.schedulers import LRSchedulerConfig
 from lerobot.common.utils.hub import HubMixin
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
+from lerobot.common.utils.utils import (
+    auto_select_torch_device,
+    is_amp_available,
+    is_torch_device_available,
+)
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 
 # Generic variable that is either PreTrainedConfig or a subclass thereof

lerobot/configs/train.py

@@ -34,11 +34,10 @@ TRAIN_CONFIG_NAME = "train_config.json"
 
 @dataclass
 class TrainPipelineConfig(HubMixin):
-    dataset: DatasetConfig
+    dataset: DatasetConfig | None = None  # NOTE: In RL, we don't need a dataset
     env: envs.EnvConfig | None = None
     policy: PreTrainedConfig | None = None
-    # Set `dir` to where you would like to save all of the run outputs. If you run another training session
-    # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
+    # Set `dir` to where you would like to save all of the run outputs. If you run another training session # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
     output_dir: Path | None = None
     job_name: str | None = None
     # Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure
@@ -107,8 +106,9 @@ class TrainPipelineConfig(HubMixin):
             train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
             self.output_dir = Path("outputs/train") / train_dir
 
-        if isinstance(self.dataset.repo_id, list):
-            raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
+        if self.dataset is not None:
+            if isinstance(self.dataset.repo_id, list):
+                raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
 
         if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
             raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
@@ -125,7 +125,10 @@ class TrainPipelineConfig(HubMixin):
         return draccus.encode(self)
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / TRAIN_CONFIG_NAME, "w") as f, draccus.config_type("json"):
+        with (
+            open(save_directory / TRAIN_CONFIG_NAME, "w") as f,
+            draccus.config_type("json"),
+        ):
             draccus.dump(self, f, indent=4)
 
     @classmethod

lerobot/configs/types.py

@@ -23,6 +23,7 @@ class FeatureType(str, Enum):
     VISUAL = "VISUAL"
     ENV = "ENV"
     ACTION = "ACTION"
+    REWARD = "REWARD"
 
 
 class NormalizationMode(str, Enum):

lerobot/scripts/configure_motor.py

@@ -38,7 +38,12 @@ def get_motor_bus_cls(brand: str) -> tuple:
             FeetechMotorsBus,
         )
 
-        return FeetechMotorsBusConfig, FeetechMotorsBus, MODEL_BAUDRATE_TABLE, SCS_SERIES_BAUDRATE_TABLE
+        return (
+            FeetechMotorsBusConfig,
+            FeetechMotorsBus,
+            MODEL_BAUDRATE_TABLE,
+            SCS_SERIES_BAUDRATE_TABLE,
+        )
 
     elif brand == "dynamixel":
         from lerobot.common.robot_devices.motors.configs import DynamixelMotorsBusConfig
@@ -48,7 +53,12 @@ def get_motor_bus_cls(brand: str) -> tuple:
             DynamixelMotorsBus,
         )
 
-        return DynamixelMotorsBusConfig, DynamixelMotorsBus, MODEL_BAUDRATE_TABLE, X_SERIES_BAUDRATE_TABLE
+        return (
+            DynamixelMotorsBusConfig,
+            DynamixelMotorsBus,
+            MODEL_BAUDRATE_TABLE,
+            X_SERIES_BAUDRATE_TABLE,
+        )
 
     else:
         raise ValueError(
@@ -164,12 +174,25 @@ def configure_motor(port, brand, model, motor_idx_des, baudrate_des):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--port", type=str, required=True, help="Motors bus port (e.g. dynamixel,feetech)")
+    parser.add_argument(
+        "--port",
+        type=str,
+        required=True,
+        help="Motors bus port (e.g. dynamixel,feetech)",
+    )
     parser.add_argument("--brand", type=str, required=True, help="Motor brand (e.g. dynamixel,feetech)")
     parser.add_argument("--model", type=str, required=True, help="Motor model (e.g. xl330-m077,sts3215)")
-    parser.add_argument("--ID", type=int, required=True, help="Desired ID of the current motor (e.g. 1,2,3)")
     parser.add_argument(
-        "--baudrate", type=int, default=1000000, help="Desired baudrate for the motor (default: 1000000)"
+        "--ID",
+        type=int,
+        required=True,
+        help="Desired ID of the current motor (e.g. 1,2,3)",
+    )
+    parser.add_argument(
+        "--baudrate",
+        type=int,
+        default=1000000,
+        help="Desired baudrate for the motor (default: 1000000)",
     )
     args = parser.parse_args()
 

lerobot/scripts/control_robot.py

@@ -274,10 +274,12 @@ def record(
     # Load pretrained policy
     policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
 
+    # Load pretrained policy
+    policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
+
     if not robot.is_connected:
         robot.connect()
-
-    listener, events = init_keyboard_listener()
+    listener, events = init_keyboard_listener(assign_rewards=cfg.assign_rewards)
 
     # Execute a few seconds without recording to:
     # 1. teleoperate the robot to move it in starting position if no policy provided,
@@ -285,7 +287,14 @@ def record(
     # 3. place the cameras windows on screen
     enable_teleoperation = policy is None
     log_say("Warmup record", cfg.play_sounds)
-    warmup_record(robot, events, enable_teleoperation, cfg.warmup_time_s, cfg.display_data, cfg.fps)
+    warmup_record(
+        robot,
+        events,
+        enable_teleoperation,
+        cfg.warmup_time_s,
+        cfg.display_data,
+        cfg.fps,
+    )
 
     if has_method(robot, "teleop_safety_stop"):
         robot.teleop_safety_stop()
@@ -350,12 +359,12 @@ def replay(
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root, episodes=[cfg.episode])
     actions = dataset.hf_dataset.select_columns("action")
-
     if not robot.is_connected:
         robot.connect()
 
     log_say("Replaying episode", cfg.play_sounds, blocking=True)
     for idx in range(dataset.num_frames):
+        current_joint_positions = robot.follower_arms["main"].read("Present_Position")
         start_episode_t = time.perf_counter()
 
         action = actions[idx]["action"]

lerobot/scripts/control_sim_robot.py

@@ -149,7 +149,11 @@ def init_sim_calibration(robot, cfg):
     axis_directions = np.array(cfg.get("axis_directions", [1]))
     offsets = np.array(cfg.get("offsets", [0])) * np.pi
 
-    return {"start_pos": start_pos, "axis_directions": axis_directions, "offsets": offsets}
+    return {
+        "start_pos": start_pos,
+        "axis_directions": axis_directions,
+        "offsets": offsets,
+    }
 
 
 def real_positions_to_sim(real_positions, axis_directions, start_pos, offsets):
@@ -197,8 +201,14 @@ def record(
     resume: bool = False,
     local_files_only: bool = False,
     run_compute_stats: bool = True,
+    assign_rewards: bool = False,
 ) -> LeRobotDataset:
     # Load pretrained policy
+
+    extra_features = (
+        {"next.reward": {"dtype": "int64", "shape": (1,), "names": None}} if assign_rewards else None
+    )
+
     policy = None
     if pretrained_policy_name_or_path is not None:
         policy, policy_fps, device, use_amp = init_policy(pretrained_policy_name_or_path, policy_overrides)
@@ -211,7 +221,7 @@ def record(
         raise ValueError("Either policy or process_action_fn has to be set to enable control in sim.")
 
     # initialize listener before sim env
-    listener, events = init_keyboard_listener()
+    listener, events = init_keyboard_listener(assign_rewards=assign_rewards)
 
     # create sim env
     env = env()
@@ -241,7 +251,11 @@ def record(
             shape = env.observation_space[key].shape
             if not key.startswith("observation.image."):
                 key = "observation.image." + key
-            features[key] = {"dtype": "video", "names": ["channels", "height", "width"], "shape": shape}
+            features[key] = {
+                "dtype": "video",
+                "names": ["channels", "height", "width"],
+                "shape": shape,
+            }
 
         for key, obs_key in state_keys_dict.items():
             features[key] = {
@@ -250,7 +264,12 @@ def record(
                 "shape": env.observation_space[obs_key].shape,
             }
 
-        features["action"] = {"dtype": "float32", "shape": env.action_space.shape, "names": None}
+        features["action"] = {
+            "dtype": "float32",
+            "shape": env.action_space.shape,
+            "names": None,
+        }
+        features = {**features, **extra_features}
 
         # Create empty dataset or load existing saved episodes
         sanity_check_dataset_name(repo_id, policy)
@@ -302,6 +321,13 @@ def record(
                 "timestamp": env_timestamp,
             }
 
+            # Overwrite environment reward with manually assigned reward
+            if assign_rewards:
+                frame["next.reward"] = events["next.reward"]
+
+                # Should success always be false to match what we do in control_utils?
+                frame["next.success"] = False
+
             for key in image_keys:
                 if not key.startswith("observation.image"):
                     frame["observation.image." + key] = observation[key]
@@ -361,7 +387,12 @@ def record(
 
 
 def replay(
-    env, root: Path, repo_id: str, episode: int, fps: int | None = None, local_files_only: bool = True
+    env,
+    root: Path,
+    repo_id: str,
+    episode: int,
+    fps: int | None = None,
+    local_files_only: bool = True,
 ):
     env = env()
 
@@ -408,7 +439,10 @@ if __name__ == "__main__":
 
     parser_record = subparsers.add_parser("record", parents=[base_parser])
     parser_record.add_argument(
-        "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
+        "--fps",
+        type=none_or_int,
+        default=None,
+        help="Frames per second (set to None to disable)",
     )
     parser_record.add_argument(
         "--root",
@@ -486,9 +520,19 @@ if __name__ == "__main__":
         default=0,
         help="Resume recording on an existing dataset.",
     )
+    parser_record.add_argument(
+        "--assign-rewards",
+        type=int,
+        default=0,
+        help="Enables the assignation of rewards to frames (by default no assignation). When enabled, assign a 0 reward to frames until the space bar is pressed which assign a 1 reward. Press the space bar a second time to assign a 0 reward. The reward assigned is reset to 0 when the episode ends.",
+    )
+
     parser_replay = subparsers.add_parser("replay", parents=[base_parser])
     parser_replay.add_argument(
-        "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
+        "--fps",
+        type=none_or_int,
+        default=None,
+        help="Frames per second (set to None to disable)",
     )
     parser_replay.add_argument(
         "--root",

lerobot/scripts/eval.py

@@ -293,7 +293,8 @@ def eval_policy(
             seeds = None
         else:
             seeds = range(
-                start_seed + (batch_ix * env.num_envs), start_seed + ((batch_ix + 1) * env.num_envs)
+                start_seed + (batch_ix * env.num_envs),
+                start_seed + ((batch_ix + 1) * env.num_envs),
             )
         rollout_data = rollout(
             env,
@@ -413,7 +414,11 @@ def eval_policy(
 
 
 def _compile_episode_data(
-    rollout_data: dict, done_indices: Tensor, start_episode_index: int, start_data_index: int, fps: float
+    rollout_data: dict,
+    done_indices: Tensor,
+    start_episode_index: int,
+    start_data_index: int,
+    fps: float,
 ) -> dict:
     """Convenience function for `eval_policy(return_episode_data=True)`
 
@@ -481,7 +486,10 @@ def eval_main(cfg: EvalPipelineConfig):
     )
     policy.eval()
 
-    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
+    with (
+        torch.no_grad(),
+        torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+    ):
         info = eval_policy(
             env,
             policy,

lerobot/scripts/eval_on_robot.py

@@ -0,0 +1,412 @@
+#!/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.
+"""Evaluate a policy by running rollouts on the real robot and computing metrics.
+
+Usage examples: evaluate a checkpoint from the LeRobot training script for 10 episodes.
+
+```
+python lerobot/scripts/eval_on_robot.py \
+    -p outputs/train/model/checkpoints/005000/pretrained_model \
+    eval.n_episodes=10
+```
+
+Test reward classifier with teleoperation (you need to press space to take over)
+```
+python lerobot/scripts/eval_on_robot.py \
+    --robot-path lerobot/configs/robot/so100.yaml \
+    --reward-classifier-pretrained-path outputs/classifier/checkpoints/best/pretrained_model \
+    --reward-classifier-config-file lerobot/configs/policy/hilserl_classifier.yaml \
+    --display-cameras 1
+```
+
+**NOTE** (michel-aractingi): This script is incomplete and it is being prepared
+for running training on the real robot.
+"""
+
+import argparse
+import logging
+import time
+
+import cv2
+import numpy as np
+import torch
+from tqdm import trange
+
+from lerobot.common.policies.policy_protocol import Policy
+from lerobot.common.robot_devices.control_utils import (
+    busy_wait,
+    is_headless,
+    reset_follower_position,
+)
+from lerobot.common.robot_devices.robots.factory import Robot, make_robot
+from lerobot.common.utils.utils import (
+    init_hydra_config,
+    init_logging,
+    log_say,
+)
+
+
+def get_classifier(pretrained_path, config_path):
+    if pretrained_path is None or config_path is None:
+        return
+
+    from lerobot.common.policies.factory import _policy_cfg_from_hydra_cfg
+    from lerobot.common.policies.hilserl.classifier.configuration_classifier import (
+        ClassifierConfig,
+    )
+    from lerobot.common.policies.hilserl.classifier.modeling_classifier import (
+        Classifier,
+    )
+
+    cfg = init_hydra_config(config_path)
+
+    classifier_config = _policy_cfg_from_hydra_cfg(ClassifierConfig, cfg)
+    classifier_config.num_cameras = len(cfg.training.image_keys)  # TODO automate these paths
+    model = Classifier(classifier_config)
+    model.load_state_dict(Classifier.from_pretrained(pretrained_path).state_dict())
+    model = model.to("mps")
+    return model
+
+
+def rollout(
+    robot: Robot,
+    policy: Policy,
+    reward_classifier,
+    fps: int,
+    control_time_s: float = 20,
+    use_amp: bool = True,
+    display_cameras: bool = False,
+) -> dict:
+    """Run a batched policy rollout on the real robot.
+
+    The return dictionary contains:
+        "robot": A a dictionary of (batch, sequence + 1, *) tensors mapped to observation
+            keys. NOTE the that this has an extra sequence element relative to the other keys in the
+            dictionary. This is because an extra observation is included for after the environment is
+            terminated or truncated.
+        "action": A (batch, sequence, action_dim) tensor of actions applied based on the observations (not
+            including the last observations).
+        "reward": A (batch, sequence) tensor of rewards received for applying the actions.
+        "success": A (batch, sequence) tensor of success conditions (the only time this can be True is upon
+            environment termination/truncation).
+        "done": A (batch, sequence) tensor of **cumulative** done conditions. For any given batch element,
+            the first True is followed by True's all the way till the end. This can be used for masking
+            extraneous elements from the sequences above.
+
+    Args:
+        robot: The robot class that defines the interface with the real robot.
+        policy: The policy. Must be a PyTorch nn module.
+
+    Returns:
+        The dictionary described above.
+    """
+    # TODO (michel-aractingi): Infer the device from policy parameters when policy is added
+    # assert isinstance(policy, nn.Module), "Policy must be a PyTorch nn module."
+    # device = get_device_from_parameters(policy)
+
+    # define keyboard listener
+    listener, events = init_keyboard_listener()
+
+    # Reset the policy. TODO (michel-aractingi) add real policy evaluation once the code is ready.
+    # policy.reset()
+
+    # NOTE: sorting to make sure the key sequence is the same during training and testing.
+    observation = robot.capture_observation()
+    image_keys = [key for key in observation if "image" in key]
+    image_keys.sort()
+
+    all_actions = []
+    all_rewards = []
+    all_successes = []
+
+    start_episode_t = time.perf_counter()
+    init_pos = robot.follower_arms["main"].read("Present_Position")
+    timestamp = 0.0
+    while timestamp < control_time_s:
+        start_loop_t = time.perf_counter()
+
+        # Apply the next action.
+        while events["pause_policy"] and not events["human_intervention_step"]:
+            busy_wait(0.5)
+
+        if events["human_intervention_step"]:
+            # take over the robot's actions
+            observation, action = robot.teleop_step(record_data=True)
+            action = action["action"]  # teleop step returns torch tensors but in a dict
+        else:
+            # explore with policy
+            with torch.inference_mode():
+                # TODO (michel-aractingi) replace this part with policy (predict_action)
+                action = robot.follower_arms["main"].read("Present_Position")
+                action = torch.from_numpy(action)
+                robot.send_action(action)
+                # action = predict_action(observation, policy, device, use_amp)
+
+        observation = robot.capture_observation()
+        images = []
+        for key in image_keys:
+            if display_cameras:
+                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+                cv2.waitKey(1)
+            images.append(observation[key].to("mps"))
+
+        reward = reward_classifier.predict_reward(images) if reward_classifier is not None else 0.0
+        all_rewards.append(reward)
+
+        # print("REWARD : ", reward)
+
+        all_actions.append(action)
+        all_successes.append(torch.tensor([False]))
+
+        dt_s = time.perf_counter() - start_loop_t
+        busy_wait(1 / fps - dt_s)
+        timestamp = time.perf_counter() - start_episode_t
+        if events["exit_early"]:
+            events["exit_early"] = False
+            events["human_intervention_step"] = False
+            events["pause_policy"] = False
+            break
+
+    reset_follower_position(robot, target_position=init_pos)
+
+    dones = torch.tensor([False] * len(all_actions))
+    dones[-1] = True
+    # Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors.
+    ret = {
+        "action": torch.stack(all_actions, dim=1),
+        "next.reward": torch.stack(all_rewards, dim=1),
+        "next.success": torch.stack(all_successes, dim=1),
+        "done": dones,
+    }
+
+    listener.stop()
+
+    return ret
+
+
+def eval_policy(
+    robot: Robot,
+    policy: torch.nn.Module,
+    fps: float,
+    n_episodes: int,
+    control_time_s: int = 20,
+    use_amp: bool = True,
+    display_cameras: bool = False,
+    reward_classifier_pretrained_path: str | None = None,
+    reward_classifier_config_file: str | None = None,
+) -> dict:
+    """
+    Args:
+        env: The batch of environments.
+        policy: The policy.
+        n_episodes: The number of episodes to evaluate.
+    Returns:
+        Dictionary with metrics and data regarding the rollouts.
+    """
+    # TODO (michel-aractingi) comment this out for testing with a fixed policy
+    # assert isinstance(policy, Policy)
+    # policy.eval()
+
+    sum_rewards = []
+    max_rewards = []
+    successes = []
+    rollouts = []
+
+    start_eval = time.perf_counter()
+    progbar = trange(n_episodes, desc="Evaluating policy on real robot")
+    reward_classifier = get_classifier(reward_classifier_pretrained_path, reward_classifier_config_file)
+
+    for _ in progbar:
+        rollout_data = rollout(
+            robot,
+            policy,
+            reward_classifier,
+            fps,
+            control_time_s,
+            use_amp,
+            display_cameras,
+        )
+
+        rollouts.append(rollout_data)
+        sum_rewards.append(sum(rollout_data["next.reward"]))
+        max_rewards.append(max(rollout_data["next.reward"]))
+        successes.append(rollout_data["next.success"][-1])
+
+    info = {
+        "per_episode": [
+            {
+                "episode_ix": i,
+                "sum_reward": sum_reward,
+                "max_reward": max_reward,
+                "pc_success": success * 100,
+            }
+            for i, (sum_reward, max_reward, success) in enumerate(
+                zip(
+                    sum_rewards[:n_episodes],
+                    max_rewards[:n_episodes],
+                    successes[:n_episodes],
+                    strict=False,
+                )
+            )
+        ],
+        "aggregated": {
+            "avg_sum_reward": float(np.nanmean(torch.cat(sum_rewards[:n_episodes]))),
+            "avg_max_reward": float(np.nanmean(torch.cat(max_rewards[:n_episodes]))),
+            "pc_success": float(np.nanmean(torch.cat(successes[:n_episodes])) * 100),
+            "eval_s": time.time() - start_eval,
+            "eval_ep_s": (time.time() - start_eval) / n_episodes,
+        },
+    }
+
+    if robot.is_connected:
+        robot.disconnect()
+
+    return info
+
+
+def init_keyboard_listener():
+    # Allow to exit early while recording an episode or resetting the environment,
+    # by tapping the right arrow key '->'. This might require a sudo permission
+    # to allow your terminal to monitor keyboard events.
+    events = {}
+    events["exit_early"] = False
+    events["rerecord_episode"] = False
+    events["pause_policy"] = False
+    events["human_intervention_step"] = False
+
+    if is_headless():
+        logging.warning(
+            "Headless environment detected. On-screen cameras display and keyboard inputs will not be available."
+        )
+        listener = None
+        return listener, events
+
+    # Only import pynput if not in a headless environment
+    from pynput import keyboard
+
+    def on_press(key):
+        try:
+            if key == keyboard.Key.right:
+                print("Right arrow key pressed. Exiting loop...")
+                events["exit_early"] = True
+            elif key == keyboard.Key.left:
+                print("Left arrow key pressed. Exiting loop and rerecord the last episode...")
+                events["rerecord_episode"] = True
+                events["exit_early"] = True
+            elif key == keyboard.Key.space:
+                # check if first space press then pause the policy for the user to get ready
+                # if second space press then the user is ready to start intervention
+                if not events["pause_policy"]:
+                    print(
+                        "Space key pressed. Human intervention required.\n"
+                        "Place the leader in similar pose to the follower and press space again."
+                    )
+                    events["pause_policy"] = True
+                    log_say(
+                        "Human intervention stage. Get ready to take over.",
+                        play_sounds=True,
+                    )
+                else:
+                    events["human_intervention_step"] = True
+                    print("Space key pressed. Human intervention starting.")
+                    log_say("Starting human intervention.", play_sounds=True)
+
+        except Exception as e:
+            print(f"Error handling key press: {e}")
+
+    listener = keyboard.Listener(on_press=on_press)
+    listener.start()
+
+    return listener, events
+
+
+if __name__ == "__main__":
+    init_logging()
+
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--robot-path",
+        type=str,
+        default="lerobot/configs/robot/koch.yaml",
+        help="Path to robot yaml file used to instantiate the robot using `make_robot` factory function.",
+    )
+    group.add_argument(
+        "--robot-overrides",
+        type=str,
+        nargs="*",
+        help="Any key=value arguments to override config values (use dots for.nested=overrides)",
+    )
+    group.add_argument(
+        "-p",
+        "--pretrained-policy-name-or-path",
+        help=(
+            "Either the repo ID of a model hosted on the Hub or a path to a directory containing weights "
+            "saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch "
+            "(useful for debugging). This argument is mutually exclusive with `--config`."
+        ),
+    )
+    group.add_argument(
+        "--config",
+        help=(
+            "Path to a yaml config you want to use for initializing a policy from scratch (useful for "
+            "debugging). This argument is mutually exclusive with `--pretrained-policy-name-or-path` (`-p`)."
+        ),
+    )
+    parser.add_argument("--revision", help="Optionally provide the Hugging Face Hub revision ID.")
+    parser.add_argument(
+        "--out-dir",
+        help=(
+            "Where to save the evaluation outputs. If not provided, outputs are saved in "
+            "outputs/eval/{timestamp}_{env_name}_{policy_name}"
+        ),
+    )
+    parser.add_argument(
+        "--display-cameras",
+        help=("Whether to display the camera feed while the rollout is happening"),
+    )
+    parser.add_argument(
+        "--reward-classifier-pretrained-path",
+        type=str,
+        default=None,
+        help="Path to the pretrained classifier weights.",
+    )
+    parser.add_argument(
+        "--reward-classifier-config-file",
+        type=str,
+        default=None,
+        help="Path to a yaml config file that is necessary to build the reward classifier model.",
+    )
+
+    args = parser.parse_args()
+
+    robot_cfg = init_hydra_config(args.robot_path, args.robot_overrides)
+    robot = make_robot(robot_cfg)
+    if not robot.is_connected:
+        robot.connect()
+
+    eval_policy(
+        robot,
+        None,
+        fps=40,
+        n_episodes=2,
+        control_time_s=100,
+        display_cameras=args.display_cameras,
+        reward_classifier_config_file=args.reward_classifier_config_file,
+        reward_classifier_pretrained_path=args.reward_classifier_pretrained_path,
+    )

lerobot/scripts/server/actor_server.py

@@ -0,0 +1,676 @@
+#!/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.
+import logging
+import os
+import time
+from functools import lru_cache
+from queue import Empty
+from statistics import mean, quantiles
+
+# from lerobot.scripts.eval import eval_policy
+import grpc
+import torch
+from torch import nn
+from torch.multiprocessing import Event, Queue
+
+# TODO: Remove the import of maniskill
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies.sac.modeling_sac import SACPolicy
+from lerobot.common.robot_devices.utils import busy_wait
+from lerobot.common.utils.random_utils import set_seed
+from lerobot.common.utils.utils import (
+    TimerManager,
+    get_safe_torch_device,
+    init_logging,
+)
+from lerobot.configs import parser
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.scripts.server import hilserl_pb2, hilserl_pb2_grpc, learner_service
+from lerobot.scripts.server.buffer import (
+    Transition,
+    bytes_to_state_dict,
+    move_state_dict_to_device,
+    move_transition_to_device,
+    python_object_to_bytes,
+    transitions_to_bytes,
+)
+from lerobot.scripts.server.gym_manipulator import make_robot_env
+from lerobot.scripts.server.network_utils import (
+    receive_bytes_in_chunks,
+    send_bytes_in_chunks,
+)
+from lerobot.scripts.server.utils import get_last_item_from_queue, setup_process_handlers
+
+ACTOR_SHUTDOWN_TIMEOUT = 30
+
+
+#################################################
+# Main entry point #
+#################################################
+
+
+@parser.wrap()
+def actor_cli(cfg: TrainPipelineConfig):
+    cfg.validate()
+    if not use_threads(cfg):
+        import torch.multiprocessing as mp
+
+        mp.set_start_method("spawn")
+
+    # Create logs directory to ensure it exists
+    log_dir = os.path.join(cfg.output_dir, "logs")
+    os.makedirs(log_dir, exist_ok=True)
+    log_file = os.path.join(log_dir, f"actor_{cfg.job_name}.log")
+
+    # Initialize logging with explicit log file
+    init_logging(log_file=log_file)
+    logging.info(f"Actor logging initialized, writing to {log_file}")
+
+    shutdown_event = setup_process_handlers(use_threads(cfg))
+
+    learner_client, grpc_channel = learner_service_client(
+        host=cfg.policy.actor_learner_config.learner_host,
+        port=cfg.policy.actor_learner_config.learner_port,
+    )
+
+    logging.info("[ACTOR] Establishing connection with Learner")
+    if not establish_learner_connection(learner_client, shutdown_event):
+        logging.error("[ACTOR] Failed to establish connection with Learner")
+        return
+
+    if not use_threads(cfg):
+        # If we use multithreading, we can reuse the channel
+        grpc_channel.close()
+        grpc_channel = None
+
+    logging.info("[ACTOR] Connection with Learner established")
+
+    parameters_queue = Queue()
+    transitions_queue = Queue()
+    interactions_queue = Queue()
+
+    concurrency_entity = None
+    if use_threads(cfg):
+        from threading import Thread
+
+        concurrency_entity = Thread
+    else:
+        from multiprocessing import Process
+
+        concurrency_entity = Process
+
+    receive_policy_process = concurrency_entity(
+        target=receive_policy,
+        args=(cfg, parameters_queue, shutdown_event, grpc_channel),
+        daemon=True,
+    )
+
+    transitions_process = concurrency_entity(
+        target=send_transitions,
+        args=(cfg, transitions_queue, shutdown_event, grpc_channel),
+        daemon=True,
+    )
+
+    interactions_process = concurrency_entity(
+        target=send_interactions,
+        args=(cfg, interactions_queue, shutdown_event, grpc_channel),
+        daemon=True,
+    )
+
+    transitions_process.start()
+    interactions_process.start()
+    receive_policy_process.start()
+
+    # HACK: FOR MANISKILL we do not have a reward classifier
+    # TODO: Remove this once we merge into main
+    reward_classifier = None
+    # if (
+    #     cfg.env.reward_classifier["pretrained_path"] is not None
+    #     and cfg.env.reward_classifier["config_path"] is not None
+    # ):
+    #     reward_classifier = get_classifier(
+    #         pretrained_path=cfg.env.reward_classifier["pretrained_path"],
+    #         config_path=cfg.env.reward_classifier["config_path"],
+    #     )
+
+    act_with_policy(
+        cfg=cfg,
+        reward_classifier=reward_classifier,
+        shutdown_event=shutdown_event,
+        parameters_queue=parameters_queue,
+        transitions_queue=transitions_queue,
+        interactions_queue=interactions_queue,
+    )
+    logging.info("[ACTOR] Policy process joined")
+
+    logging.info("[ACTOR] Closing queues")
+    transitions_queue.close()
+    interactions_queue.close()
+    parameters_queue.close()
+
+    transitions_process.join()
+    logging.info("[ACTOR] Transitions process joined")
+    interactions_process.join()
+    logging.info("[ACTOR] Interactions process joined")
+    receive_policy_process.join()
+    logging.info("[ACTOR] Receive policy process joined")
+
+    logging.info("[ACTOR] join queues")
+    transitions_queue.cancel_join_thread()
+    interactions_queue.cancel_join_thread()
+    parameters_queue.cancel_join_thread()
+
+    logging.info("[ACTOR] queues closed")
+
+
+#################################################
+# Core algorithm functions #
+#################################################
+
+
+def act_with_policy(
+    cfg: TrainPipelineConfig,
+    reward_classifier: nn.Module,
+    shutdown_event: any,  # Event,
+    parameters_queue: Queue,
+    transitions_queue: Queue,
+    interactions_queue: Queue,
+):
+    """
+    Executes policy interaction within the environment.
+
+    This function rolls out the policy in the environment, collecting interaction data and pushing it to a queue for streaming to the learner.
+    Once an episode is completed, updated network parameters received from the learner are retrieved from a queue and loaded into the network.
+
+    Args:
+        cfg: Configuration settings for the interaction process.
+        reward_classifier: Reward classifier to use for the interaction process.
+        shutdown_event: Event to check if the process should shutdown.
+        parameters_queue: Queue to receive updated network parameters from the learner.
+        transitions_queue: Queue to send transitions to the learner.
+        interactions_queue: Queue to send interactions to the learner.
+    """
+    # Initialize logging for multiprocessing
+    if not use_threads(cfg):
+        log_dir = os.path.join(cfg.output_dir, "logs")
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = os.path.join(log_dir, f"actor_policy_{os.getpid()}.log")
+        init_logging(log_file=log_file)
+        logging.info("Actor policy process logging initialized")
+
+    logging.info("make_env online")
+
+    online_env = make_robot_env(cfg=cfg.env)
+
+    set_seed(cfg.seed)
+    device = get_safe_torch_device(cfg.policy.device, log=True)
+
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging.info("make_policy")
+
+    ### Instantiate the policy in both the actor and learner processes
+    ### To avoid sending a SACPolicy object through the port, we create a policy intance
+    ### on both sides, the learner sends the updated parameters every n steps to update the actor's parameters
+    # TODO: At some point we should just need make sac policy
+    policy: SACPolicy = make_policy(
+        cfg=cfg.policy,
+        env_cfg=cfg.env,
+    )
+    policy = policy.eval()
+    assert isinstance(policy, nn.Module)
+
+    obs, info = online_env.reset()
+
+    # NOTE: For the moment we will solely handle the case of a single environment
+    sum_reward_episode = 0
+    list_transition_to_send_to_learner = []
+    list_policy_time = []
+    episode_intervention = False
+    # Add counters for intervention rate calculation
+    episode_intervention_steps = 0
+    episode_total_steps = 0
+
+    for interaction_step in range(cfg.policy.online_steps):
+        start_time = time.perf_counter()
+        if shutdown_event.is_set():
+            logging.info("[ACTOR] Shutting down act_with_policy")
+            return
+
+        if interaction_step >= cfg.policy.online_step_before_learning:
+            # Time policy inference and check if it meets FPS requirement
+            with TimerManager(
+                elapsed_time_list=list_policy_time,
+                label="Policy inference time",
+                log=False,
+            ) as timer:  # noqa: F841
+                action = policy.select_action(batch=obs)
+            policy_fps = 1.0 / (list_policy_time[-1] + 1e-9)
+
+            log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step)
+
+            next_obs, reward, done, truncated, info = online_env.step(action.squeeze(dim=0).cpu().numpy())
+        else:
+            # TODO (azouitine): Make a custom space for torch tensor
+            action = online_env.action_space.sample()
+            next_obs, reward, done, truncated, info = online_env.step(action)
+
+            # HACK: We have only one env but we want to batch it, it will be resolved with the torch box
+            action = (
+                torch.from_numpy(action[0]).to(device, non_blocking=device.type == "cuda").unsqueeze(dim=0)
+            )
+
+        sum_reward_episode += float(reward)
+        # Increment total steps counter for intervention rate
+        episode_total_steps += 1
+
+        # NOTE: We overide the action if the intervention is True, because the action applied is the intervention action
+        if "is_intervention" in info and info["is_intervention"]:
+            # TODO: Check the shape
+            # NOTE: The action space for demonstration before hand is with the full action space
+            # but sometimes for example we want to deactivate the gripper
+            action = info["action_intervention"]
+            episode_intervention = True
+            # Increment intervention steps counter
+            episode_intervention_steps += 1
+
+        # Check for NaN values in observations
+        for key, tensor in obs.items():
+            if torch.isnan(tensor).any():
+                logging.error(f"[ACTOR] NaN values found in obs[{key}] at step {interaction_step}")
+
+        list_transition_to_send_to_learner.append(
+            Transition(
+                state=obs,
+                action=action,
+                reward=reward,
+                next_state=next_obs,
+                done=done,
+                truncated=truncated,  # TODO: (azouitine) Handle truncation properly
+                complementary_info=info,  # TODO Handle information for the transition, is_demonstraction: bool
+            )
+        )
+        # assign obs to the next obs and continue the rollout
+        obs = next_obs
+
+        # HACK: We have only one env but we want to batch it, it will be resolved with the torch box
+        # Because we are using a single environment we can index at zero
+        if done or truncated:
+            # TODO: Handle logging for episode information
+            logging.info(f"[ACTOR] Global step {interaction_step}: Episode reward: {sum_reward_episode}")
+
+            update_policy_parameters(policy=policy.actor, parameters_queue=parameters_queue, device=device)
+
+            if len(list_transition_to_send_to_learner) > 0:
+                push_transitions_to_transport_queue(
+                    transitions=list_transition_to_send_to_learner,
+                    transitions_queue=transitions_queue,
+                )
+                list_transition_to_send_to_learner = []
+
+            stats = get_frequency_stats(list_policy_time)
+            list_policy_time.clear()
+
+            # Calculate intervention rate
+            intervention_rate = 0.0
+            if episode_total_steps > 0:
+                intervention_rate = episode_intervention_steps / episode_total_steps
+
+            # Send episodic reward to the learner
+            interactions_queue.put(
+                python_object_to_bytes(
+                    {
+                        "Episodic reward": sum_reward_episode,
+                        "Interaction step": interaction_step,
+                        "Episode intervention": int(episode_intervention),
+                        "Intervention rate": intervention_rate,
+                        **stats,
+                    }
+                )
+            )
+            sum_reward_episode = 0.0
+            episode_intervention = False
+            # Reset intervention counters
+            episode_intervention_steps = 0
+            episode_total_steps = 0
+            obs, info = online_env.reset()
+
+        if cfg.env.fps is not None:
+            dt_time = time.perf_counter() - start_time
+            busy_wait(1 / cfg.env.fps - dt_time)
+
+
+#################################################
+#  Communication Functions - Group all gRPC/messaging functions  #
+#################################################
+
+
+def establish_learner_connection(
+    stub,
+    shutdown_event: any,  # Event,
+    attempts=30,
+):
+    for _ in range(attempts):
+        if shutdown_event.is_set():
+            logging.info("[ACTOR] Shutting down establish_learner_connection")
+            return False
+
+        # Force a connection attempt and check state
+        try:
+            logging.info("[ACTOR] Send ready message to Learner")
+            if stub.Ready(hilserl_pb2.Empty()) == hilserl_pb2.Empty():
+                return True
+        except grpc.RpcError as e:
+            logging.error(f"[ACTOR] Waiting for Learner to be ready... {e}")
+            time.sleep(2)
+    return False
+
+
+@lru_cache(maxsize=1)
+def learner_service_client(
+    host="127.0.0.1", port=50051
+) -> tuple[hilserl_pb2_grpc.LearnerServiceStub, grpc.Channel]:
+    import json
+
+    """
+    Returns a client for the learner service.
+
+    GRPC uses HTTP/2, which is a binary protocol and multiplexes requests over a single connection.
+    So we need to create only one client and reuse it.
+    """
+
+    service_config = {
+        "methodConfig": [
+            {
+                "name": [{}],  # Applies to ALL methods in ALL services
+                "retryPolicy": {
+                    "maxAttempts": 5,  # Max retries (total attempts = 5)
+                    "initialBackoff": "0.1s",  # First retry after 0.1s
+                    "maxBackoff": "2s",  # Max wait time between retries
+                    "backoffMultiplier": 2,  # Exponential backoff factor
+                    "retryableStatusCodes": [
+                        "UNAVAILABLE",
+                        "DEADLINE_EXCEEDED",
+                    ],  # Retries on network failures
+                },
+            }
+        ]
+    }
+
+    service_config_json = json.dumps(service_config)
+
+    channel = grpc.insecure_channel(
+        f"{host}:{port}",
+        options=[
+            ("grpc.max_receive_message_length", learner_service.MAX_MESSAGE_SIZE),
+            ("grpc.max_send_message_length", learner_service.MAX_MESSAGE_SIZE),
+            ("grpc.enable_retries", 1),
+            ("grpc.service_config", service_config_json),
+        ],
+    )
+    stub = hilserl_pb2_grpc.LearnerServiceStub(channel)
+    logging.info("[ACTOR] Learner service client created")
+    return stub, channel
+
+
+def receive_policy(
+    cfg: TrainPipelineConfig,
+    parameters_queue: Queue,
+    shutdown_event: any,  # Event,
+    learner_client: hilserl_pb2_grpc.LearnerServiceStub | None = None,
+    grpc_channel: grpc.Channel | None = None,
+):
+    logging.info("[ACTOR] Start receiving parameters from the Learner")
+
+    if not use_threads(cfg):
+        # Create a process-specific log file
+        log_dir = os.path.join(cfg.output_dir, "logs")
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = os.path.join(log_dir, f"actor_receive_policy_{os.getpid()}.log")
+
+        # Initialize logging with explicit log file
+        init_logging(log_file=log_file)
+        logging.info("Actor receive policy process logging initialized")
+
+        # Setup process handlers to handle shutdown signal
+        # But use shutdown event from the main process
+        setup_process_handlers(use_threads=False)
+
+    if grpc_channel is None or learner_client is None:
+        learner_client, grpc_channel = learner_service_client(
+            host=cfg.policy.actor_learner_config.learner_host,
+            port=cfg.policy.actor_learner_config.learner_port,
+        )
+
+    try:
+        iterator = learner_client.StreamParameters(hilserl_pb2.Empty())
+        receive_bytes_in_chunks(
+            iterator,
+            parameters_queue,
+            shutdown_event,
+            log_prefix="[ACTOR] parameters",
+        )
+
+    except grpc.RpcError as e:
+        logging.error(f"[ACTOR] gRPC error: {e}")
+
+    if not use_threads(cfg):
+        grpc_channel.close()
+    logging.info("[ACTOR] Received policy loop stopped")
+
+
+def send_transitions(
+    cfg: TrainPipelineConfig,
+    transitions_queue: Queue,
+    shutdown_event: any,  # Event,
+    learner_client: hilserl_pb2_grpc.LearnerServiceStub | None = None,
+    grpc_channel: grpc.Channel | None = None,
+) -> hilserl_pb2.Empty:
+    """
+    Sends transitions to the learner.
+
+    This function continuously retrieves messages from the queue and processes:
+
+    - **Transition Data:**
+        - A batch of transitions (observation, action, reward, next observation) is collected.
+        - Transitions are moved to the CPU and serialized using PyTorch.
+        - The serialized data is wrapped in a `hilserl_pb2.Transition` message and sent to the learner.
+    """
+
+    if not use_threads(cfg):
+        # Create a process-specific log file
+        log_dir = os.path.join(cfg.output_dir, "logs")
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = os.path.join(log_dir, f"actor_transitions_{os.getpid()}.log")
+
+        # Initialize logging with explicit log file
+        init_logging(log_file=log_file)
+        logging.info("Actor transitions process logging initialized")
+
+        # Setup process handlers to handle shutdown signal
+        # But use shutdown event from the main process
+        setup_process_handlers(False)
+
+    if grpc_channel is None or learner_client is None:
+        learner_client, grpc_channel = learner_service_client(
+            host=cfg.policy.actor_learner_config.learner_host,
+            port=cfg.policy.actor_learner_config.learner_port,
+        )
+
+    try:
+        learner_client.SendTransitions(transitions_stream(shutdown_event, transitions_queue))
+    except grpc.RpcError as e:
+        logging.error(f"[ACTOR] gRPC error: {e}")
+
+    logging.info("[ACTOR] Finished streaming transitions")
+
+    if not use_threads(cfg):
+        grpc_channel.close()
+    logging.info("[ACTOR] Transitions process stopped")
+
+
+def send_interactions(
+    cfg: TrainPipelineConfig,
+    interactions_queue: Queue,
+    shutdown_event: any,  # Event,
+    learner_client: hilserl_pb2_grpc.LearnerServiceStub | None = None,
+    grpc_channel: grpc.Channel | None = None,
+) -> hilserl_pb2.Empty:
+    """
+    Sends interactions to the learner.
+
+    This function continuously retrieves messages from the queue and processes:
+
+    - **Interaction Messages:**
+        - Contains useful statistics about episodic rewards and policy timings.
+        - The message is serialized using `pickle` and sent to the learner.
+    """
+
+    if not use_threads(cfg):
+        # Create a process-specific log file
+        log_dir = os.path.join(cfg.output_dir, "logs")
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = os.path.join(log_dir, f"actor_interactions_{os.getpid()}.log")
+
+        # Initialize logging with explicit log file
+        init_logging(log_file=log_file)
+        logging.info("Actor interactions process logging initialized")
+
+        # Setup process handlers to handle shutdown signal
+        # But use shutdown event from the main process
+        setup_process_handlers(False)
+
+    if grpc_channel is None or learner_client is None:
+        learner_client, grpc_channel = learner_service_client(
+            host=cfg.policy.actor_learner_config.learner_host,
+            port=cfg.policy.actor_learner_config.learner_port,
+        )
+
+    try:
+        learner_client.SendInteractions(interactions_stream(shutdown_event, interactions_queue))
+    except grpc.RpcError as e:
+        logging.error(f"[ACTOR] gRPC error: {e}")
+
+    logging.info("[ACTOR] Finished streaming interactions")
+
+    if not use_threads(cfg):
+        grpc_channel.close()
+    logging.info("[ACTOR] Interactions process stopped")
+
+
+def transitions_stream(shutdown_event: Event, transitions_queue: Queue) -> hilserl_pb2.Empty:
+    while not shutdown_event.is_set():
+        try:
+            message = transitions_queue.get(block=True, timeout=5)
+        except Empty:
+            logging.debug("[ACTOR] Transition queue is empty")
+            continue
+
+        yield from send_bytes_in_chunks(
+            message, hilserl_pb2.Transition, log_prefix="[ACTOR] Send transitions"
+        )
+
+    return hilserl_pb2.Empty()
+
+
+def interactions_stream(
+    shutdown_event: any,  # Event,
+    interactions_queue: Queue,
+) -> hilserl_pb2.Empty:
+    while not shutdown_event.is_set():
+        try:
+            message = interactions_queue.get(block=True, timeout=5)
+        except Empty:
+            logging.debug("[ACTOR] Interaction queue is empty")
+            continue
+
+        yield from send_bytes_in_chunks(
+            message,
+            hilserl_pb2.InteractionMessage,
+            log_prefix="[ACTOR] Send interactions",
+        )
+
+    return hilserl_pb2.Empty()
+
+
+#################################################
+#  Policy functions #
+#################################################
+
+
+def update_policy_parameters(policy: SACPolicy, parameters_queue: Queue, device):
+    if not parameters_queue.empty():
+        logging.info("[ACTOR] Load new parameters from Learner.")
+        bytes_state_dict = get_last_item_from_queue(parameters_queue)
+        state_dict = bytes_to_state_dict(bytes_state_dict)
+        state_dict = move_state_dict_to_device(state_dict, device=device)
+        policy.load_state_dict(state_dict)
+
+
+#################################################
+#  Utilities functions #
+#################################################
+
+
+def push_transitions_to_transport_queue(transitions: list, transitions_queue):
+    """Send transitions to learner in smaller chunks to avoid network issues.
+
+    Args:
+        transitions: List of transitions to send
+        message_queue: Queue to send messages to learner
+        chunk_size: Size of each chunk to send
+    """
+    transition_to_send_to_learner = []
+    for transition in transitions:
+        tr = move_transition_to_device(transition=transition, device="cpu")
+        for key, value in tr["state"].items():
+            if torch.isnan(value).any():
+                logging.warning(f"Found NaN values in transition {key}")
+
+        transition_to_send_to_learner.append(tr)
+
+    transitions_queue.put(transitions_to_bytes(transition_to_send_to_learner))
+
+
+def get_frequency_stats(list_policy_time: list[float]) -> dict[str, float]:
+    stats = {}
+    list_policy_fps = [1.0 / t for t in list_policy_time]
+    if len(list_policy_fps) > 1:
+        policy_fps = mean(list_policy_fps)
+        quantiles_90 = quantiles(list_policy_fps, n=10)[-1]
+        logging.debug(f"[ACTOR] Average policy frame rate: {policy_fps}")
+        logging.debug(f"[ACTOR] Policy frame rate 90th percentile: {quantiles_90}")
+        stats = {
+            "Policy frequency [Hz]": policy_fps,
+            "Policy frequency 90th-p [Hz]": quantiles_90,
+        }
+    return stats
+
+
+def log_policy_frequency_issue(policy_fps: float, cfg: TrainPipelineConfig, interaction_step: int):
+    if policy_fps < cfg.env.fps:
+        logging.warning(
+            f"[ACTOR] Policy FPS {policy_fps:.1f} below required {cfg.env.fps} at step {interaction_step}"
+        )
+
+
+def use_threads(cfg: TrainPipelineConfig) -> bool:
+    return cfg.policy.concurrency.actor == "threads"
+
+
+if __name__ == "__main__":
+    actor_cli()

lerobot/scripts/server/crop_dataset_roi.py

@@ -0,0 +1,282 @@
+import argparse  # noqa: I001
+import json
+from copy import deepcopy
+from typing import Dict, Tuple
+from pathlib import Path
+import cv2
+
+# import torch.nn.functional as F  # noqa: N812
+import torchvision.transforms.functional as F  # type: ignore  # noqa: N812
+from tqdm import tqdm  # type: ignore
+
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+
+def select_rect_roi(img):
+    """
+    Allows the user to draw a rectangular ROI on the image.
+
+    The user must click and drag to draw the rectangle.
+    - While dragging, the rectangle is dynamically drawn.
+    - On mouse button release, the rectangle is fixed.
+    - Press 'c' to confirm the selection.
+    - Press 'r' to reset the selection.
+    - Press ESC to cancel.
+
+    Returns:
+        A tuple (top, left, height, width) representing the rectangular ROI,
+        or None if no valid ROI is selected.
+    """
+    # Create a working copy of the image
+    clone = img.copy()
+    working_img = clone.copy()
+
+    roi = None  # Will store the final ROI as (top, left, height, width)
+    drawing = False
+    ix, iy = -1, -1  # Initial click coordinates
+
+    def mouse_callback(event, x, y, flags, param):
+        nonlocal ix, iy, drawing, roi, working_img
+
+        if event == cv2.EVENT_LBUTTONDOWN:
+            # Start drawing: record starting coordinates
+            drawing = True
+            ix, iy = x, y
+
+        elif event == cv2.EVENT_MOUSEMOVE:
+            if drawing:
+                # Compute the top-left and bottom-right corners regardless of drag direction
+                top = min(iy, y)
+                left = min(ix, x)
+                bottom = max(iy, y)
+                right = max(ix, x)
+                # Show a temporary image with the current rectangle drawn
+                temp = working_img.copy()
+                cv2.rectangle(temp, (left, top), (right, bottom), (0, 255, 0), 2)
+                cv2.imshow("Select ROI", temp)
+
+        elif event == cv2.EVENT_LBUTTONUP:
+            # Finish drawing
+            drawing = False
+            top = min(iy, y)
+            left = min(ix, x)
+            bottom = max(iy, y)
+            right = max(ix, x)
+            height = bottom - top
+            width = right - left
+            roi = (top, left, height, width)  # (top, left, height, width)
+            # Draw the final rectangle on the working image and display it
+            working_img = clone.copy()
+            cv2.rectangle(working_img, (left, top), (right, bottom), (0, 255, 0), 2)
+            cv2.imshow("Select ROI", working_img)
+
+    # Create the window and set the callback
+    cv2.namedWindow("Select ROI")
+    cv2.setMouseCallback("Select ROI", mouse_callback)
+    cv2.imshow("Select ROI", working_img)
+
+    print("Instructions for ROI selection:")
+    print("  - Click and drag to draw a rectangular ROI.")
+    print("  - Press 'c' to confirm the selection.")
+    print("  - Press 'r' to reset and draw again.")
+    print("  - Press ESC to cancel the selection.")
+
+    # Wait until the user confirms with 'c', resets with 'r', or cancels with ESC
+    while True:
+        key = cv2.waitKey(1) & 0xFF
+        # Confirm ROI if one has been drawn
+        if key == ord("c") and roi is not None:
+            break
+        # Reset: clear the ROI and restore the original image
+        elif key == ord("r"):
+            working_img = clone.copy()
+            roi = None
+            cv2.imshow("Select ROI", working_img)
+        # Cancel selection for this image
+        elif key == 27:  # ESC key
+            roi = None
+            break
+
+    cv2.destroyWindow("Select ROI")
+    return roi
+
+
+def select_square_roi_for_images(images: dict) -> dict:
+    """
+    For each image in the provided dictionary, open a window to allow the user
+    to select a rectangular ROI. Returns a dictionary mapping each key to a tuple
+    (top, left, height, width) representing the ROI.
+
+    Parameters:
+        images (dict): Dictionary where keys are identifiers and values are OpenCV images.
+
+    Returns:
+        dict: Mapping of image keys to the selected rectangular ROI.
+    """
+    selected_rois = {}
+
+    for key, img in images.items():
+        if img is None:
+            print(f"Image for key '{key}' is None, skipping.")
+            continue
+
+        print(f"\nSelect rectangular ROI for image with key: '{key}'")
+        roi = select_rect_roi(img)
+
+        if roi is None:
+            print(f"No valid ROI selected for '{key}'.")
+        else:
+            selected_rois[key] = roi
+            print(f"ROI for '{key}': {roi}")
+
+    return selected_rois
+
+
+def get_image_from_lerobot_dataset(dataset: LeRobotDataset):
+    """
+    Find the first row in the dataset and extract the image in order to be used for the crop.
+    """
+    row = dataset[0]
+    image_dict = {}
+    for k in row:
+        if "image" in k:
+            image_dict[k] = deepcopy(row[k])
+    return image_dict
+
+
+def convert_lerobot_dataset_to_cropper_lerobot_dataset(
+    original_dataset: LeRobotDataset,
+    crop_params_dict: Dict[str, Tuple[int, int, int, int]],
+    new_repo_id: str,
+    new_dataset_root: str,
+    resize_size: Tuple[int, int] = (128, 128),
+) -> LeRobotDataset:
+    """
+    Converts an existing LeRobotDataset by iterating over its episodes and frames,
+    applying cropping and resizing to image observations, and saving a new dataset
+    with the transformed data.
+
+    Args:
+        original_dataset (LeRobotDataset): The source dataset.
+        crop_params_dict (Dict[str, Tuple[int, int, int, int]]):
+            A dictionary mapping observation keys to crop parameters (top, left, height, width).
+        new_repo_id (str): Repository id for the new dataset.
+        new_dataset_root (str): The root directory where the new dataset will be written.
+        resize_size (Tuple[int, int], optional): The target size (height, width) after cropping.
+            Defaults to (128, 128).
+
+    Returns:
+        LeRobotDataset: A new LeRobotDataset where the specified image observations have been cropped
+                        and resized.
+    """
+    # 1. Create a new (empty) LeRobotDataset for writing.
+    new_dataset = LeRobotDataset.create(
+        repo_id=new_repo_id,
+        fps=original_dataset.fps,
+        root=new_dataset_root,
+        robot_type=original_dataset.meta.robot_type,
+        features=original_dataset.meta.info["features"],
+        use_videos=len(original_dataset.meta.video_keys) > 0,
+    )
+
+    # Update the metadata for every image key that will be cropped:
+    # (Here we simply set the shape to be the final resize_size.)
+    for key in crop_params_dict:
+        if key in new_dataset.meta.info["features"]:
+            new_dataset.meta.info["features"][key]["shape"] = list(resize_size)
+
+    # 2. Process each episode in the original dataset.
+    episodes_info = original_dataset.meta.episodes
+    # (Sort episodes by episode_index for consistency.)
+
+    episodes_info = sorted(episodes_info, key=lambda x: x["episode_index"])
+    # Use the first task from the episode metadata (or "unknown" if not provided)
+    task = episodes_info[0]["tasks"][0] if episodes_info[0].get("tasks") else "unknown"
+
+    last_episode_index = 0
+    for sample in tqdm(original_dataset):
+        episode_index = sample.pop("episode_index")
+        if episode_index != last_episode_index:
+            new_dataset.save_episode(task, encode_videos=True)
+            last_episode_index = episode_index
+        sample.pop("frame_index")
+        # Make a shallow copy of the sample (the values—e.g. torch tensors—are assumed immutable)
+        new_sample = sample.copy()
+        # Loop over each observation key that should be cropped/resized.
+        for key, params in crop_params_dict.items():
+            if key in new_sample:
+                top, left, height, width = params
+                # Apply crop then resize.
+                cropped = F.crop(new_sample[key], top, left, height, width)
+                resized = F.resize(cropped, resize_size)
+                new_sample[key] = resized
+        # Add the transformed frame to the new dataset.
+        new_dataset.add_frame(new_sample)
+
+    # save last episode
+    new_dataset.save_episode(task, encode_videos=True)
+
+    # Optionally, consolidate the new dataset to compute statistics and update video info.
+    new_dataset.consolidate(run_compute_stats=True, keep_image_files=True)
+
+    new_dataset.push_to_hub(tags=None)
+
+    return new_dataset
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Crop rectangular ROIs from a LeRobot dataset.")
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default="lerobot",
+        help="The repository id of the LeRobot dataset to process.",
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default=None,
+        help="The root directory of the LeRobot dataset.",
+    )
+    parser.add_argument(
+        "--crop-params-path",
+        type=str,
+        default=None,
+        help="The path to the JSON file containing the ROIs.",
+    )
+    args = parser.parse_args()
+
+    local_files_only = args.root is not None
+    dataset = LeRobotDataset(repo_id=args.repo_id, root=args.root, local_files_only=local_files_only)
+
+    images = get_image_from_lerobot_dataset(dataset)
+    images = {k: v.cpu().permute(1, 2, 0).numpy() for k, v in images.items()}
+    images = {k: (v * 255).astype("uint8") for k, v in images.items()}
+
+    if args.crop_params_path is None:
+        rois = select_square_roi_for_images(images)
+    else:
+        with open(args.crop_params_path) as f:
+            rois = json.load(f)
+
+    # Print the selected rectangular ROIs
+    print("\nSelected Rectangular Regions of Interest (top, left, height, width):")
+    for key, roi in rois.items():
+        print(f"{key}: {roi}")
+
+    new_repo_id = args.repo_id + "_cropped_resized"
+    new_dataset_root = Path(str(dataset.root) + "_cropped_resized")
+
+    croped_resized_dataset = convert_lerobot_dataset_to_cropper_lerobot_dataset(
+        original_dataset=dataset,
+        crop_params_dict=rois,
+        new_repo_id=new_repo_id,
+        new_dataset_root=new_dataset_root,
+        resize_size=(128, 128),
+    )
+
+    meta_dir = new_dataset_root / "meta"
+    meta_dir.mkdir(exist_ok=True)
+
+    with open(meta_dir / "crop_params.json", "w") as f:
+        json.dump(rois, f, indent=4)

lerobot/scripts/server/end_effector_control_utils.py

@@ -0,0 +1,781 @@
+import argparse
+import logging
+import sys
+import time
+
+import numpy as np
+import torch
+
+from lerobot.common.robot_devices.utils import busy_wait
+from lerobot.scripts.server.kinematics import RobotKinematics
+
+logging.basicConfig(level=logging.INFO)
+
+
+class InputController:
+    """Base class for input controllers that generate motion deltas."""
+
+    def __init__(self, x_step_size=0.01, y_step_size=0.01, z_step_size=0.01):
+        """
+        Initialize the controller.
+
+        Args:
+            x_step_size: Base movement step size in meters
+            y_step_size: Base movement step size in meters
+            z_step_size: Base movement step size in meters
+        """
+        self.x_step_size = x_step_size
+        self.y_step_size = y_step_size
+        self.z_step_size = z_step_size
+        self.running = True
+        self.episode_end_status = None  # None, "success", or "failure"
+        self.intervention_flag = False
+        self.open_gripper_command = False
+        self.close_gripper_command = False
+
+    def start(self):
+        """Start the controller and initialize resources."""
+        pass
+
+    def stop(self):
+        """Stop the controller and release resources."""
+        pass
+
+    def get_deltas(self):
+        """Get the current movement deltas (dx, dy, dz) in meters."""
+        return 0.0, 0.0, 0.0
+
+    def should_quit(self):
+        """Return True if the user has requested to quit."""
+        return not self.running
+
+    def update(self):
+        """Update controller state - call this once per frame."""
+        pass
+
+    def __enter__(self):
+        """Support for use in 'with' statements."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Ensure resources are released when exiting 'with' block."""
+        self.stop()
+
+    def get_episode_end_status(self):
+        """
+        Get the current episode end status.
+
+        Returns:
+            None if episode should continue, "success" or "failure" otherwise
+        """
+        status = self.episode_end_status
+        self.episode_end_status = None  # Reset after reading
+        return status
+
+    def should_intervene(self):
+        """Return True if intervention flag was set."""
+        return self.intervention_flag
+
+    def gripper_command(self):
+        """Return the current gripper command."""
+        if self.open_gripper_command == self.close_gripper_command:
+            return "no-op"
+        elif self.open_gripper_command:
+            return "open"
+        elif self.close_gripper_command:
+            return "close"
+
+
+class KeyboardController(InputController):
+    """Generate motion deltas from keyboard input."""
+
+    def __init__(self, x_step_size=0.01, y_step_size=0.01, z_step_size=0.01):
+        super().__init__(x_step_size, y_step_size, z_step_size)
+        self.key_states = {
+            "forward_x": False,
+            "backward_x": False,
+            "forward_y": False,
+            "backward_y": False,
+            "forward_z": False,
+            "backward_z": False,
+            "quit": False,
+            "success": False,
+            "failure": False,
+        }
+        self.listener = None
+
+    def start(self):
+        """Start the keyboard listener."""
+        from pynput import keyboard
+
+        def on_press(key):
+            try:
+                if key == keyboard.Key.up:
+                    self.key_states["forward_x"] = True
+                elif key == keyboard.Key.down:
+                    self.key_states["backward_x"] = True
+                elif key == keyboard.Key.left:
+                    self.key_states["forward_y"] = True
+                elif key == keyboard.Key.right:
+                    self.key_states["backward_y"] = True
+                elif key == keyboard.Key.shift:
+                    self.key_states["backward_z"] = True
+                elif key == keyboard.Key.shift_r:
+                    self.key_states["forward_z"] = True
+                elif key == keyboard.Key.esc:
+                    self.key_states["quit"] = True
+                    self.running = False
+                    return False
+                elif key == keyboard.Key.enter:
+                    self.key_states["success"] = True
+                    self.episode_end_status = "success"
+                elif key == keyboard.Key.backspace:
+                    self.key_states["failure"] = True
+                    self.episode_end_status = "failure"
+            except AttributeError:
+                pass
+
+        def on_release(key):
+            try:
+                if key == keyboard.Key.up:
+                    self.key_states["forward_x"] = False
+                elif key == keyboard.Key.down:
+                    self.key_states["backward_x"] = False
+                elif key == keyboard.Key.left:
+                    self.key_states["forward_y"] = False
+                elif key == keyboard.Key.right:
+                    self.key_states["backward_y"] = False
+                elif key == keyboard.Key.shift:
+                    self.key_states["backward_z"] = False
+                elif key == keyboard.Key.shift_r:
+                    self.key_states["forward_z"] = False
+                elif key == keyboard.Key.enter:
+                    self.key_states["success"] = False
+                elif key == keyboard.Key.backspace:
+                    self.key_states["failure"] = False
+            except AttributeError:
+                pass
+
+        self.listener = keyboard.Listener(on_press=on_press, on_release=on_release)
+        self.listener.start()
+
+        print("Keyboard controls:")
+        print("  Arrow keys: Move in X-Y plane")
+        print("  Shift and Shift_R: Move in Z axis")
+        print("  Enter: End episode with SUCCESS")
+        print("  Backspace: End episode with FAILURE")
+        print("  ESC: Exit")
+
+    def stop(self):
+        """Stop the keyboard listener."""
+        if self.listener and self.listener.is_alive():
+            self.listener.stop()
+
+    def get_deltas(self):
+        """Get the current movement deltas from keyboard state."""
+        delta_x = delta_y = delta_z = 0.0
+
+        if self.key_states["forward_x"]:
+            delta_x += self.x_step_size
+        if self.key_states["backward_x"]:
+            delta_x -= self.x_step_size
+        if self.key_states["forward_y"]:
+            delta_y += self.y_step_size
+        if self.key_states["backward_y"]:
+            delta_y -= self.y_step_size
+        if self.key_states["forward_z"]:
+            delta_z += self.z_step_size
+        if self.key_states["backward_z"]:
+            delta_z -= self.z_step_size
+
+        return delta_x, delta_y, delta_z
+
+    def should_quit(self):
+        """Return True if ESC was pressed."""
+        return self.key_states["quit"]
+
+    def should_save(self):
+        """Return True if Enter was pressed (save episode)."""
+        return self.key_states["success"] or self.key_states["failure"]
+
+
+class GamepadController(InputController):
+    """Generate motion deltas from gamepad input."""
+
+    def __init__(self, x_step_size=0.01, y_step_size=0.01, z_step_size=0.01, deadzone=0.1):
+        super().__init__(x_step_size, y_step_size, z_step_size)
+        self.deadzone = deadzone
+        self.joystick = None
+        self.intervention_flag = False
+
+    def start(self):
+        """Initialize pygame and the gamepad."""
+        import pygame
+
+        pygame.init()
+        pygame.joystick.init()
+
+        if pygame.joystick.get_count() == 0:
+            logging.error("No gamepad detected. Please connect a gamepad and try again.")
+            self.running = False
+            return
+
+        self.joystick = pygame.joystick.Joystick(0)
+        self.joystick.init()
+        logging.info(f"Initialized gamepad: {self.joystick.get_name()}")
+
+        print("Gamepad controls:")
+        print("  Left analog stick: Move in X-Y plane")
+        print("  Right analog stick (vertical): Move in Z axis")
+        print("  B/Circle button: Exit")
+        print("  Y/Triangle button: End episode with SUCCESS")
+        print("  A/Cross button: End episode with FAILURE")
+        print("  X/Square button: Rerecord episode")
+
+    def stop(self):
+        """Clean up pygame resources."""
+        import pygame
+
+        if pygame.joystick.get_init():
+            if self.joystick:
+                self.joystick.quit()
+            pygame.joystick.quit()
+        pygame.quit()
+
+    def update(self):
+        """Process pygame events to get fresh gamepad readings."""
+        import pygame
+
+        for event in pygame.event.get():
+            if event.type == pygame.JOYBUTTONDOWN:
+                if event.button == 3:
+                    self.episode_end_status = "success"
+                # A button (1) for failure
+                elif event.button == 1:
+                    self.episode_end_status = "failure"
+                # X button (0) for rerecord
+                elif event.button == 0:
+                    self.episode_end_status = "rerecord_episode"
+
+                # RB button (6) for opening gripper
+                elif event.button == 6:
+                    self.open_gripper_command = True
+
+                # LT button (7) for closing gripper
+                elif event.button == 7:
+                    self.close_gripper_command = True
+
+            # Reset episode status on button release
+            elif event.type == pygame.JOYBUTTONUP:
+                if event.button in [0, 2, 3]:
+                    self.episode_end_status = None
+
+                elif event.button == 6:
+                    self.open_gripper_command = False
+
+                elif event.button == 7:
+                    self.close_gripper_command = False
+
+            # Check for RB button (typically button 5) for intervention flag
+            if self.joystick.get_button(5):
+                self.intervention_flag = True
+            else:
+                self.intervention_flag = False
+
+    def get_deltas(self):
+        """Get the current movement deltas from gamepad state."""
+        import pygame
+
+        try:
+            # Read joystick axes
+            # Left stick X and Y (typically axes 0 and 1)
+            x_input = self.joystick.get_axis(0)  # Left/Right
+            y_input = self.joystick.get_axis(1)  # Up/Down (often inverted)
+
+            # Right stick Y (typically axis 3 or 4)
+            z_input = self.joystick.get_axis(3)  # Up/Down for Z
+
+            # Apply deadzone to avoid drift
+            x_input = 0 if abs(x_input) < self.deadzone else x_input
+            y_input = 0 if abs(y_input) < self.deadzone else y_input
+            z_input = 0 if abs(z_input) < self.deadzone else z_input
+
+            # Calculate deltas (note: may need to invert axes depending on controller)
+            delta_x = -y_input * self.y_step_size  # Forward/backward
+            delta_y = -x_input * self.x_step_size  # Left/right
+            delta_z = -z_input * self.z_step_size  # Up/down
+
+            return delta_x, delta_y, delta_z
+
+        except pygame.error:
+            logging.error("Error reading gamepad. Is it still connected?")
+            return 0.0, 0.0, 0.0
+
+
+class GamepadControllerHID(InputController):
+    """Generate motion deltas from gamepad input using HIDAPI."""
+
+    def __init__(
+        self,
+        x_step_size=0.01,
+        y_step_size=0.01,
+        z_step_size=0.01,
+        deadzone=0.1,
+        vendor_id=0x046D,
+        product_id=0xC219,
+    ):
+        """
+        Initialize the HID gamepad controller.
+
+        Args:
+            step_size: Base movement step size in meters
+            z_scale: Scaling factor for Z-axis movement
+            deadzone: Joystick deadzone to prevent drift
+            vendor_id: USB vendor ID of the gamepad (default: Logitech)
+            product_id: USB product ID of the gamepad (default: RumblePad 2)
+        """
+        super().__init__(x_step_size, y_step_size, z_step_size)
+        self.deadzone = deadzone
+        self.vendor_id = vendor_id
+        self.product_id = product_id
+        self.device = None
+        self.device_info = None
+
+        # Movement values (normalized from -1.0 to 1.0)
+        self.left_x = 0.0
+        self.left_y = 0.0
+        self.right_x = 0.0
+        self.right_y = 0.0
+
+        # Button states
+        self.buttons = {}
+        self.quit_requested = False
+        self.save_requested = False
+
+    def find_device(self):
+        """Look for the gamepad device by vendor and product ID."""
+        import hid
+
+        devices = hid.enumerate()
+        for device in devices:
+            if device["vendor_id"] == self.vendor_id and device["product_id"] == self.product_id:
+                logging.info(f"Found gamepad: {device.get('product_string', 'Unknown')}")
+                return device
+
+        logging.error(
+            f"No gamepad with vendor ID 0x{self.vendor_id:04X} and product ID 0x{self.product_id:04X} found"
+        )
+        return None
+
+    def start(self):
+        """Connect to the gamepad using HIDAPI."""
+        import hid
+
+        self.device_info = self.find_device()
+        if not self.device_info:
+            self.running = False
+            return
+
+        try:
+            logging.info(f"Connecting to gamepad at path: {self.device_info['path']}")
+            self.device = hid.device()
+            self.device.open_path(self.device_info["path"])
+            self.device.set_nonblocking(1)
+
+            manufacturer = self.device.get_manufacturer_string()
+            product = self.device.get_product_string()
+            logging.info(f"Connected to {manufacturer} {product}")
+
+            logging.info("Gamepad controls (HID mode):")
+            logging.info("  Left analog stick: Move in X-Y plane")
+            logging.info("  Right analog stick: Move in Z axis (vertical)")
+            logging.info("  Button 1/B/Circle: Exit")
+            logging.info("  Button 2/A/Cross: End episode with SUCCESS")
+            logging.info("  Button 3/X/Square: End episode with FAILURE")
+
+        except OSError as e:
+            logging.error(f"Error opening gamepad: {e}")
+            logging.error("You might need to run this with sudo/admin privileges on some systems")
+            self.running = False
+
+    def stop(self):
+        """Close the HID device connection."""
+        if self.device:
+            self.device.close()
+            self.device = None
+
+    def update(self):
+        """
+        Read and process the latest gamepad data.
+        Due to an issue with the HIDAPI, we need to read the read the device several times in order to get a stable reading
+        """
+        for _ in range(10):
+            self._update()
+
+    def _update(self):
+        """Read and process the latest gamepad data."""
+        if not self.device or not self.running:
+            return
+
+        try:
+            # Read data from the gamepad
+            data = self.device.read(64)
+            if data:
+                # Interpret gamepad data - this will vary by controller model
+                # These offsets are for the Logitech RumblePad 2
+                if len(data) >= 8:
+                    # Normalize joystick values from 0-255 to -1.0-1.0
+                    self.left_x = (data[1] - 128) / 128.0
+                    self.left_y = (data[2] - 128) / 128.0
+                    self.right_x = (data[3] - 128) / 128.0
+                    self.right_y = (data[4] - 128) / 128.0
+
+                    # Apply deadzone
+                    self.left_x = 0 if abs(self.left_x) < self.deadzone else self.left_x
+                    self.left_y = 0 if abs(self.left_y) < self.deadzone else self.left_y
+                    self.right_x = 0 if abs(self.right_x) < self.deadzone else self.right_x
+                    self.right_y = 0 if abs(self.right_y) < self.deadzone else self.right_y
+
+                    # Parse button states (byte 5 in the Logitech RumblePad 2)
+                    buttons = data[5]
+
+                    # Check if RB is pressed then the intervention flag should be set
+                    self.intervention_flag = data[6] in [2, 6, 10, 14]
+
+                    # Check if RT is pressed
+                    self.open_gripper_command = data[6] in [8, 10, 12]
+
+                    # Check if LT is pressed
+                    self.close_gripper_command = data[6] in [4, 6, 12]
+
+                    # Check if Y/Triangle button (bit 7) is pressed for saving
+                    # Check if X/Square button (bit 5) is pressed for failure
+                    # Check if A/Cross button (bit 4) is pressed for rerecording
+                    if buttons & 1 << 7:
+                        self.episode_end_status = "success"
+                    elif buttons & 1 << 5:
+                        self.episode_end_status = "failure"
+                    elif buttons & 1 << 4:
+                        self.episode_end_status = "rerecord_episode"
+                    else:
+                        self.episode_end_status = None
+
+        except OSError as e:
+            logging.error(f"Error reading from gamepad: {e}")
+
+    def get_deltas(self):
+        """Get the current movement deltas from gamepad state."""
+        # Calculate deltas - invert as needed based on controller orientation
+        delta_x = -self.left_y * self.x_step_size  # Forward/backward
+        delta_y = -self.left_x * self.y_step_size  # Left/right
+        delta_z = -self.right_y * self.z_step_size  # Up/down
+
+        return delta_x, delta_y, delta_z
+
+    def should_quit(self):
+        """Return True if quit button was pressed."""
+        return self.quit_requested
+
+    def should_save(self):
+        """Return True if save button was pressed."""
+        return self.save_requested
+
+
+def test_forward_kinematics(robot, fps=10):
+    logging.info("Testing Forward Kinematics")
+    timestep = time.perf_counter()
+    kinematics = RobotKinematics(robot.robot_type)
+    while time.perf_counter() - timestep < 60.0:
+        loop_start_time = time.perf_counter()
+        robot.teleop_step()
+        obs = robot.capture_observation()
+        joint_positions = obs["observation.state"].cpu().numpy()
+        ee_pos = kinematics.fk_gripper_tip(joint_positions)
+        logging.info(f"EE Position: {ee_pos[:3, 3]}")
+        busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+
+def test_inverse_kinematics(robot, fps=10):
+    logging.info("Testing Inverse Kinematics")
+    timestep = time.perf_counter()
+    while time.perf_counter() - timestep < 60.0:
+        loop_start_time = time.perf_counter()
+        obs = robot.capture_observation()
+        joint_positions = obs["observation.state"].cpu().numpy()
+        ee_pos = RobotKinematics.fk_gripper_tip(joint_positions)
+        desired_ee_pos = ee_pos
+        target_joint_state = RobotKinematics.ik(joint_positions, desired_ee_pos, position_only=True)
+        robot.send_action(torch.from_numpy(target_joint_state))
+        logging.info(f"Target Joint State: {target_joint_state}")
+        busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+
+def teleoperate_inverse_kinematics_with_leader(robot, fps=10):
+    logging.info("Testing Inverse Kinematics")
+    kinematics = RobotKinematics(robot.robot_type)
+    timestep = time.perf_counter()
+    while time.perf_counter() - timestep < 60.0:
+        loop_start_time = time.perf_counter()
+        obs = robot.capture_observation()
+        joint_positions = obs["observation.state"].cpu().numpy()
+        ee_pos = kinematics.fk_gripper_tip(joint_positions)
+
+        leader_joint_positions = robot.leader_arms["main"].read("Present_Position")
+        leader_ee = kinematics.fk_gripper_tip(leader_joint_positions)
+
+        desired_ee_pos = leader_ee
+        target_joint_state = kinematics.ik(joint_positions, desired_ee_pos, position_only=True)
+        robot.send_action(torch.from_numpy(target_joint_state))
+        logging.info(f"Leader EE: {leader_ee[:3, 3]}, Follower EE: {ee_pos[:3, 3]}")
+        busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+
+def teleoperate_delta_inverse_kinematics_with_leader(robot, fps=10):
+    logging.info("Testing Delta End-Effector Control")
+    timestep = time.perf_counter()
+
+    # Initial position capture
+    obs = robot.capture_observation()
+    joint_positions = obs["observation.state"].cpu().numpy()
+
+    kinematics = RobotKinematics(robot.robot_type)
+
+    leader_joint_positions = robot.leader_arms["main"].read("Present_Position")
+    initial_leader_ee = kinematics.fk_gripper_tip(leader_joint_positions)
+
+    desired_ee_pos = np.diag(np.ones(4))
+
+    while time.perf_counter() - timestep < 60.0:
+        loop_start_time = time.perf_counter()
+
+        # Get leader state for teleoperation
+        leader_joint_positions = robot.leader_arms["main"].read("Present_Position")
+        leader_ee = kinematics.fk_gripper_tip(leader_joint_positions)
+
+        # Get current state
+        # obs = robot.capture_observation()
+        # joint_positions = obs["observation.state"].cpu().numpy()
+        joint_positions = robot.follower_arms["main"].read("Present_Position")
+        current_ee_pos = kinematics.fk_gripper_tip(joint_positions)
+
+        # Calculate delta between leader and follower end-effectors
+        # Scaling factor can be adjusted for sensitivity
+        scaling_factor = 1.0
+        ee_delta = (leader_ee - initial_leader_ee) * scaling_factor
+
+        # Apply delta to current position
+        desired_ee_pos[0, 3] = current_ee_pos[0, 3] + ee_delta[0, 3]
+        desired_ee_pos[1, 3] = current_ee_pos[1, 3] + ee_delta[1, 3]
+        desired_ee_pos[2, 3] = current_ee_pos[2, 3] + ee_delta[2, 3]
+
+        if np.any(np.abs(ee_delta[:3, 3]) > 0.01):
+            # Compute joint targets via inverse kinematics
+            target_joint_state = kinematics.ik(joint_positions, desired_ee_pos, position_only=True)
+
+            initial_leader_ee = leader_ee.copy()
+
+            # Send command to robot
+            robot.send_action(torch.from_numpy(target_joint_state))
+
+            # Logging
+            logging.info(f"Current EE: {current_ee_pos[:3, 3]}, Desired EE: {desired_ee_pos[:3, 3]}")
+            logging.info(f"Delta EE: {ee_delta[:3, 3]}")
+
+        busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+
+def teleoperate_delta_inverse_kinematics(robot, controller, fps=10, bounds=None, fk_func=None):
+    """
+    Control a robot using delta end-effector movements from any input controller.
+
+    Args:
+        robot: Robot instance to control
+        controller: InputController instance (keyboard, gamepad, etc.)
+        fps: Control frequency in Hz
+        bounds: Optional position limits
+        fk_func: Forward kinematics function to use
+    """
+    if fk_func is None:
+        fk_func = RobotKinematics.fk_gripper_tip
+
+    logging.info(f"Testing Delta End-Effector Control with {controller.__class__.__name__}")
+
+    # Initial position capture
+    obs = robot.capture_observation()
+    joint_positions = obs["observation.state"].cpu().numpy()
+    kinematics = RobotKinematics(robot.robot_type)
+    current_ee_pos = kinematics.fk_gripper_tip(joint_positions)
+
+    # Initialize desired position with current position
+    desired_ee_pos = np.eye(4)  # Identity matrix
+
+    timestep = time.perf_counter()
+    with controller:
+        while not controller.should_quit() and time.perf_counter() - timestep < 60.0:
+            loop_start_time = time.perf_counter()
+
+            # Process input events
+            controller.update()
+
+            # Get currrent robot state
+            joint_positions = robot.follower_arms["main"].read("Present_Position")
+            current_ee_pos = kinematics.fk_gripper_tip(joint_positions)
+
+            # Get movement deltas from the controller
+            delta_x, delta_y, delta_z = controller.get_deltas()
+
+            # Update desired position
+            desired_ee_pos[0, 3] = current_ee_pos[0, 3] + delta_x
+            desired_ee_pos[1, 3] = current_ee_pos[1, 3] + delta_y
+            desired_ee_pos[2, 3] = current_ee_pos[2, 3] + delta_z
+
+            # Apply bounds if provided
+            if bounds is not None:
+                desired_ee_pos[:3, 3] = np.clip(desired_ee_pos[:3, 3], bounds["min"], bounds["max"])
+
+            # Only send commands if there's actual movement
+            if any([abs(v) > 0.001 for v in [delta_x, delta_y, delta_z]]):
+                # Compute joint targets via inverse kinematics
+                target_joint_state = kinematics.ik(joint_positions, desired_ee_pos, position_only=True)
+
+                # Send command to robot
+                robot.send_action(torch.from_numpy(target_joint_state))
+
+            busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+
+def teleoperate_gym_env(env, controller, fps: int = 30):
+    """
+    Control a robot through a gym environment using keyboard inputs.
+
+    Args:
+        env: A gym environment created with make_robot_env
+        fps: Target control frequency
+    """
+
+    logging.info("Testing Keyboard Control of Gym Environment")
+    print("Keyboard controls:")
+    print("  Arrow keys: Move in X-Y plane")
+    print("  Shift and Shift_R: Move in Z axis")
+    print("  ESC: Exit")
+
+    # Reset the environment to get initial observation
+    obs, info = env.reset()
+
+    try:
+        with controller:
+            while not controller.should_quit():
+                loop_start_time = time.perf_counter()
+
+                # Process input events
+                controller.update()
+
+                # Get movement deltas from the controller
+                delta_x, delta_y, delta_z = controller.get_deltas()
+
+                # Create the action vector
+                action = np.array([delta_x, delta_y, delta_z])
+
+                # Skip if no movement
+                if any([abs(v) > 0.001 for v in [delta_x, delta_y, delta_z]]):
+                    # Step the environment - pass action as a tensor with intervention flag
+                    action_tensor = torch.from_numpy(action.astype(np.float32))
+                    obs, reward, terminated, truncated, info = env.step((action_tensor, False))
+
+                    # Log information
+                    logging.info(f"Action: [{delta_x:.4f}, {delta_y:.4f}, {delta_z:.4f}]")
+                    logging.info(f"Reward: {reward}")
+
+                    # Reset if episode ended
+                    if terminated or truncated:
+                        logging.info("Episode ended, resetting environment")
+                        obs, info = env.reset()
+
+                # Maintain target frame rate
+                busy_wait(1 / fps - (time.perf_counter() - loop_start_time))
+
+    finally:
+        # Close the environment
+        env.close()
+
+
+if __name__ == "__main__":
+    from lerobot.common.envs.configs import EEActionSpaceConfig, EnvWrapperConfig, HILSerlRobotEnvConfig
+    from lerobot.common.robot_devices.robots.configs import RobotConfig
+    from lerobot.common.robot_devices.robots.utils import make_robot_from_config
+    from lerobot.scripts.server.gym_manipulator import make_robot_env
+
+    parser = argparse.ArgumentParser(description="Test end-effector control")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="keyboard",
+        choices=[
+            "keyboard",
+            "gamepad",
+            "keyboard_gym",
+            "gamepad_gym",
+            "leader",
+            "leader_abs",
+        ],
+        help="Control mode to use",
+    )
+    parser.add_argument(
+        "--robot-type",
+        type=str,
+        default="so100",
+        help="Robot type (so100, koch, aloha, etc.)",
+    )
+
+    args = parser.parse_args()
+
+    robot_config = RobotConfig.get_choice_class(args.robot_type)(mock=False)
+    robot = make_robot_from_config(robot_config)
+
+    if not robot.is_connected:
+        robot.connect()
+
+    # Example bounds
+    bounds = {
+        "max": np.array([0.32170487, 0.201285, 0.10273342]),
+        "min": np.array([0.16631757, -0.08237468, 0.03364977]),
+    }
+
+    try:
+        # Determine controller type based on mode prefix
+        controller = None
+        if args.mode.startswith("keyboard"):
+            controller = KeyboardController(x_step_size=0.01, y_step_size=0.01, z_step_size=0.05)
+        elif args.mode.startswith("gamepad"):
+            if sys.platform == "darwin":
+                controller = GamepadControllerHID(x_step_size=0.01, y_step_size=0.01, z_step_size=0.05)
+            else:
+                controller = GamepadController(x_step_size=0.01, y_step_size=0.01, z_step_size=0.05)
+
+        # Handle mode categories
+        if args.mode in ["keyboard", "gamepad"]:
+            # Direct robot control modes
+            teleoperate_delta_inverse_kinematics(robot, controller, bounds=bounds, fps=10)
+
+        elif args.mode in ["keyboard_gym", "gamepad_gym"]:
+            # Gym environment control modes
+            cfg = HILSerlRobotEnvConfig(robot=robot_config, wrapper=EnvWrapperConfig())
+            cfg.wrapper.ee_action_space_params = EEActionSpaceConfig(
+                x_step_size=0.03, y_step_size=0.03, z_step_size=0.03, bounds=bounds
+            )
+            cfg.wrapper.ee_action_space_params.use_gamepad = False
+            cfg.device = "cpu"
+            env = make_robot_env(cfg, robot)
+            teleoperate_gym_env(env, controller, fps=cfg.fps)
+
+        elif args.mode == "leader":
+            # Leader-follower modes don't use controllers
+            teleoperate_delta_inverse_kinematics_with_leader(robot)
+
+        elif args.mode == "leader_abs":
+            teleoperate_inverse_kinematics_with_leader(robot)
+
+    finally:
+        if robot.is_connected:
+            robot.disconnect()

lerobot/scripts/server/find_joint_limits.py

@@ -0,0 +1,140 @@
+import argparse
+import time
+
+import cv2
+import numpy as np
+
+from lerobot.common.robot_devices.control_utils import is_headless
+from lerobot.common.robot_devices.robots.configs import RobotConfig
+from lerobot.common.robot_devices.robots.utils import make_robot_from_config
+from lerobot.configs import parser
+from lerobot.scripts.server.kinematics import RobotKinematics
+
+follower_port = "/dev/tty.usbmodem58760431631"
+leader_port = "/dev/tty.usbmodem58760433331"
+
+
+def find_joint_bounds(
+    robot,
+    control_time_s=30,
+    display_cameras=False,
+):
+    if not robot.is_connected:
+        robot.connect()
+
+    start_episode_t = time.perf_counter()
+    pos_list = []
+    while True:
+        observation, action = robot.teleop_step(record_data=True)
+
+        # Wait for 5 seconds to stabilize the robot initial position
+        if time.perf_counter() - start_episode_t < 5:
+            continue
+
+        pos_list.append(robot.follower_arms["main"].read("Present_Position"))
+
+        if display_cameras and not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+            cv2.waitKey(1)
+
+        if time.perf_counter() - start_episode_t > control_time_s:
+            max = np.max(np.stack(pos_list), 0)
+            min = np.min(np.stack(pos_list), 0)
+            print(f"Max angle position per joint {max}")
+            print(f"Min angle position per joint {min}")
+            break
+
+
+def find_ee_bounds(
+    robot,
+    control_time_s=30,
+    display_cameras=False,
+):
+    if not robot.is_connected:
+        robot.connect()
+
+    start_episode_t = time.perf_counter()
+    ee_list = []
+    while True:
+        observation, action = robot.teleop_step(record_data=True)
+
+        # Wait for 5 seconds to stabilize the robot initial position
+        if time.perf_counter() - start_episode_t < 5:
+            continue
+
+        kinematics = RobotKinematics(robot.robot_type)
+        joint_positions = robot.follower_arms["main"].read("Present_Position")
+        print(f"Joint positions: {joint_positions}")
+        ee_list.append(kinematics.fk_gripper_tip(joint_positions)[:3, 3])
+
+        if display_cameras and not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+            cv2.waitKey(1)
+
+        if time.perf_counter() - start_episode_t > control_time_s:
+            max = np.max(np.stack(ee_list), 0)
+            min = np.min(np.stack(ee_list), 0)
+            print(f"Max ee position {max}")
+            print(f"Min ee position {min}")
+            break
+
+
+def make_robot(robot_type="so100"):
+    """
+    Create a robot instance using the appropriate robot config class.
+
+    Args:
+        robot_type: Robot type string (e.g., "so100", "koch", "aloha")
+
+    Returns:
+        Robot instance
+    """
+
+    # Get the appropriate robot config class based on robot_type
+    robot_config = RobotConfig.get_choice_class(robot_type)(mock=False)
+    robot_config.leader_arms["main"].port = leader_port
+    robot_config.follower_arms["main"].port = follower_port
+
+    return make_robot_from_config(robot_config)
+
+
+if __name__ == "__main__":
+    # Create argparse for script-specific arguments
+    parser = argparse.ArgumentParser(add_help=False)  # Set add_help=False to avoid conflict
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="joint",
+        choices=["joint", "ee"],
+        help="Mode to run the script in. Can be 'joint' or 'ee'.",
+    )
+    parser.add_argument(
+        "--control-time-s",
+        type=int,
+        default=30,
+        help="Time step to use for control.",
+    )
+    parser.add_argument(
+        "--robot-type",
+        type=str,
+        default="so100",
+        help="Robot type (so100, koch, aloha, etc.)",
+    )
+
+    # Only parse known args, leaving robot config args for Hydra if used
+    args = parser.parse_args()
+
+    # Create robot with the appropriate config
+    robot = make_robot(args.robot_type)
+
+    if args.mode == "joint":
+        find_joint_bounds(robot, args.control_time_s)
+    elif args.mode == "ee":
+        find_ee_bounds(robot, args.control_time_s)
+
+    if robot.is_connected:
+        robot.disconnect()

lerobot/scripts/server/hilserl.proto

@@ -0,0 +1,55 @@
+// !/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.
+syntax = "proto3";
+
+package hil_serl;
+
+// LearnerService: the Actor calls this to push transitions.
+// The Learner implements this service.
+service LearnerService {
+  // Actor -> Learner to store transitions
+  rpc SendInteractionMessage(InteractionMessage) returns (Empty);
+  rpc StreamParameters(Empty) returns (stream Parameters);
+  rpc SendTransitions(stream Transition) returns (Empty);
+  rpc SendInteractions(stream InteractionMessage) returns (Empty);
+  rpc Ready(Empty) returns (Empty);
+}
+
+enum TransferState {
+    TRANSFER_UNKNOWN = 0;
+    TRANSFER_BEGIN = 1;
+    TRANSFER_MIDDLE = 2;
+    TRANSFER_END = 3;
+}
+
+// Messages
+message Transition {
+  TransferState transfer_state = 1;
+  bytes data = 2;
+}
+
+message Parameters {
+  TransferState transfer_state = 1;
+  bytes data = 2;
+}
+
+message InteractionMessage {
+  TransferState transfer_state = 1;
+  bytes data = 2;
+}
+
+message Empty {}

lerobot/scripts/server/hilserl_pb2.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: hilserl.proto
+# Protobuf Python Version: 5.29.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    29,
+    0,
+    '',
+    'hilserl.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rhilserl.proto\x12\x08hil_serl\"K\n\nTransition\x12/\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x17.hil_serl.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"K\n\nParameters\x12/\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x17.hil_serl.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"S\n\x12InteractionMessage\x12/\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x17.hil_serl.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\x07\n\x05\x45mpty*`\n\rTransferState\x12\x14\n\x10TRANSFER_UNKNOWN\x10\x00\x12\x12\n\x0eTRANSFER_BEGIN\x10\x01\x12\x13\n\x0fTRANSFER_MIDDLE\x10\x02\x12\x10\n\x0cTRANSFER_END\x10\x03\x32\xc2\x02\n\x0eLearnerService\x12G\n\x16SendInteractionMessage\x12\x1c.hil_serl.InteractionMessage\x1a\x0f.hil_serl.Empty\x12;\n\x10StreamParameters\x12\x0f.hil_serl.Empty\x1a\x14.hil_serl.Parameters0\x01\x12:\n\x0fSendTransitions\x12\x14.hil_serl.Transition\x1a\x0f.hil_serl.Empty(\x01\x12\x43\n\x10SendInteractions\x12\x1c.hil_serl.InteractionMessage\x1a\x0f.hil_serl.Empty(\x01\x12)\n\x05Ready\x12\x0f.hil_serl.Empty\x1a\x0f.hil_serl.Emptyb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'hilserl_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_TRANSFERSTATE']._serialized_start=275
+  _globals['_TRANSFERSTATE']._serialized_end=371
+  _globals['_TRANSITION']._serialized_start=27
+  _globals['_TRANSITION']._serialized_end=102
+  _globals['_PARAMETERS']._serialized_start=104
+  _globals['_PARAMETERS']._serialized_end=179
+  _globals['_INTERACTIONMESSAGE']._serialized_start=181
+  _globals['_INTERACTIONMESSAGE']._serialized_end=264
+  _globals['_EMPTY']._serialized_start=266
+  _globals['_EMPTY']._serialized_end=273
+  _globals['_LEARNERSERVICE']._serialized_start=374
+  _globals['_LEARNERSERVICE']._serialized_end=696
+# @@protoc_insertion_point(module_scope)

lerobot/scripts/server/hilserl_pb2_grpc.py

@@ -0,0 +1,276 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import hilserl_pb2 as hilserl__pb2
+
+GRPC_GENERATED_VERSION = '1.70.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in hilserl_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class LearnerServiceStub(object):
+    """LearnerService: the Actor calls this to push transitions.
+    The Learner implements this service.
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.SendInteractionMessage = channel.unary_unary(
+                '/hil_serl.LearnerService/SendInteractionMessage',
+                request_serializer=hilserl__pb2.InteractionMessage.SerializeToString,
+                response_deserializer=hilserl__pb2.Empty.FromString,
+                _registered_method=True)
+        self.StreamParameters = channel.unary_stream(
+                '/hil_serl.LearnerService/StreamParameters',
+                request_serializer=hilserl__pb2.Empty.SerializeToString,
+                response_deserializer=hilserl__pb2.Parameters.FromString,
+                _registered_method=True)
+        self.SendTransitions = channel.stream_unary(
+                '/hil_serl.LearnerService/SendTransitions',
+                request_serializer=hilserl__pb2.Transition.SerializeToString,
+                response_deserializer=hilserl__pb2.Empty.FromString,
+                _registered_method=True)
+        self.SendInteractions = channel.stream_unary(
+                '/hil_serl.LearnerService/SendInteractions',
+                request_serializer=hilserl__pb2.InteractionMessage.SerializeToString,
+                response_deserializer=hilserl__pb2.Empty.FromString,
+                _registered_method=True)
+        self.Ready = channel.unary_unary(
+                '/hil_serl.LearnerService/Ready',
+                request_serializer=hilserl__pb2.Empty.SerializeToString,
+                response_deserializer=hilserl__pb2.Empty.FromString,
+                _registered_method=True)
+
+
+class LearnerServiceServicer(object):
+    """LearnerService: the Actor calls this to push transitions.
+    The Learner implements this service.
+    """
+
+    def SendInteractionMessage(self, request, context):
+        """Actor -> Learner to store transitions
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def StreamParameters(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def SendTransitions(self, request_iterator, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def SendInteractions(self, request_iterator, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Ready(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_LearnerServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'SendInteractionMessage': grpc.unary_unary_rpc_method_handler(
+                    servicer.SendInteractionMessage,
+                    request_deserializer=hilserl__pb2.InteractionMessage.FromString,
+                    response_serializer=hilserl__pb2.Empty.SerializeToString,
+            ),
+            'StreamParameters': grpc.unary_stream_rpc_method_handler(
+                    servicer.StreamParameters,
+                    request_deserializer=hilserl__pb2.Empty.FromString,
+                    response_serializer=hilserl__pb2.Parameters.SerializeToString,
+            ),
+            'SendTransitions': grpc.stream_unary_rpc_method_handler(
+                    servicer.SendTransitions,
+                    request_deserializer=hilserl__pb2.Transition.FromString,
+                    response_serializer=hilserl__pb2.Empty.SerializeToString,
+            ),
+            'SendInteractions': grpc.stream_unary_rpc_method_handler(
+                    servicer.SendInteractions,
+                    request_deserializer=hilserl__pb2.InteractionMessage.FromString,
+                    response_serializer=hilserl__pb2.Empty.SerializeToString,
+            ),
+            'Ready': grpc.unary_unary_rpc_method_handler(
+                    servicer.Ready,
+                    request_deserializer=hilserl__pb2.Empty.FromString,
+                    response_serializer=hilserl__pb2.Empty.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'hil_serl.LearnerService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('hil_serl.LearnerService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class LearnerService(object):
+    """LearnerService: the Actor calls this to push transitions.
+    The Learner implements this service.
+    """
+
+    @staticmethod
+    def SendInteractionMessage(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/hil_serl.LearnerService/SendInteractionMessage',
+            hilserl__pb2.InteractionMessage.SerializeToString,
+            hilserl__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def StreamParameters(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_stream(
+            request,
+            target,
+            '/hil_serl.LearnerService/StreamParameters',
+            hilserl__pb2.Empty.SerializeToString,
+            hilserl__pb2.Parameters.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def SendTransitions(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_unary(
+            request_iterator,
+            target,
+            '/hil_serl.LearnerService/SendTransitions',
+            hilserl__pb2.Transition.SerializeToString,
+            hilserl__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def SendInteractions(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_unary(
+            request_iterator,
+            target,
+            '/hil_serl.LearnerService/SendInteractions',
+            hilserl__pb2.InteractionMessage.SerializeToString,
+            hilserl__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Ready(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/hil_serl.LearnerService/Ready',
+            hilserl__pb2.Empty.SerializeToString,
+            hilserl__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

lerobot/scripts/server/kinematics.py

@@ -0,0 +1,521 @@
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+
+def skew_symmetric(w):
+    """Creates the skew-symmetric matrix from a 3D vector."""
+    return np.array([[0, -w[2], w[1]], [w[2], 0, -w[0]], [-w[1], w[0], 0]])
+
+
+def rodrigues_rotation(w, theta):
+    """Computes the rotation matrix using Rodrigues' formula."""
+    w_hat = skew_symmetric(w)
+    return np.eye(3) + np.sin(theta) * w_hat + (1 - np.cos(theta)) * w_hat @ w_hat
+
+
+def screw_axis_to_transform(S, theta):
+    """Converts a screw axis to a 4x4 transformation matrix."""
+    S_w = S[:3]
+    S_v = S[3:]
+    if np.allclose(S_w, 0) and np.linalg.norm(S_v) == 1:  # Pure translation
+        T = np.eye(4)
+        T[:3, 3] = S_v * theta
+    elif np.linalg.norm(S_w) == 1:  # Rotation and translation
+        w_hat = skew_symmetric(S_w)
+        R = np.eye(3) + np.sin(theta) * w_hat + (1 - np.cos(theta)) * w_hat @ w_hat
+        t = (np.eye(3) * theta + (1 - np.cos(theta)) * w_hat + (theta - np.sin(theta)) * w_hat @ w_hat) @ S_v
+        T = np.eye(4)
+        T[:3, :3] = R
+        T[:3, 3] = t
+    else:
+        raise ValueError("Invalid screw axis parameters")
+    return T
+
+
+def pose_difference_se3(pose1, pose2):
+    """
+    Calculates the SE(3) difference between two 4x4 homogeneous transformation matrices.
+
+    pose1 - pose2
+
+    Args:
+        pose1: A 4x4 numpy array representing the first pose.
+        pose2: A 4x4 numpy array representing the second pose.
+
+    Returns:
+        A tuple (translation_diff, rotation_diff) where:
+        - translation_diff is a 3x1 numpy array representing the translational difference.
+        - rotation_diff is a 3x1 numpy array representing the rotational difference in axis-angle representation.
+    """
+
+    # Extract rotation matrices from poses
+    R1 = pose1[:3, :3]
+    R2 = pose2[:3, :3]
+
+    # Calculate translational difference
+    translation_diff = pose1[:3, 3] - pose2[:3, 3]
+
+    # Calculate rotational difference using scipy's Rotation library
+    R_diff = Rotation.from_matrix(R1 @ R2.T)
+    rotation_diff = R_diff.as_rotvec()  # Convert to axis-angle representation
+
+    return np.concatenate([translation_diff, rotation_diff])
+
+
+def se3_error(target_pose, current_pose):
+    pos_error = target_pose[:3, 3] - current_pose[:3, 3]
+    R_target = target_pose[:3, :3]
+    R_current = current_pose[:3, :3]
+    R_error = R_target @ R_current.T
+    rot_error = Rotation.from_matrix(R_error).as_rotvec()
+    return np.concatenate([pos_error, rot_error])
+
+
+class RobotKinematics:
+    """Robot kinematics class supporting multiple robot models."""
+
+    # Robot measurements dictionary
+    ROBOT_MEASUREMENTS = {
+        "koch": {
+            "gripper": [0.239, -0.001, 0.024],
+            "wrist": [0.209, 0, 0.024],
+            "forearm": [0.108, 0, 0.02],
+            "humerus": [0, 0, 0.036],
+            "shoulder": [0, 0, 0],
+            "base": [0, 0, 0.02],
+        },
+        "so100": {
+            "gripper": [0.320, 0, 0.050],
+            "wrist": [0.278, 0, 0.050],
+            "forearm": [0.143, 0, 0.044],
+            "humerus": [0.031, 0, 0.072],
+            "shoulder": [0, 0, 0],
+            "base": [0, 0, 0.02],
+        },
+        "moss": {
+            "gripper": [0.246, 0.013, 0.111],
+            "wrist": [0.245, 0.002, 0.064],
+            "forearm": [0.122, 0, 0.064],
+            "humerus": [0.001, 0.001, 0.063],
+            "shoulder": [0, 0, 0],
+            "base": [0, 0, 0.02],
+        },
+    }
+
+    def __init__(self, robot_type="so100"):
+        """Initialize kinematics for the specified robot type.
+
+        Args:
+            robot_type: String specifying the robot model ("koch", "so100", or "moss")
+        """
+        if robot_type not in self.ROBOT_MEASUREMENTS:
+            raise ValueError(
+                f"Unknown robot type: {robot_type}. Available types: {list(self.ROBOT_MEASUREMENTS.keys())}"
+            )
+
+        self.robot_type = robot_type
+        self.measurements = self.ROBOT_MEASUREMENTS[robot_type]
+
+        # Initialize all transformation matrices and screw axes
+        self._setup_transforms()
+
+    def _create_translation_matrix(self, x=0, y=0, z=0):
+        """Create a 4x4 translation matrix."""
+        return np.array([[1, 0, 0, x], [0, 1, 0, y], [0, 0, 1, z], [0, 0, 0, 1]])
+
+    def _setup_transforms(self):
+        """Setup all transformation matrices and screw axes for the robot."""
+        # Set up rotation matrices (constant across robot types)
+
+        # Gripper orientation
+        self.gripper_X0 = np.array(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 1, 0],
+                [0, -1, 0, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+
+        # Wrist orientation
+        self.wrist_X0 = np.array(
+            [
+                [0, -1, 0, 0],
+                [1, 0, 0, 0],
+                [0, 0, 1, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+
+        # Base orientation
+        self.base_X0 = np.array(
+            [
+                [0, 0, 1, 0],
+                [1, 0, 0, 0],
+                [0, 1, 0, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+
+        # Gripper
+        # Screw axis of gripper frame wrt base frame
+        self.S_BG = np.array(
+            [
+                1,
+                0,
+                0,
+                0,
+                self.measurements["gripper"][2],
+                -self.measurements["gripper"][1],
+            ]
+        )
+
+        # Gripper origin to centroid transform
+        self.X_GoGc = self._create_translation_matrix(x=0.07)
+
+        # Gripper origin to tip transform
+        self.X_GoGt = self._create_translation_matrix(x=0.12)
+
+        # 0-position gripper frame pose wrt base
+        self.X_BoGo = self._create_translation_matrix(
+            x=self.measurements["gripper"][0],
+            y=self.measurements["gripper"][1],
+            z=self.measurements["gripper"][2],
+        )
+
+        # Wrist
+        # Screw axis of wrist frame wrt base frame
+        self.S_BR = np.array([0, 1, 0, -self.measurements["wrist"][2], 0, self.measurements["wrist"][0]])
+
+        # 0-position origin to centroid transform
+        self.X_RoRc = self._create_translation_matrix(x=0.0035, y=-0.002)
+
+        # 0-position wrist frame pose wrt base
+        self.X_BR = self._create_translation_matrix(
+            x=self.measurements["wrist"][0],
+            y=self.measurements["wrist"][1],
+            z=self.measurements["wrist"][2],
+        )
+
+        # Forearm
+        # Screw axis of forearm frame wrt base frame
+        self.S_BF = np.array(
+            [
+                0,
+                1,
+                0,
+                -self.measurements["forearm"][2],
+                0,
+                self.measurements["forearm"][0],
+            ]
+        )
+
+        # Forearm origin + centroid transform
+        self.X_FoFc = self._create_translation_matrix(x=0.036)
+
+        # 0-position forearm frame pose wrt base
+        self.X_BF = self._create_translation_matrix(
+            x=self.measurements["forearm"][0],
+            y=self.measurements["forearm"][1],
+            z=self.measurements["forearm"][2],
+        )
+
+        # Humerus
+        # Screw axis of humerus frame wrt base frame
+        self.S_BH = np.array(
+            [
+                0,
+                -1,
+                0,
+                self.measurements["humerus"][2],
+                0,
+                -self.measurements["humerus"][0],
+            ]
+        )
+
+        # Humerus origin to centroid transform
+        self.X_HoHc = self._create_translation_matrix(x=0.0475)
+
+        # 0-position humerus frame pose wrt base
+        self.X_BH = self._create_translation_matrix(
+            x=self.measurements["humerus"][0],
+            y=self.measurements["humerus"][1],
+            z=self.measurements["humerus"][2],
+        )
+
+        # Shoulder
+        # Screw axis of shoulder frame wrt Base frame
+        self.S_BS = np.array([0, 0, -1, 0, 0, 0])
+
+        # Shoulder origin to centroid transform
+        self.X_SoSc = self._create_translation_matrix(x=-0.017, z=0.0235)
+
+        # 0-position shoulder frame pose wrt base
+        self.X_BS = self._create_translation_matrix(
+            x=self.measurements["shoulder"][0],
+            y=self.measurements["shoulder"][1],
+            z=self.measurements["shoulder"][2],
+        )
+
+        # Base
+        # Base origin to centroid transform
+        self.X_BoBc = self._create_translation_matrix(y=0.015)
+
+        # World to base transform
+        self.X_WoBo = self._create_translation_matrix(
+            x=self.measurements["base"][0],
+            y=self.measurements["base"][1],
+            z=self.measurements["base"][2],
+        )
+
+        # Pre-compute gripper post-multiplication matrix
+        self._fk_gripper_post = self.X_GoGc @ self.X_BoGo @ self.gripper_X0
+
+    def fk_base(self):
+        """Forward kinematics for the base frame."""
+        return self.X_WoBo @ self.X_BoBc @ self.base_X0
+
+    def fk_shoulder(self, robot_pos_deg):
+        """Forward kinematics for the shoulder frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return self.X_WoBo @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0]) @ self.X_SoSc @ self.X_BS
+
+    def fk_humerus(self, robot_pos_deg):
+        """Forward kinematics for the humerus frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return (
+            self.X_WoBo
+            @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0])
+            @ screw_axis_to_transform(self.S_BH, robot_pos_rad[1])
+            @ self.X_HoHc
+            @ self.X_BH
+        )
+
+    def fk_forearm(self, robot_pos_deg):
+        """Forward kinematics for the forearm frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return (
+            self.X_WoBo
+            @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0])
+            @ screw_axis_to_transform(self.S_BH, robot_pos_rad[1])
+            @ screw_axis_to_transform(self.S_BF, robot_pos_rad[2])
+            @ self.X_FoFc
+            @ self.X_BF
+        )
+
+    def fk_wrist(self, robot_pos_deg):
+        """Forward kinematics for the wrist frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return (
+            self.X_WoBo
+            @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0])
+            @ screw_axis_to_transform(self.S_BH, robot_pos_rad[1])
+            @ screw_axis_to_transform(self.S_BF, robot_pos_rad[2])
+            @ screw_axis_to_transform(self.S_BR, robot_pos_rad[3])
+            @ self.X_RoRc
+            @ self.X_BR
+            @ self.wrist_X0
+        )
+
+    def fk_gripper(self, robot_pos_deg):
+        """Forward kinematics for the gripper frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return (
+            self.X_WoBo
+            @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0])
+            @ screw_axis_to_transform(self.S_BH, robot_pos_rad[1])
+            @ screw_axis_to_transform(self.S_BF, robot_pos_rad[2])
+            @ screw_axis_to_transform(self.S_BR, robot_pos_rad[3])
+            @ screw_axis_to_transform(self.S_BG, robot_pos_rad[4])
+            @ self._fk_gripper_post
+        )
+
+    def fk_gripper_tip(self, robot_pos_deg):
+        """Forward kinematics for the gripper tip frame."""
+        robot_pos_rad = robot_pos_deg / 180 * np.pi
+        return (
+            self.X_WoBo
+            @ screw_axis_to_transform(self.S_BS, robot_pos_rad[0])
+            @ screw_axis_to_transform(self.S_BH, robot_pos_rad[1])
+            @ screw_axis_to_transform(self.S_BF, robot_pos_rad[2])
+            @ screw_axis_to_transform(self.S_BR, robot_pos_rad[3])
+            @ screw_axis_to_transform(self.S_BG, robot_pos_rad[4])
+            @ self.X_GoGt
+            @ self.X_BoGo
+            @ self.gripper_X0
+        )
+
+    def compute_jacobian(self, robot_pos_deg, fk_func=None):
+        """Finite differences to compute the Jacobian.
+        J(i, j) represents how the ith component of the end-effector's velocity changes wrt a small change
+        in the jth joint's velocity.
+
+        Args:
+            robot_pos_deg: Current joint positions in degrees
+            fk_func: Forward kinematics function to use (defaults to fk_gripper)
+        """
+        if fk_func is None:
+            fk_func = self.fk_gripper
+
+        eps = 1e-8
+        jac = np.zeros(shape=(6, 5))
+        delta = np.zeros(len(robot_pos_deg[:-1]), dtype=np.float64)
+        for el_ix in range(len(robot_pos_deg[:-1])):
+            delta *= 0
+            delta[el_ix] = eps / 2
+            Sdot = (
+                pose_difference_se3(
+                    fk_func(robot_pos_deg[:-1] + delta),
+                    fk_func(robot_pos_deg[:-1] - delta),
+                )
+                / eps
+            )
+            jac[:, el_ix] = Sdot
+        return jac
+
+    def compute_positional_jacobian(self, robot_pos_deg, fk_func=None):
+        """Finite differences to compute the positional Jacobian.
+        J(i, j) represents how the ith component of the end-effector's position changes wrt a small change
+        in the jth joint's velocity.
+
+        Args:
+            robot_pos_deg: Current joint positions in degrees
+            fk_func: Forward kinematics function to use (defaults to fk_gripper)
+        """
+        if fk_func is None:
+            fk_func = self.fk_gripper
+
+        eps = 1e-8
+        jac = np.zeros(shape=(3, 5))
+        delta = np.zeros(len(robot_pos_deg[:-1]), dtype=np.float64)
+        for el_ix in range(len(robot_pos_deg[:-1])):
+            delta *= 0
+            delta[el_ix] = eps / 2
+            Sdot = (
+                fk_func(robot_pos_deg[:-1] + delta)[:3, 3] - fk_func(robot_pos_deg[:-1] - delta)[:3, 3]
+            ) / eps
+            jac[:, el_ix] = Sdot
+        return jac
+
+    def ik(self, current_joint_state, desired_ee_pose, position_only=True, fk_func=None):
+        """Inverse kinematics using gradient descent.
+
+        Args:
+            current_joint_state: Initial joint positions in degrees
+            desired_ee_pose: Target end-effector pose as a 4x4 transformation matrix
+            position_only: If True, only match end-effector position, not orientation
+            fk_func: Forward kinematics function to use (defaults to fk_gripper)
+
+        Returns:
+            Joint positions in degrees that achieve the desired end-effector pose
+        """
+        if fk_func is None:
+            fk_func = self.fk_gripper
+
+        # Do gradient descent.
+        max_iterations = 5
+        learning_rate = 1
+        for _ in range(max_iterations):
+            current_ee_pose = fk_func(current_joint_state)
+            if not position_only:
+                error = se3_error(desired_ee_pose, current_ee_pose)
+                jac = self.compute_jacobian(current_joint_state, fk_func)
+            else:
+                error = desired_ee_pose[:3, 3] - current_ee_pose[:3, 3]
+                jac = self.compute_positional_jacobian(current_joint_state, fk_func)
+            delta_angles = np.linalg.pinv(jac) @ error
+            current_joint_state[:-1] += learning_rate * delta_angles
+
+            if np.linalg.norm(error) < 5e-3:
+                return current_joint_state
+        return current_joint_state
+
+
+if __name__ == "__main__":
+    import time
+
+    def run_test(robot_type):
+        """Run test suite for a specific robot type."""
+        print(f"\n--- Testing {robot_type.upper()} Robot ---")
+
+        # Initialize kinematics for this robot
+        robot = RobotKinematics(robot_type)
+
+        # Test 1: Forward kinematics consistency
+        print("Test 1: Forward kinematics consistency")
+        test_angles = np.array([30, 45, -30, 20, 10, 0])  # Example joint angles in degrees
+
+        # Calculate FK for different joints
+        shoulder_pose = robot.fk_shoulder(test_angles)
+        humerus_pose = robot.fk_humerus(test_angles)
+        forearm_pose = robot.fk_forearm(test_angles)
+        wrist_pose = robot.fk_wrist(test_angles)
+        gripper_pose = robot.fk_gripper(test_angles)
+        gripper_tip_pose = robot.fk_gripper_tip(test_angles)
+
+        # Check that poses form a consistent kinematic chain (positions should be progressively further from origin)
+        distances = [
+            np.linalg.norm(shoulder_pose[:3, 3]),
+            np.linalg.norm(humerus_pose[:3, 3]),
+            np.linalg.norm(forearm_pose[:3, 3]),
+            np.linalg.norm(wrist_pose[:3, 3]),
+            np.linalg.norm(gripper_pose[:3, 3]),
+            np.linalg.norm(gripper_tip_pose[:3, 3]),
+        ]
+
+        # Check if distances generally increase along the chain
+        is_consistent = all(distances[i] <= distances[i + 1] for i in range(len(distances) - 1))
+        print(f"  Pose distances from origin: {[round(d, 3) for d in distances]}")
+        print(f"  Kinematic chain consistency: {'PASSED' if is_consistent else 'FAILED'}")
+
+        # Test 2: Jacobian computation
+        print("Test 2: Jacobian computation")
+        jacobian = robot.compute_jacobian(test_angles)
+        positional_jacobian = robot.compute_positional_jacobian(test_angles)
+
+        # Check shapes
+        jacobian_shape_ok = jacobian.shape == (6, 5)
+        pos_jacobian_shape_ok = positional_jacobian.shape == (3, 5)
+
+        print(f"  Jacobian shape: {'PASSED' if jacobian_shape_ok else 'FAILED'}")
+        print(f"  Positional Jacobian shape: {'PASSED' if pos_jacobian_shape_ok else 'FAILED'}")
+
+        # Test 3: Inverse kinematics
+        print("Test 3: Inverse kinematics (position only)")
+
+        # Generate target pose from known joint angles
+        original_angles = np.array([10, 20, 30, -10, 5, 0])
+        target_pose = robot.fk_gripper(original_angles)
+
+        # Start IK from a different position
+        initial_guess = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+
+        # Measure IK performance
+        start_time = time.time()
+        computed_angles = robot.ik(initial_guess.copy(), target_pose)
+        ik_time = time.time() - start_time
+
+        # Compute resulting pose from IK solution
+        result_pose = robot.fk_gripper(computed_angles)
+
+        # Calculate position error
+        pos_error = np.linalg.norm(target_pose[:3, 3] - result_pose[:3, 3])
+        passed = pos_error < 0.01  # Accept errors less than 1cm
+
+        print(f"  IK computation time: {ik_time:.4f} seconds")
+        print(f"  Position error: {pos_error:.4f}")
+        print(f"  IK position accuracy: {'PASSED' if passed else 'FAILED'}")
+
+        return is_consistent and jacobian_shape_ok and pos_jacobian_shape_ok and passed
+
+    # Run tests for all robot types
+    results = {}
+    for robot_type in ["koch", "so100", "moss"]:
+        results[robot_type] = run_test(robot_type)
+
+    # Print overall summary
+    print("\n=== Test Summary ===")
+    all_passed = all(results.values())
+    for robot_type, passed in results.items():
+        print(f"{robot_type.upper()}: {'PASSED' if passed else 'FAILED'}")
+    print(f"\nOverall: {'ALL TESTS PASSED' if all_passed else 'SOME TESTS FAILED'}")

lerobot/scripts/server/learner_service.py

@@ -0,0 +1,80 @@
+import logging
+from multiprocessing import Event, Queue
+
+import hilserl_pb2  # type: ignore
+import hilserl_pb2_grpc  # type: ignore
+
+from lerobot.scripts.server.network_utils import receive_bytes_in_chunks, send_bytes_in_chunks
+
+MAX_MESSAGE_SIZE = 4 * 1024 * 1024  # 4 MB
+MAX_WORKERS = 3  # Stream parameters, send transitions and interactions
+STUTDOWN_TIMEOUT = 10
+
+
+class LearnerService(hilserl_pb2_grpc.LearnerServiceServicer):
+    def __init__(
+        self,
+        shutdown_event: Event,
+        parameters_queue: Queue,
+        seconds_between_pushes: float,
+        transition_queue: Queue,
+        interaction_message_queue: Queue,
+    ):
+        self.shutdown_event = shutdown_event
+        self.parameters_queue = parameters_queue
+        self.seconds_between_pushes = seconds_between_pushes
+        self.transition_queue = transition_queue
+        self.interaction_message_queue = interaction_message_queue
+
+    def StreamParameters(self, request, context):
+        # TODO: authorize the request
+        logging.info("[LEARNER] Received request to stream parameters from the Actor")
+
+        while not self.shutdown_event.is_set():
+            logging.info("[LEARNER] Push parameters to the Actor")
+            buffer = self.parameters_queue.get()
+
+            yield from send_bytes_in_chunks(
+                buffer,
+                hilserl_pb2.Parameters,
+                log_prefix="[LEARNER] Sending parameters",
+                silent=True,
+            )
+
+            logging.info("[LEARNER] Parameters sent")
+
+            self.shutdown_event.wait(self.seconds_between_pushes)
+
+        logging.info("[LEARNER] Stream parameters finished")
+        return hilserl_pb2.Empty()
+
+    def SendTransitions(self, request_iterator, _context):
+        # TODO: authorize the request
+        logging.info("[LEARNER] Received request to receive transitions from the Actor")
+
+        receive_bytes_in_chunks(
+            request_iterator,
+            self.transition_queue,
+            self.shutdown_event,
+            log_prefix="[LEARNER] transitions",
+        )
+
+        logging.debug("[LEARNER] Finished receiving transitions")
+        return hilserl_pb2.Empty()
+
+    def SendInteractions(self, request_iterator, _context):
+        # TODO: authorize the request
+        logging.info("[LEARNER] Received request to receive interactions from the Actor")
+
+        receive_bytes_in_chunks(
+            request_iterator,
+            self.interaction_message_queue,
+            self.shutdown_event,
+            log_prefix="[LEARNER] interactions",
+        )
+
+        logging.debug("[LEARNER] Finished receiving interactions")
+        return hilserl_pb2.Empty()
+
+    def Ready(self, request, context):
+        return hilserl_pb2.Empty()

lerobot/scripts/server/maniskill_manipulator.py

@@ -0,0 +1,268 @@
+from typing import Any
+
+import einops
+import gymnasium as gym
+import numpy as np
+import torch
+from mani_skill.utils.wrappers.record import RecordEpisode
+from mani_skill.vector.wrappers.gymnasium import ManiSkillVectorEnv
+
+from lerobot.common.envs.configs import ManiskillEnvConfig
+
+
+def preprocess_maniskill_observation(
+    observations: dict[str, np.ndarray],
+) -> dict[str, torch.Tensor]:
+    """Convert environment observation to LeRobot format observation.
+    Args:
+        observation: Dictionary of observation batches from a Gym vector environment.
+    Returns:
+        Dictionary of observation batches with keys renamed to LeRobot format and values as tensors.
+    """
+    # map to expected inputs for the policy
+    return_observations = {}
+    # TODO: You have to merge all tensors from agent key and extra key
+    # You don't keep sensor param key in the observation
+    # And you keep sensor data rgb
+    q_pos = observations["agent"]["qpos"]
+    q_vel = observations["agent"]["qvel"]
+    tcp_pos = observations["extra"]["tcp_pose"]
+    img = observations["sensor_data"]["base_camera"]["rgb"]
+
+    _, h, w, c = img.shape
+    assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
+
+    # sanity check that images are uint8
+    assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"
+
+    # convert to channel first of type float32 in range [0,1]
+    img = einops.rearrange(img, "b h w c -> b c h w").contiguous()
+    img = img.type(torch.float32)
+    img /= 255
+
+    state = torch.cat([q_pos, q_vel, tcp_pos], dim=-1)
+
+    return_observations["observation.image"] = img
+    return_observations["observation.state"] = state
+    return return_observations
+
+
+class ManiSkillObservationWrapper(gym.ObservationWrapper):
+    def __init__(self, env, device: torch.device = "cuda"):
+        super().__init__(env)
+        if isinstance(device, str):
+            device = torch.device(device)
+        self.device = device
+
+    def observation(self, observation):
+        observation = preprocess_maniskill_observation(observation)
+        observation = {k: v.to(self.device) for k, v in observation.items()}
+        return observation
+
+
+class ManiSkillCompat(gym.Wrapper):
+    def __init__(self, env):
+        super().__init__(env)
+        new_action_space_shape = env.action_space.shape[-1]
+        new_low = np.squeeze(env.action_space.low, axis=0)
+        new_high = np.squeeze(env.action_space.high, axis=0)
+        self.action_space = gym.spaces.Box(low=new_low, high=new_high, shape=(new_action_space_shape,))
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        options = {}
+        return super().reset(seed=seed, options=options)
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        reward = reward.item()
+        terminated = terminated.item()
+        truncated = truncated.item()
+        return obs, reward, terminated, truncated, info
+
+
+class ManiSkillActionWrapper(gym.ActionWrapper):
+    def __init__(self, env):
+        super().__init__(env)
+        self.action_space = gym.spaces.Tuple(spaces=(env.action_space, gym.spaces.Discrete(2)))
+
+    def action(self, action):
+        action, telop = action
+        return action
+
+
+class ManiSkillMultiplyActionWrapper(gym.Wrapper):
+    def __init__(self, env, multiply_factor: float = 1):
+        super().__init__(env)
+        self.multiply_factor = multiply_factor
+        action_space_agent: gym.spaces.Box = env.action_space[0]
+        action_space_agent.low = action_space_agent.low * multiply_factor
+        action_space_agent.high = action_space_agent.high * multiply_factor
+        self.action_space = gym.spaces.Tuple(spaces=(action_space_agent, gym.spaces.Discrete(2)))
+
+    def step(self, action):
+        if isinstance(action, tuple):
+            action, telop = action
+        else:
+            telop = 0
+        action = action / self.multiply_factor
+        obs, reward, terminated, truncated, info = self.env.step((action, telop))
+        return obs, reward, terminated, truncated, info
+
+
+class BatchCompatibleWrapper(gym.ObservationWrapper):
+    """Ensures observations are batch-compatible by adding a batch dimension if necessary."""
+
+    def __init__(self, env):
+        super().__init__(env)
+
+    def observation(self, observation: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        for key in observation:
+            if "image" in key and observation[key].dim() == 3:
+                observation[key] = observation[key].unsqueeze(0)
+            if "state" in key and observation[key].dim() == 1:
+                observation[key] = observation[key].unsqueeze(0)
+        return observation
+
+
+class TimeLimitWrapper(gym.Wrapper):
+    """Adds a time limit to the environment based on fps and control_time."""
+
+    def __init__(self, env, control_time_s, fps):
+        super().__init__(env)
+        self.control_time_s = control_time_s
+        self.fps = fps
+        self.max_episode_steps = int(self.control_time_s * self.fps)
+        self.current_step = 0
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        self.current_step += 1
+
+        if self.current_step >= self.max_episode_steps:
+            terminated = True
+
+        return obs, reward, terminated, truncated, info
+
+    def reset(self, *, seed=None, options=None):
+        self.current_step = 0
+        return super().reset(seed=seed, options=options)
+
+
+class ManiskillMockGripperWrapper(gym.Wrapper):
+    def __init__(self, env, nb_discrete_actions: int = 3):
+        super().__init__(env)
+        new_shape = env.action_space[0].shape[0] + 1
+        new_low = np.concatenate([env.action_space[0].low, [0]])
+        new_high = np.concatenate([env.action_space[0].high, [nb_discrete_actions - 1]])
+        action_space_agent = gym.spaces.Box(low=new_low, high=new_high, shape=(new_shape,))
+        self.action_space = gym.spaces.Tuple((action_space_agent, env.action_space[1]))
+
+    def step(self, action):
+        if isinstance(action, tuple):
+            action_agent, telop_action = action
+        else:
+            telop_action = 0
+            action_agent = action
+        real_action = action_agent[:-1]
+        final_action = (real_action, telop_action)
+        obs, reward, terminated, truncated, info = self.env.step(final_action)
+        return obs, reward, terminated, truncated, info
+
+
+def make_maniskill(
+    cfg: ManiskillEnvConfig,
+    n_envs: int | None = None,
+) -> gym.Env:
+    """
+    Factory function to create a ManiSkill environment with standard wrappers.
+
+    Args:
+        cfg: Configuration for the ManiSkill environment
+        n_envs: Number of parallel environments
+
+    Returns:
+        A wrapped ManiSkill environment
+    """
+    env = gym.make(
+        cfg.task,
+        obs_mode=cfg.obs_type,
+        control_mode=cfg.control_mode,
+        render_mode=cfg.render_mode,
+        sensor_configs={"width": cfg.image_size, "height": cfg.image_size},
+        num_envs=n_envs,
+    )
+
+    # Add video recording if enabled
+    if cfg.video_record.enabled:
+        env = RecordEpisode(
+            env,
+            output_dir=cfg.video_record.record_dir,
+            save_trajectory=True,
+            trajectory_name=cfg.video_record.trajectory_name,
+            save_video=True,
+            video_fps=30,
+        )
+
+    # Add observation and image processing
+    env = ManiSkillObservationWrapper(env, device=cfg.device)
+    env = ManiSkillVectorEnv(env, ignore_terminations=True, auto_reset=False)
+    env._max_episode_steps = env.max_episode_steps = cfg.episode_length
+    env.unwrapped.metadata["render_fps"] = cfg.fps
+
+    # Add compatibility wrappers
+    env = ManiSkillCompat(env)
+    env = ManiSkillActionWrapper(env)
+    env = ManiSkillMultiplyActionWrapper(env, multiply_factor=0.03)  # Scale actions for better control
+    if cfg.mock_gripper:
+        env = ManiskillMockGripperWrapper(env, nb_discrete_actions=3)
+
+    return env
+
+
+# @parser.wrap()
+# def main(cfg: TrainPipelineConfig):
+#     """Main function to run the ManiSkill environment."""
+#     # Create the ManiSkill environment
+#     env = make_maniskill(cfg.env, n_envs=1)
+
+#     # Reset the environment
+#     obs, info = env.reset()
+
+#     # Run a simple interaction loop
+#     sum_reward = 0
+#     for i in range(100):
+#         # Sample a random action
+#         action = env.action_space.sample()
+
+#         # Step the environment
+#         start_time = time.perf_counter()
+#         obs, reward, terminated, truncated, info = env.step(action)
+#         step_time = time.perf_counter() - start_time
+#         sum_reward += reward
+#         # Log information
+
+#         # Reset if episode terminated
+#         if terminated or truncated:
+#             logging.info(f"Step {i}, reward: {sum_reward}, step time: {step_time}s")
+#             sum_reward = 0
+#             obs, info = env.reset()
+
+#     # Close the environment
+#     env.close()
+
+
+# if __name__ == "__main__":
+#     logging.basicConfig(level=logging.INFO)
+#     main()
+
+if __name__ == "__main__":
+    import draccus
+
+    config = ManiskillEnvConfig()
+    draccus.set_config_type("json")
+    draccus.dump(
+        config=config,
+        stream=open(file="run_config.json", mode="w"),
+    )

lerobot/scripts/server/network_utils.py

@@ -0,0 +1,95 @@
+#!/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.
+
+import io
+import logging
+from multiprocessing import Event, Queue
+from typing import Any
+
+from lerobot.scripts.server import hilserl_pb2
+
+CHUNK_SIZE = 2 * 1024 * 1024  # 2 MB
+
+
+def bytes_buffer_size(buffer: io.BytesIO) -> int:
+    buffer.seek(0, io.SEEK_END)
+    result = buffer.tell()
+    buffer.seek(0)
+    return result
+
+
+def send_bytes_in_chunks(buffer: bytes, message_class: Any, log_prefix: str = "", silent: bool = True):
+    buffer = io.BytesIO(buffer)
+    size_in_bytes = bytes_buffer_size(buffer)
+
+    sent_bytes = 0
+
+    logging_method = logging.info if not silent else logging.debug
+
+    logging_method(f"{log_prefix} Buffer size {size_in_bytes / 1024 / 1024} MB with")
+
+    while sent_bytes < size_in_bytes:
+        transfer_state = hilserl_pb2.TransferState.TRANSFER_MIDDLE
+
+        if sent_bytes + CHUNK_SIZE >= size_in_bytes:
+            transfer_state = hilserl_pb2.TransferState.TRANSFER_END
+        elif sent_bytes == 0:
+            transfer_state = hilserl_pb2.TransferState.TRANSFER_BEGIN
+
+        size_to_read = min(CHUNK_SIZE, size_in_bytes - sent_bytes)
+        chunk = buffer.read(size_to_read)
+
+        yield message_class(transfer_state=transfer_state, data=chunk)
+        sent_bytes += size_to_read
+        logging_method(f"{log_prefix} Sent {sent_bytes}/{size_in_bytes} bytes with state {transfer_state}")
+
+    logging_method(f"{log_prefix} Published {sent_bytes / 1024 / 1024} MB")
+
+
+def receive_bytes_in_chunks(iterator, queue: Queue, shutdown_event: Event, log_prefix: str = ""):
+    bytes_buffer = io.BytesIO()
+    step = 0
+
+    logging.info(f"{log_prefix} Starting receiver")
+    for item in iterator:
+        logging.debug(f"{log_prefix} Received item")
+        if shutdown_event.is_set():
+            logging.info(f"{log_prefix} Shutting down receiver")
+            return
+
+        if item.transfer_state == hilserl_pb2.TransferState.TRANSFER_BEGIN:
+            bytes_buffer.seek(0)
+            bytes_buffer.truncate(0)
+            bytes_buffer.write(item.data)
+            logging.debug(f"{log_prefix} Received data at step 0")
+            step = 0
+            continue
+        elif item.transfer_state == hilserl_pb2.TransferState.TRANSFER_MIDDLE:
+            bytes_buffer.write(item.data)
+            step += 1
+            logging.debug(f"{log_prefix} Received data at step {step}")
+        elif item.transfer_state == hilserl_pb2.TransferState.TRANSFER_END:
+            bytes_buffer.write(item.data)
+            logging.debug(f"{log_prefix} Received data at step end size {bytes_buffer_size(bytes_buffer)}")
+
+            queue.put(bytes_buffer.getvalue())
+
+            bytes_buffer.seek(0)
+            bytes_buffer.truncate(0)
+            step = 0
+
+            logging.debug(f"{log_prefix} Queue updated")

lerobot/scripts/server/utils.py

@@ -0,0 +1,73 @@
+#!/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.
+
+import logging
+import signal
+import sys
+from queue import Empty
+
+from torch.multiprocessing import Queue
+
+shutdown_event_counter = 0
+
+
+def setup_process_handlers(use_threads: bool) -> any:
+    if use_threads:
+        from threading import Event
+    else:
+        from multiprocessing import Event
+
+    shutdown_event = Event()
+
+    # Define signal handler
+    def signal_handler(signum, frame):
+        logging.info("Shutdown signal received. Cleaning up...")
+        shutdown_event.set()
+        global shutdown_event_counter
+        shutdown_event_counter += 1
+
+        if shutdown_event_counter > 1:
+            logging.info("Force shutdown")
+            sys.exit(1)
+
+    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
+    signal.signal(signal.SIGTERM, signal_handler)  # Termination request (kill)
+    signal.signal(signal.SIGHUP, signal_handler)  # Terminal closed/Hangup
+    signal.signal(signal.SIGQUIT, signal_handler)  # Ctrl+\
+
+    def signal_handler(signum, frame):
+        logging.info("Shutdown signal received. Cleaning up...")
+        shutdown_event.set()
+
+    return shutdown_event
+
+
+def get_last_item_from_queue(queue: Queue):
+    item = queue.get()
+    counter = 1
+
+    # Drain queue and keep only the most recent parameters
+    try:
+        while True:
+            item = queue.get_nowait()
+            counter += 1
+    except Empty:
+        pass
+
+    logging.debug(f"Drained {counter} items from queue")
+
+    return item

lerobot/scripts/train.py

@@ -53,6 +53,84 @@ from lerobot.configs.train import TrainPipelineConfig
 from lerobot.scripts.eval import eval_policy
 
 
+def make_optimizer_and_scheduler(cfg, policy):
+    if cfg.policy.name == "act":
+        optimizer_params_dicts = [
+            {
+                "params": [
+                    p
+                    for n, p in policy.named_parameters()
+                    if not n.startswith("model.backbone") and p.requires_grad
+                ]
+            },
+            {
+                "params": [
+                    p
+                    for n, p in policy.named_parameters()
+                    if n.startswith("model.backbone") and p.requires_grad
+                ],
+                "lr": cfg.training.lr_backbone,
+            },
+        ]
+        optimizer = torch.optim.AdamW(
+            optimizer_params_dicts,
+            lr=cfg.training.lr,
+            weight_decay=cfg.training.weight_decay,
+        )
+        lr_scheduler = None
+    elif cfg.policy.name == "diffusion":
+        optimizer = torch.optim.Adam(
+            policy.diffusion.parameters(),
+            cfg.training.lr,
+            cfg.training.adam_betas,
+            cfg.training.adam_eps,
+            cfg.training.adam_weight_decay,
+        )
+        from diffusers.optimization import get_scheduler
+
+        lr_scheduler = get_scheduler(
+            cfg.training.lr_scheduler,
+            optimizer=optimizer,
+            num_warmup_steps=cfg.training.lr_warmup_steps,
+            num_training_steps=cfg.training.offline_steps,
+        )
+    elif policy.name == "tdmpc":
+        optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
+        lr_scheduler = None
+
+    elif policy.name == "sac":
+        optimizer = torch.optim.Adam(
+            [
+                {"params": policy.actor.parameters(), "lr": policy.config.actor_lr},
+                {
+                    "params": policy.critic_ensemble.parameters(),
+                    "lr": policy.config.critic_lr,
+                },
+                {
+                    "params": policy.temperature.parameters(),
+                    "lr": policy.config.temperature_lr,
+                },
+            ]
+        )
+        lr_scheduler = None
+
+    elif cfg.policy.name == "vqbet":
+        from lerobot.common.policies.vqbet.modeling_vqbet import (
+            VQBeTOptimizer,
+            VQBeTScheduler,
+        )
+
+        optimizer = VQBeTOptimizer(policy, cfg)
+        lr_scheduler = VQBeTScheduler(optimizer, cfg)
+    elif cfg.policy.name == "hilserl_classifier":
+        optimizer = torch.optim.AdamW(policy.parameters(), cfg.policy.learning_rate)
+        lr_scheduler = None
+    else:
+        raise NotImplementedError()
+
+    return optimizer, lr_scheduler
+
+
 def update_policy(
     train_metrics: MetricsTracker,
     policy: PreTrainedPolicy,
@@ -196,7 +274,11 @@ def train(cfg: TrainPipelineConfig):
     }
 
     train_tracker = MetricsTracker(
-        cfg.batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
+        cfg.batch_size,
+        dataset.num_frames,
+        dataset.num_episodes,
+        train_metrics,
+        initial_step=step,
     )
 
     logging.info("Start offline training on a fixed dataset")
@@ -267,7 +349,11 @@ def train(cfg: TrainPipelineConfig):
                 "eval_s": AverageMeter("eval_s", ":.3f"),
             }
             eval_tracker = MetricsTracker(
-                cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
+                cfg.batch_size,
+                dataset.num_frames,
+                dataset.num_episodes,
+                eval_metrics,
+                initial_step=step,
             )
             eval_tracker.eval_s = eval_info["aggregated"].pop("eval_s")
             eval_tracker.avg_sum_reward = eval_info["aggregated"].pop("avg_sum_reward")

lerobot/scripts/train_hilserl_classifier.py

@@ -0,0 +1,466 @@
+# 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.
+import logging
+import time
+from contextlib import nullcontext
+from pprint import pformat
+
+import numpy as np
+import torch
+import torch.nn as nn
+import wandb
+from deepdiff import DeepDiff
+from omegaconf import DictConfig, OmegaConf
+from termcolor import colored
+from torch import optim
+from torch.autograd import profiler
+from torch.cuda.amp import GradScaler
+from torch.utils.data import DataLoader, RandomSampler, WeightedRandomSampler
+from tqdm import tqdm
+
+from lerobot.common.datasets.factory import resolve_delta_timestamps
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.policies.hilserl.classifier.configuration_classifier import ClassifierConfig
+from lerobot.common.policies.hilserl.classifier.modeling_classifier import Classifier
+from lerobot.common.utils.utils import (
+    format_big_number,
+    get_safe_torch_device,
+    init_hydra_config,
+    init_logging,
+    set_global_seed,
+)
+from lerobot.scripts.server.buffer import random_shift
+
+
+def get_model(cfg, logger):  # noqa I001
+    classifier_config = _policy_cfg_from_hydra_cfg(ClassifierConfig, cfg)
+    model = Classifier(classifier_config)
+    if cfg.resume:
+        model.load_state_dict(Classifier.from_pretrained(str(logger.last_pretrained_model_dir)).state_dict())
+    return model
+
+
+def create_balanced_sampler(dataset, cfg):
+    # Get underlying dataset if using Subset
+    original_dataset = dataset.dataset if isinstance(dataset, torch.utils.data.Subset) else dataset
+
+    # Get indices if using Subset (for slicing)
+    indices = dataset.indices if isinstance(dataset, torch.utils.data.Subset) else None
+
+    # Get labels from Hugging Face dataset
+    if indices is not None:
+        # Get subset of labels using Hugging Face's select()
+        hf_subset = original_dataset.hf_dataset.select(indices)
+        labels = hf_subset[cfg.training.label_key]
+    else:
+        # Get all labels directly
+        labels = original_dataset.hf_dataset[cfg.training.label_key]
+
+    labels = torch.stack(labels)
+    _, counts = torch.unique(labels, return_counts=True)
+    class_weights = 1.0 / counts.float()
+    sample_weights = class_weights[labels]
+
+    return WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True)
+
+
+def support_amp(device: torch.device, cfg: DictConfig) -> bool:
+    # Check if the device supports AMP
+    # Here is an example of the issue that says that MPS doesn't support AMP properply
+    return cfg.training.use_amp and device.type in ("cuda", "cpu")
+
+
+def train_epoch(model, train_loader, criterion, optimizer, grad_scaler, device, logger, step, cfg):
+    # Single epoch training loop with AMP support and progress tracking
+    model.train()
+    correct = 0
+    total = 0
+
+    pbar = tqdm(train_loader, desc="Training")
+    for batch_idx, batch in enumerate(pbar):
+        start_time = time.perf_counter()
+        images = [batch[img_key].to(device) for img_key in cfg.training.image_keys]
+        images = [random_shift(img, 4) for img in images]
+        labels = batch[cfg.training.label_key].float().to(device)
+
+        # Forward pass with optional AMP
+        with torch.autocast(device_type=device.type) if support_amp(device, cfg) else nullcontext():
+            outputs = model(images)
+            loss = criterion(outputs.logits, labels)
+
+        # Backward pass with gradient scaling if AMP enabled
+        optimizer.zero_grad()
+        if cfg.training.use_amp:
+            grad_scaler.scale(loss).backward()
+            grad_scaler.step(optimizer)
+            grad_scaler.update()
+        else:
+            loss.backward()
+            optimizer.step()
+
+        # Track metrics
+        if model.config.num_classes == 2:
+            predictions = (torch.sigmoid(outputs.logits) > 0.5).float()
+        else:
+            predictions = torch.argmax(outputs.logits, dim=1)
+        correct += (predictions == labels).sum().item()
+        total += labels.size(0)
+
+        current_acc = 100 * correct / total
+        train_info = {
+            "loss": loss.item(),
+            "accuracy": current_acc,
+            "dataloading_s": time.perf_counter() - start_time,
+        }
+
+        logger.log_dict(train_info, step + batch_idx, mode="train")
+        pbar.set_postfix({"loss": f"{loss.item():.4f}", "acc": f"{current_acc:.2f}%"})
+
+
+def validate(model, val_loader, criterion, device, logger, cfg):
+    # Validation loop with metric tracking and sample logging
+    model.eval()
+    correct = 0
+    total = 0
+    batch_start_time = time.perf_counter()
+    samples = []
+    running_loss = 0
+    inference_times = []
+
+    with (
+        torch.no_grad(),
+        torch.autocast(device_type=device.type) if support_amp(device, cfg) else nullcontext(),
+    ):
+        for batch in tqdm(val_loader, desc="Validation"):
+            images = [batch[img_key].to(device) for img_key in cfg.training.image_keys]
+            labels = batch[cfg.training.label_key].float().to(device)
+
+            if cfg.training.profile_inference_time and logger._cfg.wandb.enable:
+                with (
+                    profiler.profile(record_shapes=True) as prof,
+                    profiler.record_function("model_inference"),
+                ):
+                    outputs = model(images)
+                inference_times.append(
+                    next(x for x in prof.key_averages() if x.key == "model_inference").cpu_time
+                )
+            else:
+                outputs = model(images)
+
+            loss = criterion(outputs.logits, labels)
+
+            # Track metrics
+            if model.config.num_classes == 2:
+                predictions = (torch.sigmoid(outputs.logits) > 0.5).float()
+            else:
+                predictions = torch.argmax(outputs.logits, dim=1)
+            correct += (predictions == labels).sum().item()
+            total += labels.size(0)
+            running_loss += loss.item()
+
+            # Log sample predictions for visualization
+            if len(samples) < cfg.eval.num_samples_to_log:
+                for i in range(min(cfg.eval.num_samples_to_log - len(samples), len(images))):
+                    if model.config.num_classes == 2:
+                        confidence = round(outputs.probabilities[i].item(), 3)
+                    else:
+                        confidence = [round(prob, 3) for prob in outputs.probabilities[i].tolist()]
+                    samples.append(
+                        {
+                            **{
+                                f"image_{img_key}": wandb.Image(images[img_idx][i].cpu())
+                                for img_idx, img_key in enumerate(cfg.training.image_keys)
+                            },
+                            "true_label": labels[i].item(),
+                            "predicted": predictions[i].item(),
+                            "confidence": confidence,
+                        }
+                    )
+
+    accuracy = 100 * correct / total
+    avg_loss = running_loss / len(val_loader)
+    print(f"Average validation loss {avg_loss}, and accuracy {accuracy}")
+
+    eval_info = {
+        "loss": avg_loss,
+        "accuracy": accuracy,
+        "eval_s": time.perf_counter() - batch_start_time,
+        "eval/prediction_samples": wandb.Table(
+            data=[list(s.values()) for s in samples],
+            columns=list(samples[0].keys()),
+        )
+        if logger._cfg.wandb.enable
+        else None,
+    }
+
+    if len(inference_times) > 0:
+        eval_info["inference_time_avg"] = np.mean(inference_times)
+        eval_info["inference_time_median"] = np.median(inference_times)
+        eval_info["inference_time_std"] = np.std(inference_times)
+        eval_info["inference_time_batch_size"] = val_loader.batch_size
+
+        print(
+            f"Inference mean time: {eval_info['inference_time_avg']:.2f} us, median: {eval_info['inference_time_median']:.2f} us, std: {eval_info['inference_time_std']:.2f} us, with {len(inference_times)} iterations on {device.type} device, batch size: {eval_info['inference_time_batch_size']}"
+        )
+
+    return accuracy, eval_info
+
+
+def benchmark_inference_time(model, dataset, logger, cfg, device, step):
+    if not cfg.training.profile_inference_time:
+        return
+
+    iters = cfg.training.profile_inference_time_iters
+    inference_times = []
+
+    loader = DataLoader(
+        dataset,
+        batch_size=1,
+        num_workers=cfg.training.num_workers,
+        sampler=RandomSampler(dataset),
+        pin_memory=True,
+    )
+
+    model.eval()
+    with torch.no_grad():
+        for _ in tqdm(range(iters), desc="Benchmarking inference time"):
+            x = next(iter(loader))
+            x = [x[img_key].to(device) for img_key in cfg.training.image_keys]
+
+            # Warm up
+            for _ in range(10):
+                _ = model(x)
+
+            # sync the device
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            elif device.type == "mps":
+                torch.mps.synchronize()
+
+            with (
+                profiler.profile(record_shapes=True) as prof,
+                profiler.record_function("model_inference"),
+            ):
+                _ = model(x)
+
+            inference_times.append(
+                next(x for x in prof.key_averages() if x.key == "model_inference").cpu_time
+            )
+
+    inference_times = np.array(inference_times)
+    avg, median, std = (
+        inference_times.mean(),
+        np.median(inference_times),
+        inference_times.std(),
+    )
+    print(
+        f"Inference time mean: {avg:.2f} us, median: {median:.2f} us, std: {std:.2f} us, with {iters} iterations on {device.type} device"
+    )
+    if logger._cfg.wandb.enable:
+        logger.log_dict(
+            {
+                "inference_time_benchmark_avg": avg,
+                "inference_time_benchmark_median": median,
+                "inference_time_benchmark_std": std,
+            },
+            step + 1,
+            mode="eval",
+        )
+
+    return avg, median, std
+
+
+def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = None) -> None:
+    if out_dir is None:
+        raise NotImplementedError()
+    if job_name is None:
+        raise NotImplementedError()
+
+    # Main training pipeline with support for resuming training
+    init_logging()
+    logging.info(OmegaConf.to_yaml(cfg))
+
+    # Initialize training environment
+    device = get_safe_torch_device(cfg.device, log=True)
+    set_global_seed(cfg.seed)
+
+    # Setup dataset and dataloaders
+    dataset = LeRobotDataset(
+        cfg.dataset_repo_id,
+        root=cfg.dataset_root,
+        local_files_only=cfg.local_files_only,
+    )
+    logging.info(f"Dataset size: {len(dataset)}")
+
+    n_total = len(dataset)
+    n_train = int(cfg.train_split_proportion * len(dataset))
+    train_dataset = torch.utils.data.Subset(dataset, range(0, n_train))
+    val_dataset = torch.utils.data.Subset(dataset, range(n_train, n_total))
+
+    sampler = create_balanced_sampler(train_dataset, cfg)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=cfg.training.batch_size,
+        num_workers=cfg.training.num_workers,
+        sampler=sampler,
+        pin_memory=device.type == "cuda",
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=cfg.eval.batch_size,
+        shuffle=False,
+        num_workers=cfg.training.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+
+    # Resume training if requested
+    step = 0
+    best_val_acc = 0
+
+    if cfg.resume:
+        if not Logger.get_last_checkpoint_dir(out_dir).exists():
+            raise RuntimeError(
+                "You have set resume=True, but there is no model checkpoint in "
+                f"{Logger.get_last_checkpoint_dir(out_dir)}"
+            )
+        checkpoint_cfg_path = str(Logger.get_last_pretrained_model_dir(out_dir) / "config.yaml")
+        logging.info(
+            colored(
+                "You have set resume=True, indicating that you wish to resume a run",
+                color="yellow",
+                attrs=["bold"],
+            )
+        )
+        # Load and validate checkpoint configuration
+        checkpoint_cfg = init_hydra_config(checkpoint_cfg_path)
+        # Check for differences between the checkpoint configuration and provided configuration.
+        # Hack to resolve the delta_timestamps ahead of time in order to properly diff.
+        resolve_delta_timestamps(cfg)
+        diff = DeepDiff(OmegaConf.to_container(checkpoint_cfg), OmegaConf.to_container(cfg))
+        # Ignore the `resume` and parameters.
+        if "values_changed" in diff and "root['resume']" in diff["values_changed"]:
+            del diff["values_changed"]["root['resume']"]
+        if len(diff) > 0:
+            logging.warning(
+                "At least one difference was detected between the checkpoint configuration and "
+                f"the provided configuration: \n{pformat(diff)}\nNote that the checkpoint configuration "
+                "takes precedence.",
+            )
+        # Use the checkpoint config instead of the provided config (but keep `resume` parameter).
+        cfg = checkpoint_cfg
+        cfg.resume = True
+
+    # Initialize model and training components
+    model = get_model(cfg=cfg, logger=logger).to(device)
+
+    optimizer = optim.AdamW(model.parameters(), lr=cfg.training.learning_rate)
+    # Use BCEWithLogitsLoss for binary classification and CrossEntropyLoss for multi-class
+    criterion = nn.BCEWithLogitsLoss() if model.config.num_classes == 2 else nn.CrossEntropyLoss()
+    grad_scaler = GradScaler(enabled=cfg.training.use_amp)
+
+    # Log model parameters
+    num_learnable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    num_total_params = sum(p.numel() for p in model.parameters())
+    logging.info(f"Learnable parameters: {format_big_number(num_learnable_params)}")
+    logging.info(f"Total parameters: {format_big_number(num_total_params)}")
+
+    if cfg.resume:
+        step = logger.load_last_training_state(optimizer, None)
+
+    # Training loop with validation and checkpointing
+    for epoch in range(cfg.training.num_epochs):
+        logging.info(f"\nEpoch {epoch + 1}/{cfg.training.num_epochs}")
+
+        train_epoch(
+            model,
+            train_loader,
+            criterion,
+            optimizer,
+            grad_scaler,
+            device,
+            logger,
+            step,
+            cfg,
+        )
+
+        # Periodic validation
+        if cfg.training.eval_freq > 0 and (epoch + 1) % cfg.training.eval_freq == 0:
+            val_acc, eval_info = validate(
+                model,
+                val_loader,
+                criterion,
+                device,
+                logger,
+                cfg,
+            )
+            logger.log_dict(eval_info, step + len(train_loader), mode="eval")
+
+            # Save best model
+            if val_acc > best_val_acc:
+                best_val_acc = val_acc
+                logger.save_checkpoint(
+                    train_step=step + len(train_loader),
+                    policy=model,
+                    optimizer=optimizer,
+                    scheduler=None,
+                    identifier="best",
+                )
+
+        # Periodic checkpointing
+        if cfg.training.save_checkpoint and (epoch + 1) % cfg.training.save_freq == 0:
+            logger.save_checkpoint(
+                train_step=step + len(train_loader),
+                policy=model,
+                optimizer=optimizer,
+                scheduler=None,
+                identifier=f"{epoch + 1:06d}",
+            )
+
+        step += len(train_loader)
+
+    benchmark_inference_time(model, dataset, logger, cfg, device, step)
+
+    logging.info("Training completed")
+
+
+@hydra.main(
+    version_base="1.2",
+    config_name="hilserl_classifier",
+    config_path="../configs/policy",
+)
+def train_cli(cfg: dict):
+    train(
+        cfg,
+        out_dir=hydra.core.hydra_config.HydraConfig.get().run.dir,
+        job_name=hydra.core.hydra_config.HydraConfig.get().job.name,
+    )
+
+
+def train_notebook(
+    out_dir=None,
+    job_name=None,
+    config_name="hilserl_classifier",
+    config_path="../configs/policy",
+):
+    from hydra import compose, initialize
+
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    initialize(config_path=config_path)
+    cfg = compose(config_name=config_name)
+    train(cfg, out_dir=out_dir, job_name=job_name)
+
+
+if __name__ == "__main__":
+    train_cli()

lerobot/scripts/train_sac.py

@@ -0,0 +1,594 @@
+#!/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.
+import functools
+import logging
+import random
+from pprint import pformat
+from typing import Callable, Optional, Sequence, TypedDict
+
+import hydra
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from tqdm import tqdm
+
+# TODO: Remove the import of maniskill
+from lerobot.common.datasets.factory import make_dataset
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.envs.factory import make_maniskill_env
+from lerobot.common.envs.utils import preprocess_maniskill_observation
+from lerobot.common.logger import Logger, log_output_dir
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies.sac.modeling_sac import SACPolicy
+from lerobot.common.utils.utils import (
+    format_big_number,
+    get_safe_torch_device,
+    init_logging,
+    set_global_seed,
+)
+
+
+def make_optimizers_and_scheduler(cfg, policy):
+    optimizer_actor = torch.optim.Adam(
+        # NOTE: Handle the case of shared encoder where the encoder weights are not optimized with the gradient of the actor
+        params=policy.actor.parameters_to_optimize,
+        lr=policy.config.actor_lr,
+    )
+    optimizer_critic = torch.optim.Adam(
+        params=policy.critic_ensemble.parameters(), lr=policy.config.critic_lr
+    )
+    # We wrap policy log temperature in list because this is a torch tensor and not a nn.Module
+    optimizer_temperature = torch.optim.Adam(params=[policy.log_alpha], lr=policy.config.critic_lr)
+    lr_scheduler = None
+    optimizers = {
+        "actor": optimizer_actor,
+        "critic": optimizer_critic,
+        "temperature": optimizer_temperature,
+    }
+    return optimizers, lr_scheduler
+
+
+class Transition(TypedDict):
+    state: dict[str, torch.Tensor]
+    action: torch.Tensor
+    reward: float
+    next_state: dict[str, torch.Tensor]
+    done: bool
+    complementary_info: dict[str, torch.Tensor] = None
+
+
+class BatchTransition(TypedDict):
+    state: dict[str, torch.Tensor]
+    action: torch.Tensor
+    reward: torch.Tensor
+    next_state: dict[str, torch.Tensor]
+    done: torch.Tensor
+
+
+def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
+    """
+    Perform a per-image random crop over a batch of images in a vectorized way.
+    (Same as shown previously.)
+    """
+    B, C, H, W = images.shape
+    crop_h, crop_w = output_size
+
+    if crop_h > H or crop_w > W:
+        raise ValueError(
+            f"Requested crop size ({crop_h}, {crop_w}) is bigger than the image size ({H}, {W})."
+        )
+
+    tops = torch.randint(0, H - crop_h + 1, (B,), device=images.device)
+    lefts = torch.randint(0, W - crop_w + 1, (B,), device=images.device)
+
+    rows = torch.arange(crop_h, device=images.device).unsqueeze(0) + tops.unsqueeze(1)
+    cols = torch.arange(crop_w, device=images.device).unsqueeze(0) + lefts.unsqueeze(1)
+
+    rows = rows.unsqueeze(2).expand(-1, -1, crop_w)  # (B, crop_h, crop_w)
+    cols = cols.unsqueeze(1).expand(-1, crop_h, -1)  # (B, crop_h, crop_w)
+
+    images_hwcn = images.permute(0, 2, 3, 1)  # (B, H, W, C)
+
+    # Gather pixels
+    cropped_hwcn = images_hwcn[torch.arange(B, device=images.device).view(B, 1, 1), rows, cols, :]
+    # cropped_hwcn => (B, crop_h, crop_w, C)
+
+    cropped = cropped_hwcn.permute(0, 3, 1, 2)  # (B, C, crop_h, crop_w)
+    return cropped
+
+
+def random_shift(images: torch.Tensor, pad: int = 4):
+    """Vectorized random shift, imgs: (B,C,H,W), pad: #pixels"""
+    _, _, h, w = images.shape
+    images = F.pad(input=images, pad=(pad, pad, pad, pad), mode="replicate")
+    return random_crop_vectorized(images=images, output_size=(h, w))
+
+
+class ReplayBuffer:
+    def __init__(
+        self,
+        capacity: int,
+        device: str = "cuda:0",
+        state_keys: Optional[Sequence[str]] = None,
+        image_augmentation_function: Optional[Callable] = None,
+        use_drq: bool = True,
+    ):
+        """
+        Args:
+            capacity (int): Maximum number of transitions to store in the buffer.
+            device (str): The device where the tensors will be moved ("cuda:0" or "cpu").
+            state_keys (List[str]): The list of keys that appear in `state` and `next_state`.
+            image_augmentation_function (Optional[Callable]): A function that takes a batch of images
+                and returns a batch of augmented images. If None, a default augmentation function is used.
+            use_drq (bool): Whether to use the default DRQ image augmentation style, when sampling in the buffer.
+        """
+        self.capacity = capacity
+        self.device = device
+        self.memory: list[Transition] = []
+        self.position = 0
+
+        # If no state_keys provided, default to an empty list
+        # (you can handle this differently if needed)
+        self.state_keys = state_keys if state_keys is not None else []
+        if image_augmentation_function is None:
+            self.image_augmentation_function = functools.partial(random_shift, pad=4)
+        self.use_drq = use_drq
+
+    def add(
+        self,
+        state: dict[str, torch.Tensor],
+        action: torch.Tensor,
+        reward: float,
+        next_state: dict[str, torch.Tensor],
+        done: bool,
+        complementary_info: Optional[dict[str, torch.Tensor]] = None,
+    ):
+        """Saves a transition."""
+        if len(self.memory) < self.capacity:
+            self.memory.append(None)
+
+        # Create and store the Transition
+        self.memory[self.position] = Transition(
+            state=state,
+            action=action,
+            reward=reward,
+            next_state=next_state,
+            done=done,
+            complementary_info=complementary_info,
+        )
+        self.position: int = (self.position + 1) % self.capacity
+
+    # TODO: ADD image_augmentation and use_drq arguments in this function in order to instantiate the class with them
+    @classmethod
+    def from_lerobot_dataset(
+        cls,
+        lerobot_dataset: LeRobotDataset,
+        device: str = "cuda:0",
+        state_keys: Optional[Sequence[str]] = None,
+    ) -> "ReplayBuffer":
+        """
+        Convert a LeRobotDataset into a ReplayBuffer.
+
+        Args:
+            lerobot_dataset (LeRobotDataset): The dataset to convert.
+            device (str): The device . Defaults to "cuda:0".
+            state_keys (Optional[Sequence[str]], optional): The list of keys that appear in `state` and `next_state`.
+            Defaults to None.
+
+        Returns:
+            ReplayBuffer: The replay buffer with offline dataset transitions.
+        """
+        # We convert the LeRobotDataset into a replay buffer, because it is more efficient to sample from
+        # a replay buffer than from a lerobot dataset.
+        replay_buffer = cls(capacity=len(lerobot_dataset), device=device, state_keys=state_keys)
+        list_transition = cls._lerobotdataset_to_transitions(dataset=lerobot_dataset, state_keys=state_keys)
+        # Fill the replay buffer with the lerobot dataset transitions
+        for data in list_transition:
+            replay_buffer.add(
+                state=data["state"],
+                action=data["action"],
+                reward=data["reward"],
+                next_state=data["next_state"],
+                done=data["done"],
+            )
+        return replay_buffer
+
+    @staticmethod
+    def _lerobotdataset_to_transitions(
+        dataset: LeRobotDataset,
+        state_keys: Optional[Sequence[str]] = None,
+    ) -> list[Transition]:
+        """
+        Convert a LeRobotDataset into a list of RL (s, a, r, s', done) transitions.
+
+        Args:
+            dataset (LeRobotDataset):
+                The dataset to convert. Each item in the dataset is expected to have
+                at least the following keys:
+                {
+                    "action": ...
+                    "next.reward": ...
+                    "next.done": ...
+                    "episode_index": ...
+                }
+                plus whatever your 'state_keys' specify.
+
+            state_keys (Optional[Sequence[str]]):
+                The dataset keys to include in 'state' and 'next_state'. Their names
+                will be kept as-is in the output transitions. E.g.
+                ["observation.state", "observation.environment_state"].
+                If None, you must handle or define default keys.
+
+        Returns:
+            transitions (List[Transition]):
+                A list of Transition dictionaries with the same length as `dataset`.
+        """
+
+        # If not provided, you can either raise an error or define a default:
+        if state_keys is None:
+            raise ValueError("You must provide a list of keys in `state_keys` that define your 'state'.")
+
+        transitions: list[Transition] = []
+        num_frames = len(dataset)
+
+        for i in tqdm(range(num_frames)):
+            current_sample = dataset[i]
+
+            # ----- 1) Current state -----
+            current_state: dict[str, torch.Tensor] = {}
+            for key in state_keys:
+                val = current_sample[key]
+                current_state[key] = val.unsqueeze(0)  # Add batch dimension
+
+            # ----- 2) Action -----
+            action = current_sample["action"].unsqueeze(0)  # Add batch dimension
+
+            # ----- 3) Reward and done -----
+            reward = float(current_sample["next.reward"].item())  # ensure float
+            done = bool(current_sample["next.done"].item())  # ensure bool
+
+            # ----- 4) Next state -----
+            # If not done and the next sample is in the same episode, we pull the next sample's state.
+            # Otherwise (done=True or next sample crosses to a new episode), next_state = current_state.
+            next_state = current_state  # default
+            if not done and (i < num_frames - 1):
+                next_sample = dataset[i + 1]
+                if next_sample["episode_index"] == current_sample["episode_index"]:
+                    # Build next_state from the same keys
+                    next_state_data: dict[str, torch.Tensor] = {}
+                    for key in state_keys:
+                        val = next_sample[key]
+                        next_state_data[key] = val.unsqueeze(0)  # Add batch dimension
+                    next_state = next_state_data
+
+            # ----- Construct the Transition -----
+            transition = Transition(
+                state=current_state,
+                action=action,
+                reward=reward,
+                next_state=next_state,
+                done=done,
+            )
+            transitions.append(transition)
+
+        return transitions
+
+    def sample(self, batch_size: int) -> BatchTransition:
+        """Sample a random batch of transitions and collate them into batched tensors."""
+        list_of_transitions = random.sample(self.memory, batch_size)
+
+        # -- Build batched states --
+        batch_state = {}
+        for key in self.state_keys:
+            batch_state[key] = torch.cat([t["state"][key] for t in list_of_transitions], dim=0).to(
+                self.device
+            )
+            if key.startswith("observation.image") and self.use_drq:
+                batch_state[key] = self.image_augmentation_function(batch_state[key])
+
+        # -- Build batched actions --
+        batch_actions = torch.cat([t["action"] for t in list_of_transitions]).to(self.device)
+
+        # -- Build batched rewards --
+        batch_rewards = torch.tensor([t["reward"] for t in list_of_transitions], dtype=torch.float32).to(
+            self.device
+        )
+
+        # -- Build batched next states --
+        batch_next_state = {}
+        for key in self.state_keys:
+            batch_next_state[key] = torch.cat([t["next_state"][key] for t in list_of_transitions], dim=0).to(
+                self.device
+            )
+            if key.startswith("observation.image") and self.use_drq:
+                batch_next_state[key] = self.image_augmentation_function(batch_next_state[key])
+
+        # -- Build batched dones --
+        batch_dones = torch.tensor([t["done"] for t in list_of_transitions], dtype=torch.float32).to(
+            self.device
+        )
+        batch_dones = torch.tensor([t["done"] for t in list_of_transitions], dtype=torch.float32).to(
+            self.device
+        )
+
+        # Return a BatchTransition typed dict
+        return BatchTransition(
+            state=batch_state,
+            action=batch_actions,
+            reward=batch_rewards,
+            next_state=batch_next_state,
+            done=batch_dones,
+        )
+
+
+def concatenate_batch_transitions(
+    left_batch_transitions: BatchTransition, right_batch_transition: BatchTransition
+) -> BatchTransition:
+    """NOTE: Be careful it change the left_batch_transitions in place"""
+    left_batch_transitions["state"] = {
+        key: torch.cat(
+            [
+                left_batch_transitions["state"][key],
+                right_batch_transition["state"][key],
+            ],
+            dim=0,
+        )
+        for key in left_batch_transitions["state"]
+    }
+    left_batch_transitions["action"] = torch.cat(
+        [left_batch_transitions["action"], right_batch_transition["action"]], dim=0
+    )
+    left_batch_transitions["reward"] = torch.cat(
+        [left_batch_transitions["reward"], right_batch_transition["reward"]], dim=0
+    )
+    left_batch_transitions["next_state"] = {
+        key: torch.cat(
+            [
+                left_batch_transitions["next_state"][key],
+                right_batch_transition["next_state"][key],
+            ],
+            dim=0,
+        )
+        for key in left_batch_transitions["next_state"]
+    }
+    left_batch_transitions["done"] = torch.cat(
+        [left_batch_transitions["done"], right_batch_transition["done"]], dim=0
+    )
+    return left_batch_transitions
+
+
+def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = None):
+    if out_dir is None:
+        raise NotImplementedError()
+    if job_name is None:
+        raise NotImplementedError()
+
+    init_logging()
+    logging.info(pformat(OmegaConf.to_container(cfg)))
+
+    # Create an env dedicated to online episodes collection from policy rollout.
+    # online_env = make_env(cfg, n_envs=cfg.training.online_rollout_batch_size)
+    # NOTE: Off policy algorithm are efficient enought to use a single environment
+    logging.info("make_env online")
+    # online_env = make_env(cfg, n_envs=1)
+    # TODO: Remove the import of maniskill and unifiy with make env
+    online_env = make_maniskill_env(cfg, n_envs=1)
+    if cfg.training.eval_freq > 0:
+        logging.info("make_env eval")
+        # eval_env = make_env(cfg, n_envs=1)
+        # TODO: Remove the import of maniskill and unifiy with make env
+        eval_env = make_maniskill_env(cfg, n_envs=1)
+
+    # TODO: Add a way to resume training
+
+    # log metrics to terminal and wandb
+    logger = Logger(cfg, out_dir, wandb_job_name=job_name)
+
+    set_global_seed(cfg.seed)
+
+    # Check device is available
+    device = get_safe_torch_device(cfg.device, log=True)
+
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging.info("make_policy")
+    # TODO: At some point we should just need make sac policy
+    policy: SACPolicy = make_policy(
+        hydra_cfg=cfg,
+        # dataset_stats=offline_dataset.meta.stats if not cfg.resume else None,
+        # Hack: But if we do online traning, we do not need dataset_stats
+        dataset_stats=None,
+        pretrained_policy_name_or_path=str(logger.last_pretrained_model_dir) if cfg.resume else None,
+        device=device,
+    )
+    assert isinstance(policy, nn.Module)
+
+    optimizers, lr_scheduler = make_optimizers_and_scheduler(cfg, policy)
+
+    # TODO: Handle resume
+
+    num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
+    num_total_params = sum(p.numel() for p in policy.parameters())
+
+    log_output_dir(out_dir)
+    logging.info(f"{cfg.env.task=}")
+    logging.info(f"{cfg.training.online_steps=}")
+    logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
+    logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
+
+    obs, info = online_env.reset()
+
+    # HACK for maniskill
+    # obs = preprocess_observation(obs)
+    obs = preprocess_maniskill_observation(obs)
+    obs = {key: obs[key].to(device, non_blocking=True) for key in obs}
+
+    replay_buffer = ReplayBuffer(
+        capacity=cfg.training.online_buffer_capacity,
+        device=device,
+        state_keys=cfg.policy.input_shapes.keys(),
+    )
+
+    batch_size = cfg.training.batch_size
+
+    if cfg.dataset_repo_id is not None:
+        logging.info("make_dataset offline buffer")
+        offline_dataset = make_dataset(cfg)
+        logging.info("Convertion to a offline replay buffer")
+        offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
+            offline_dataset, device=device, state_keys=cfg.policy.input_shapes.keys()
+        )
+        batch_size: int = batch_size // 2  # We will sample from both replay buffer
+
+    # NOTE: For the moment we will solely handle the case of a single environment
+    sum_reward_episode = 0
+
+    for interaction_step in range(cfg.training.online_steps):
+        # NOTE: At some point we should use a  wrapper to handle the observation
+
+        if interaction_step >= cfg.training.online_step_before_learning:
+            action = policy.select_action(batch=obs)
+            next_obs, reward, done, truncated, info = online_env.step(action.cpu().numpy())
+        else:
+            action = online_env.action_space.sample()
+            next_obs, reward, done, truncated, info = online_env.step(action)
+            # HACK
+            action = torch.tensor(action, dtype=torch.float32).to(device, non_blocking=True)
+
+        # HACK: For maniskill
+        # next_obs = preprocess_observation(next_obs)
+        next_obs = preprocess_maniskill_observation(next_obs)
+        next_obs = {key: next_obs[key].to(device, non_blocking=True) for key in obs}
+        sum_reward_episode += float(reward[0])
+        # Because we are using a single environment
+        # we can safely assume that the episode is done
+        if done[0] or truncated[0]:
+            logging.info(f"Global step {interaction_step}: Episode reward: {sum_reward_episode}")
+            logger.log_dict({"Sum episode reward": sum_reward_episode}, interaction_step)
+            sum_reward_episode = 0
+            # HACK: This is for maniskill
+            logging.info(
+                f"global step {interaction_step}: episode success: {info['success'].float().item()} \n"
+            )
+            logger.log_dict({"Episode success": info["success"].float().item()}, interaction_step)
+
+        replay_buffer.add(
+            state=obs,
+            action=action,
+            reward=float(reward[0]),
+            next_state=next_obs,
+            done=done[0],
+        )
+        obs = next_obs
+
+        if interaction_step < cfg.training.online_step_before_learning:
+            continue
+        for _ in range(cfg.policy.utd_ratio - 1):
+            batch = replay_buffer.sample(batch_size)
+            if cfg.dataset_repo_id is not None:
+                batch_offline = offline_replay_buffer.sample(batch_size)
+                batch = concatenate_batch_transitions(batch, batch_offline)
+
+            actions = batch["action"]
+            rewards = batch["reward"]
+            observations = batch["state"]
+            next_observations = batch["next_state"]
+            done = batch["done"]
+
+            loss_critic = policy.compute_loss_critic(
+                observations=observations,
+                actions=actions,
+                rewards=rewards,
+                next_observations=next_observations,
+                done=done,
+            )
+            optimizers["critic"].zero_grad()
+            loss_critic.backward()
+            optimizers["critic"].step()
+
+        batch = replay_buffer.sample(batch_size)
+        if cfg.dataset_repo_id is not None:
+            batch_offline = offline_replay_buffer.sample(batch_size)
+            batch = concatenate_batch_transitions(
+                left_batch_transitions=batch, right_batch_transition=batch_offline
+            )
+
+        actions = batch["action"]
+        rewards = batch["reward"]
+        observations = batch["state"]
+        next_observations = batch["next_state"]
+        done = batch["done"]
+
+        loss_critic = policy.compute_loss_critic(
+            observations=observations,
+            actions=actions,
+            rewards=rewards,
+            next_observations=next_observations,
+            done=done,
+        )
+        optimizers["critic"].zero_grad()
+        loss_critic.backward()
+        optimizers["critic"].step()
+
+        training_infos = {}
+        training_infos["loss_critic"] = loss_critic.item()
+
+        if interaction_step % cfg.training.policy_update_freq == 0:
+            # TD3 Trick
+            for _ in range(cfg.training.policy_update_freq):
+                loss_actor = policy.compute_loss_actor(observations=observations)
+
+                optimizers["actor"].zero_grad()
+                loss_actor.backward()
+                optimizers["actor"].step()
+
+                training_infos["loss_actor"] = loss_actor.item()
+
+                loss_temperature = policy.compute_loss_temperature(observations=observations)
+                optimizers["temperature"].zero_grad()
+                loss_temperature.backward()
+                optimizers["temperature"].step()
+
+                training_infos["loss_temperature"] = loss_temperature.item()
+
+        if interaction_step % cfg.training.log_freq == 0:
+            logger.log_dict(training_infos, interaction_step, mode="train")
+
+        policy.update_target_networks()
+
+
+@hydra.main(version_base="1.2", config_name="default", config_path="../configs")
+def train_cli(cfg: dict):
+    train(
+        cfg,
+        out_dir=hydra.core.hydra_config.HydraConfig.get().run.dir,
+        job_name=hydra.core.hydra_config.HydraConfig.get().job.name,
+    )
+
+
+def train_notebook(out_dir=None, job_name=None, config_name="default", config_path="../configs"):
+    from hydra import compose, initialize
+
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    initialize(config_path=config_path)
+    cfg = compose(config_name=config_name)
+    train(cfg, out_dir=out_dir, job_name=job_name)
+
+
+if __name__ == "__main__":
+    train_cli()

lerobot/scripts/visualize_dataset_html.py

@@ -81,7 +81,11 @@ def run_server(
     static_folder: Path,
     template_folder: Path,
 ):
-    app = Flask(__name__, static_folder=static_folder.resolve(), template_folder=template_folder.resolve())
+    app = Flask(
+        __name__,
+        static_folder=static_folder.resolve(),
+        template_folder=template_folder.resolve(),
+    )
     app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0  # specifying not to cache
 
     @app.route("/")
@@ -175,7 +179,9 @@ def run_server(
             ]
             videos_info = [
                 {
-                    "url": url_for("static", filename=str(video_path).replace("\\", "/")),
+                    "url": url_for(
+                        "static", filename=str(video_path).replace("\\", "/")
+                    ),
                     "filename": video_path.parent.name,
                 }
                 for video_path in video_paths
@@ -197,7 +203,8 @@ def run_server(
             ]
 
             response = requests.get(
-                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
+                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl",
+                timeout=5,
             )
             response.raise_for_status()
             # Split into lines and parse each line as JSON
@@ -282,7 +289,8 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
         repo_id = dataset.repo_id
 
         url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/" + dataset.data_path.format(
-            episode_chunk=int(episode_index) // dataset.chunks_size, episode_index=episode_index
+            episode_chunk=int(episode_index) // dataset.chunks_size,
+            episode_index=episode_index,
         )
         df = pd.read_parquet(url)
         data = df[selected_columns]  # Select specific columns
@@ -331,7 +339,8 @@ def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) ->
 
 def get_dataset_info(repo_id: str) -> IterableNamespace:
     response = requests.get(
-        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
+        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json",
+        timeout=5,
     )
     response.raise_for_status()  # Raises an HTTPError for bad responses
     dataset_info = response.json()

pyproject.toml

@@ -53,6 +53,7 @@ dependencies = [
     "einops>=0.8.0",
     "flask>=3.0.3",
     "gdown>=5.1.0",
+    "grpcio>=1.70.0",
     "gymnasium==0.29.1",                                                 # TODO(rcadene, aliberts): Make gym 1.0.0 work
     "h5py>=3.10.0",
     "huggingface-hub[hf-transfer,cli]>=0.27.1 ; python_version < '4.0'",
@@ -63,6 +64,7 @@ dependencies = [
     "opencv-python-headless>=4.9.0",
     "packaging>=24.2",
     "av>=12.0.5",
+    "protobuf>=5.29.3",
     "pymunk>=6.6.0",
     "pynput>=1.7.7",
     "pyzmq>=26.2.1",
@@ -70,7 +72,9 @@ dependencies = [
     "termcolor>=2.4.0",
     "torch>=2.2.1",
     "torchcodec>=0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
+    "torchmetrics>=1.6.0",
     "torchvision>=0.21.0",
+    "transformers>=4.47.0",
     "wandb>=0.16.3",
     "zarr>=2.17.0",
 ]
@@ -83,7 +87,9 @@ dora = [
 ]
 dynamixel = ["dynamixel-sdk>=3.7.31", "pynput>=1.7.7"]
 feetech = ["feetech-servo-sdk>=1.0.0", "pynput>=1.7.7"]
+hilserl = ["transformers>=4.48.0", "torchmetrics>=1.6.0"]
 intelrealsense = ["pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'"]
+mani_skill = ["mani-skill"]
 pi0 = ["transformers>=4.48.0"]
 pusht = ["gym-pusht>=0.1.5 ; python_version < '4.0'"]
 stretch = [
@@ -97,13 +103,40 @@ umi = ["imagecodecs>=2024.1.1"]
 video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
 xarm = ["gym-xarm>=0.1.1 ; python_version < '4.0'"]
 
+
 [tool.poetry]
 requires-poetry = ">=2.1"
 
 [tool.ruff]
 line-length = 110
 target-version = "py310"
-exclude = ["tests/artifacts/**/*.safetensors"]
+exclude = [
+    "tests/data",
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "venv",
+    "*_pb2.py",
+    "*_pb2_grpc.py",
+]
+
 
 [tool.ruff.lint]
 select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"]

tests/cameras/test_cameras.py

@@ -37,7 +37,10 @@ pytest -sx 'tests/test_cameras.py::test_camera[intelrealsense-True]'
 import numpy as np
 import pytest
 
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
 from tests.utils import TEST_CAMERA_TYPES, make_camera, require_camera
 
 # Maximum absolute difference between two consecutive images recorded by a camera.
@@ -112,7 +115,11 @@ def test_camera(request, camera_type, mock):
     )
     # TODO(rcadene): properly set `rtol`
     np.testing.assert_allclose(
-        color_image, async_color_image, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
+        color_image,
+        async_color_image,
+        rtol=1e-5,
+        atol=MAX_PIXEL_DIFFERENCE,
+        err_msg=error_msg,
     )
 
     # Test disconnecting
@@ -131,7 +138,11 @@ def test_camera(request, camera_type, mock):
     assert camera.color_mode == "bgr"
     bgr_color_image = camera.read()
     np.testing.assert_allclose(
-        color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
+        color_image,
+        bgr_color_image[:, :, [2, 1, 0]],
+        rtol=1e-5,
+        atol=MAX_PIXEL_DIFFERENCE,
+        err_msg=error_msg,
     )
     del camera
 
@@ -166,7 +177,11 @@ def test_camera(request, camera_type, mock):
         rot_color_image = camera.read()
 
         np.testing.assert_allclose(
-            rot_color_image, manual_rot_img, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
+            rot_color_image,
+            manual_rot_img,
+            rtol=1e-5,
+            atol=MAX_PIXEL_DIFFERENCE,
+            err_msg=error_msg,
         )
         del camera
 
@@ -200,7 +215,9 @@ def test_save_images_from_cameras(tmp_path, request, camera_type, mock):
     if camera_type == "opencv":
         from lerobot.common.robot_devices.cameras.opencv import save_images_from_cameras
     elif camera_type == "intelrealsense":
-        from lerobot.common.robot_devices.cameras.intelrealsense import save_images_from_cameras
+        from lerobot.common.robot_devices.cameras.intelrealsense import (
+            save_images_from_cameras,
+        )
 
     # Small `record_time_s` to speedup unit tests
     save_images_from_cameras(tmp_path, record_time_s=0.02, mock=mock)
@@ -209,7 +226,13 @@ def test_save_images_from_cameras(tmp_path, request, camera_type, mock):
 @pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
 @require_camera
 def test_camera_rotation(request, camera_type, mock):
-    config_kwargs = {"camera_type": camera_type, "mock": mock, "width": 640, "height": 480, "fps": 30}
+    config_kwargs = {
+        "camera_type": camera_type,
+        "mock": mock,
+        "width": 640,
+        "height": 480,
+        "fps": 30,
+    }
 
     # No rotation.
     camera = make_camera(**config_kwargs, rotation=None)

tests/conftest.py

@@ -14,9 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
 import traceback
 
 import pytest
+import torch
 from serial import SerialException
 
 from lerobot import available_cameras, available_motors, available_robots
@@ -86,3 +88,19 @@ def patch_builtins_input(monkeypatch):
             print(text)
 
     monkeypatch.setattr("builtins.input", print_text)
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--seed",
+        action="store",
+        default="42",
+        help="Set random seed for reproducibility",
+    )
+
+
+@pytest.fixture(autouse=True)
+def set_random_seed(request):
+    seed = int(request.config.getoption("--seed"))
+    random.seed(seed)  # Python random
+    torch.manual_seed(seed)  # PyTorch

tests/datasets/test_compute_stats.py

@@ -61,7 +61,10 @@ def test_sample_indices():
     assert len(indices) == estimate_num_samples(10)
 
 
-@patch("lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
+@patch(
+    "lerobot.common.datasets.compute_stats.load_image_as_numpy",
+    side_effect=mock_load_image_as_numpy,
+)
 def test_sample_images(mock_load):
     image_paths = [f"image_{i}.jpg" for i in range(100)]
     images = sample_images(image_paths)
@@ -145,7 +148,8 @@ def test_compute_episode_stats():
     }
 
     with patch(
-        "lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy
+        "lerobot.common.datasets.compute_stats.load_image_as_numpy",
+        side_effect=mock_load_image_as_numpy,
     ):
         stats = compute_episode_stats(episode_data, features)
 
@@ -233,7 +237,13 @@ def test_aggregate_stats():
                 "std": [2.87, 5.87, 8.87],
                 "count": 10,
             },
-            "observation.state": {"min": 1, "max": 10, "mean": 5.5, "std": 2.87, "count": 10},
+            "observation.state": {
+                "min": 1,
+                "max": 10,
+                "mean": 5.5,
+                "std": 2.87,
+                "count": 10,
+            },
             "extra_key_0": {"min": 5, "max": 25, "mean": 15, "std": 6, "count": 6},
         },
         {
@@ -244,7 +254,13 @@ def test_aggregate_stats():
                 "std": [3.42, 2.42, 1.42],
                 "count": 15,
             },
-            "observation.state": {"min": 2, "max": 15, "mean": 8.5, "std": 3.42, "count": 15},
+            "observation.state": {
+                "min": 2,
+                "max": 15,
+                "mean": 8.5,
+                "std": 3.42,
+                "count": 15,
+            },
             "extra_key_1": {"min": 0, "max": 20, "mean": 10, "std": 5, "count": 5},
         },
     ]
@@ -306,6 +322,9 @@ def test_aggregate_stats():
         np.testing.assert_allclose(results[fkey]["max"], expected_agg_stats[fkey]["max"])
         np.testing.assert_allclose(results[fkey]["mean"], expected_agg_stats[fkey]["mean"])
         np.testing.assert_allclose(
-            results[fkey]["std"], expected_agg_stats[fkey]["std"], atol=1e-04, rtol=1e-04
+            results[fkey]["std"],
+            expected_agg_stats[fkey]["std"],
+            atol=1e-04,
+            rtol=1e-04,
         )
         np.testing.assert_allclose(results[fkey]["count"], expected_agg_stats[fkey]["count"])

tests/datasets/test_datasets.py

@@ -104,7 +104,8 @@ def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
     with pytest.raises(
-        ValueError, match="Feature mismatch in `frame` dictionary:\nMissing features: {'task'}\n"
+        ValueError,
+        match="Feature mismatch in `frame` dictionary:\nMissing features: {'task'}\n",
     ):
         dataset.add_frame({"state": torch.randn(1)})
 
@@ -113,7 +114,8 @@ def test_add_frame_missing_feature(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
     with pytest.raises(
-        ValueError, match="Feature mismatch in `frame` dictionary:\nMissing features: {'state'}\n"
+        ValueError,
+        match="Feature mismatch in `frame` dictionary:\nMissing features: {'state'}\n",
     ):
         dataset.add_frame({"task": "Dummy task"})
 
@@ -122,7 +124,8 @@ def test_add_frame_extra_feature(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
     with pytest.raises(
-        ValueError, match="Feature mismatch in `frame` dictionary:\nExtra features: {'extra'}\n"
+        ValueError,
+        match="Feature mismatch in `frame` dictionary:\nExtra features: {'extra'}\n",
     ):
         dataset.add_frame({"state": torch.randn(1), "task": "Dummy task", "extra": "dummy_extra"})
 
@@ -131,7 +134,8 @@ def test_add_frame_wrong_type(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
     with pytest.raises(
-        ValueError, match="The feature 'state' of dtype 'float16' is not of the expected dtype 'float32'.\n"
+        ValueError,
+        match="The feature 'state' of dtype 'float16' is not of the expected dtype 'float32'.\n",
     ):
         dataset.add_frame({"state": torch.randn(1, dtype=torch.float16), "task": "Dummy task"})
 

tests/datasets/test_delta_timestamps.py

@@ -54,7 +54,9 @@ def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> dict[str, np.n
 
 @pytest.fixture(scope="module")
 def synced_timestamps_factory(hf_dataset_factory):
-    def _create_synced_timestamps(fps: int = 30) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def _create_synced_timestamps(
+        fps: int = 30,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         hf_dataset = hf_dataset_factory(fps=fps)
         timestamps = torch.stack(hf_dataset["timestamp"]).numpy()
         episode_indices = torch.stack(hf_dataset["episode_index"]).numpy()
@@ -91,7 +93,9 @@ def slightly_off_timestamps_factory(synced_timestamps_factory):
 @pytest.fixture(scope="module")
 def valid_delta_timestamps_factory():
     def _create_valid_delta_timestamps(
-        fps: int = 30, keys: list = DUMMY_MOTOR_FEATURES, min_max_range: tuple[int, int] = (-10, 10)
+        fps: int = 30,
+        keys: list = DUMMY_MOTOR_FEATURES,
+        min_max_range: tuple[int, int] = (-10, 10),
     ) -> dict:
         delta_timestamps = {key: [i * (1 / fps) for i in range(*min_max_range)] for key in keys}
         return delta_timestamps

tests/datasets/test_image_transforms.py

@@ -33,7 +33,9 @@ from lerobot.scripts.visualize_image_transforms import (
     save_all_transforms,
     save_each_transform,
 )
-from tests.artifacts.image_transforms.save_image_transforms_to_safetensors import ARTIFACT_DIR
+from tests.artifacts.image_transforms.save_image_transforms_to_safetensors import (
+    ARTIFACT_DIR,
+)
 from tests.utils import require_x86_64_kernel
 
 
@@ -91,7 +93,8 @@ def test_get_image_transforms_brightness(img_tensor_factory, min_max):
 def test_get_image_transforms_contrast(img_tensor_factory, min_max):
     img_tensor = img_tensor_factory()
     tf_cfg = ImageTransformsConfig(
-        enable=True, tfs={"contrast": ImageTransformConfig(type="ColorJitter", kwargs={"contrast": min_max})}
+        enable=True,
+        tfs={"contrast": ImageTransformConfig(type="ColorJitter", kwargs={"contrast": min_max})},
     )
     tf_actual = ImageTransforms(tf_cfg)
     tf_expected = v2.ColorJitter(contrast=min_max)
@@ -114,7 +117,8 @@ def test_get_image_transforms_saturation(img_tensor_factory, min_max):
 def test_get_image_transforms_hue(img_tensor_factory, min_max):
     img_tensor = img_tensor_factory()
     tf_cfg = ImageTransformsConfig(
-        enable=True, tfs={"hue": ImageTransformConfig(type="ColorJitter", kwargs={"hue": min_max})}
+        enable=True,
+        tfs={"hue": ImageTransformConfig(type="ColorJitter", kwargs={"hue": min_max})},
     )
     tf_actual = ImageTransforms(tf_cfg)
     tf_expected = v2.ColorJitter(hue=min_max)
@@ -367,7 +371,11 @@ def test_save_each_transform(img_tensor_factory, tmp_path):
         assert any(transform_dir.iterdir()), f"No transformed images found in {transform} directory."
 
         # Check for specific files within each transform directory
-        expected_files = [f"{i}.png" for i in range(1, n_examples + 1)] + ["min.png", "max.png", "mean.png"]
+        expected_files = [f"{i}.png" for i in range(1, n_examples + 1)] + [
+            "min.png",
+            "max.png",
+            "mean.png",
+        ]
         for file_name in expected_files:
             assert (transform_dir / file_name).exists(), (
                 f"{file_name} was not found in {transform} directory."

tests/datasets/test_online_buffer.py

@@ -227,7 +227,9 @@ def test_compute_sampler_weights_trivial(
         )
 
     weights = compute_sampler_weights(
-        offline_dataset, online_dataset=online_dataset, online_sampling_ratio=online_sampling_ratio
+        offline_dataset,
+        online_dataset=online_dataset,
+        online_sampling_ratio=online_sampling_ratio,
     )
     if offline_dataset_size == 0 or online_dataset_size == 0:
         expected_weights = torch.ones(offline_dataset_size + online_dataset_size)
@@ -246,10 +248,13 @@ def test_compute_sampler_weights_nontrivial_ratio(lerobot_dataset_factory, tmp_p
     online_dataset.add_data(make_spoof_data_frames(n_episodes=4, n_frames_per_episode=2))
     online_sampling_ratio = 0.8
     weights = compute_sampler_weights(
-        offline_dataset, online_dataset=online_dataset, online_sampling_ratio=online_sampling_ratio
+        offline_dataset,
+        online_dataset=online_dataset,
+        online_sampling_ratio=online_sampling_ratio,
     )
     torch.testing.assert_close(
-        weights, torch.tensor([0.05, 0.05, 0.05, 0.05, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
+        weights,
+        torch.tensor([0.05, 0.05, 0.05, 0.05, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]),
     )
 
 
@@ -259,10 +264,14 @@ def test_compute_sampler_weights_nontrivial_ratio_and_drop_last_n(lerobot_datase
     online_dataset, _ = make_new_buffer()
     online_dataset.add_data(make_spoof_data_frames(n_episodes=4, n_frames_per_episode=2))
     weights = compute_sampler_weights(
-        offline_dataset, online_dataset=online_dataset, online_sampling_ratio=0.8, online_drop_n_last_frames=1
+        offline_dataset,
+        online_dataset=online_dataset,
+        online_sampling_ratio=0.8,
+        online_drop_n_last_frames=1,
     )
     torch.testing.assert_close(
-        weights, torch.tensor([0.05, 0.05, 0.05, 0.05, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0])
+        weights,
+        torch.tensor([0.05, 0.05, 0.05, 0.05, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0]),
     )
 
 

tests/datasets/test_sampler.py

@@ -15,7 +15,9 @@
 # limitations under the License.
 from datasets import Dataset
 
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
+from lerobot.common.datasets.push_dataset_to_hub.utils import (
+    calculate_episode_data_index,
+)
 from lerobot.common.datasets.sampler import EpisodeAwareSampler
 from lerobot.common.datasets.utils import (
     hf_transform_to_torch,

tests/datasets/test_utils.py

@@ -18,8 +18,13 @@ import torch
 from datasets import Dataset
 from huggingface_hub import DatasetCard
 
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch
+from lerobot.common.datasets.push_dataset_to_hub.utils import (
+    calculate_episode_data_index,
+)
+from lerobot.common.datasets.utils import (
+    create_lerobot_dataset_card,
+    hf_transform_to_torch,
+)
 
 
 def test_default_parameters():

tests/fixtures/constants.py

@@ -20,17 +20,39 @@ DUMMY_MOTOR_FEATURES = {
     "action": {
         "dtype": "float32",
         "shape": (6,),
-        "names": ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"],
+        "names": [
+            "shoulder_pan",
+            "shoulder_lift",
+            "elbow_flex",
+            "wrist_flex",
+            "wrist_roll",
+            "gripper",
+        ],
     },
     "state": {
         "dtype": "float32",
         "shape": (6,),
-        "names": ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"],
+        "names": [
+            "shoulder_pan",
+            "shoulder_lift",
+            "elbow_flex",
+            "wrist_flex",
+            "wrist_roll",
+            "gripper",
+        ],
     },
 }
 DUMMY_CAMERA_FEATURES = {
-    "laptop": {"shape": (480, 640, 3), "names": ["height", "width", "channels"], "info": None},
-    "phone": {"shape": (480, 640, 3), "names": ["height", "width", "channels"], "info": None},
+    "laptop": {
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channels"],
+        "info": None,
+    },
+    "phone": {
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channels"],
+        "info": None,
+    },
 }
 DEFAULT_FPS = 30
 DUMMY_VIDEO_INFO = {

tests/fixtures/dataset_factories.py

@@ -23,7 +23,11 @@ import PIL.Image
 import pytest
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import (
+    CODEBASE_VERSION,
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+)
 from lerobot.common.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_FEATURES,
@@ -197,7 +201,10 @@ def tasks_factory():
     def _create_tasks(total_tasks: int = 3) -> int:
         tasks = {}
         for task_index in range(total_tasks):
-            task_dict = {"task_index": task_index, "task": f"Perform action {task_index}."}
+            task_dict = {
+                "task_index": task_index,
+                "task": f"Perform action {task_index}.",
+            }
             tasks[task_index] = task_dict
         return tasks
 
@@ -275,7 +282,10 @@ def hf_dataset_factory(features_factory, tasks_factory, episodes_factory, img_ar
             timestamp_col = np.concatenate((timestamp_col, np.arange(ep_dict["length"]) / fps))
             frame_index_col = np.concatenate((frame_index_col, np.arange(ep_dict["length"], dtype=int)))
             episode_index_col = np.concatenate(
-                (episode_index_col, np.full(ep_dict["length"], ep_dict["episode_index"], dtype=int))
+                (
+                    episode_index_col,
+                    np.full(ep_dict["length"], ep_dict["episode_index"], dtype=int),
+                )
             )
             ep_task_index = get_task_index(tasks, ep_dict["tasks"][0])
             task_index = np.concatenate((task_index, np.full(ep_dict["length"], ep_task_index, dtype=int)))
@@ -340,7 +350,9 @@ def lerobot_dataset_metadata_factory(
             tasks = tasks_factory(total_tasks=info["total_tasks"])
         if not episodes:
             episodes = episodes_factory(
-                total_episodes=info["total_episodes"], total_frames=info["total_frames"], tasks=tasks
+                total_episodes=info["total_episodes"],
+                total_frames=info["total_frames"],
+                tasks=tasks,
             )
 
         mock_snapshot_download = mock_snapshot_download_factory(
@@ -392,7 +404,9 @@ def lerobot_dataset_factory(
     ) -> LeRobotDataset:
         if not info:
             info = info_factory(
-                total_episodes=total_episodes, total_frames=total_frames, total_tasks=total_tasks
+                total_episodes=total_episodes,
+                total_frames=total_frames,
+                total_tasks=total_tasks,
             )
         if not stats:
             stats = stats_factory(features=info["features"])

tests/fixtures/files.py

@@ -102,7 +102,10 @@ def episode_path(episodes_factory):
 @pytest.fixture(scope="session")
 def single_episode_parquet_path(hf_dataset_factory, info_factory):
     def _create_single_episode_parquet(
-        dir: Path, ep_idx: int = 0, hf_dataset: datasets.Dataset | None = None, info: dict | None = None
+        dir: Path,
+        ep_idx: int = 0,
+        hf_dataset: datasets.Dataset | None = None,
+        info: dict | None = None,
     ) -> Path:
         if not info:
             info = info_factory()

tests/fixtures/hub.py

@@ -67,7 +67,9 @@ def mock_snapshot_download_factory(
             tasks = tasks_factory(total_tasks=info["total_tasks"])
         if not episodes:
             episodes = episodes_factory(
-                total_episodes=info["total_episodes"], total_frames=info["total_frames"], tasks=tasks
+                total_episodes=info["total_episodes"],
+                total_frames=info["total_frames"],
+                tasks=tasks,
             )
         if not hf_dataset:
             hf_dataset = hf_dataset_factory(tasks=tasks, episodes=episodes, fps=info["fps"])
@@ -93,7 +95,13 @@ def mock_snapshot_download_factory(
 
             # List all possible files
             all_files = []
-            meta_files = [INFO_PATH, STATS_PATH, EPISODES_STATS_PATH, TASKS_PATH, EPISODES_PATH]
+            meta_files = [
+                INFO_PATH,
+                STATS_PATH,
+                EPISODES_STATS_PATH,
+                TASKS_PATH,
+                EPISODES_PATH,
+            ]
             all_files.extend(meta_files)
 
             data_files = []
@@ -105,7 +113,9 @@ def mock_snapshot_download_factory(
             all_files.extend(data_files)
 
             allowed_files = filter_repo_objects(
-                all_files, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns
+                all_files,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
             )
 
             # Create allowed files

tests/motors/test_motors.py

@@ -43,7 +43,10 @@ import time
 import numpy as np
 import pytest
 
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
 from lerobot.scripts.find_motors_bus_port import find_port
 from tests.utils import TEST_MOTOR_TYPES, make_motors_bus, require_motor