Port SAC WIP (#581)

Co-authored-by: KeWang1017 <ke.wang@helloleap.ai>

.pre-commit-config.yaml

@@ -3,7 +3,7 @@ default_language_version:
     python: python3.10
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: check-added-large-files
       - id: debug-statements
@@ -14,11 +14,11 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.16.0
+    rev: v3.19.0
     hooks:
     -   id: pyupgrade
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.2
+    rev: v0.8.2
     hooks:
       - id: ruff
         args: [--fix]
@@ -32,6 +32,6 @@ repos:
           - "--check"
           - "--no-update"
   - repo: https://github.com/gitleaks/gitleaks
-    rev: v8.18.4
+    rev: v8.21.2
     hooks:
       - id: gitleaks

examples/6_add_image_transforms.py

@@ -10,10 +10,10 @@ from torchvision.transforms import ToPILImage, v2
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 
-dataset_repo_id = "lerobot/aloha_static_tape"
+dataset_repo_id = "lerobot/aloha_static_screw_driver"
 
 # Create a LeRobotDataset with no transformations
-dataset = LeRobotDataset(dataset_repo_id)
+dataset = LeRobotDataset(dataset_repo_id, episodes=[0])
 # This is equivalent to `dataset = LeRobotDataset(dataset_repo_id, image_transforms=None)`
 
 # Get the index of the first observation in the first episode
@@ -28,12 +28,13 @@ transforms = v2.Compose(
     [
         v2.ColorJitter(brightness=(0.5, 1.5)),
         v2.ColorJitter(contrast=(0.5, 1.5)),
+        v2.ColorJitter(hue=(-0.1, 0.1)),
         v2.RandomAdjustSharpness(sharpness_factor=2, p=1),
     ]
 )
 
 # Create another LeRobotDataset with the defined transformations
-transformed_dataset = LeRobotDataset(dataset_repo_id, image_transforms=transforms)
+transformed_dataset = LeRobotDataset(dataset_repo_id, episodes=[0], image_transforms=transforms)
 
 # Get a frame from the transformed dataset
 transformed_frame = transformed_dataset[first_idx][transformed_dataset.meta.camera_keys[0]]

examples/9_use_aloha.md

@@ -56,7 +56,7 @@ python lerobot/scripts/control_robot.py teleoperate \
     --robot-overrides max_relative_target=5
 ```
 
-By adding `--robot-overrides max_relative_target=5`, we override the default value for `max_relative_target` defined in `lerobot/configs/robot/aloha.yaml`. It is expected to be `5` to limit the magnitude of the movement for more safety, but the teloperation won't be smooth. When you feel confident, you can disable this limit by adding `--robot-overrides max_relative_target=null` to the command line:
+By adding `--robot-overrides max_relative_target=5`, we override the default value for `max_relative_target` defined in `lerobot/configs/robot/aloha.yaml`. It is expected to be `5` to limit the magnitude of the movement for more safety, but the teleoperation won't be smooth. When you feel confident, you can disable this limit by adding `--robot-overrides max_relative_target=null` to the command line:
 ```bash
 python lerobot/scripts/control_robot.py teleoperate \
     --robot-path lerobot/configs/robot/aloha.yaml \

lerobot/common/datasets/lerobot_dataset.py

@@ -298,7 +298,7 @@ class LeRobotDatasetMetadata:
                     f"Some cameras in your {robot.robot_type} robot don't have an fps matching the fps of your dataset."
                     "In this case, frames from lower fps cameras will be repeated to fill in the blanks."
                 )
-        elif robot_type is None or features is None:
+        elif features is None:
             raise ValueError(
                 "Dataset features must either come from a Robot or explicitly passed upon creation."
             )

lerobot/common/robot_devices/control_utils.py

@@ -122,12 +122,12 @@ def predict_action(observation, policy, device, use_amp):
 
 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 
+    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 
+        assign_rewards (bool): If True, allows annotating the collected trajectory
         with a binary reward at the end of the episode to indicate success.
     """
     events = {}
@@ -199,7 +199,7 @@ def init_policy(pretrained_policy_name_or_path, policy_overrides):
 def warmup_record(
     robot,
     events,
-    enable_teloperation,
+    enable_teleoperation,
     warmup_time_s,
     display_cameras,
     fps,
@@ -210,7 +210,7 @@ def warmup_record(
         display_cameras=display_cameras,
         events=events,
         fps=fps,
-        teleoperate=enable_teloperation,
+        teleoperate=enable_teleoperation,
     )
 
 
@@ -287,6 +287,8 @@ def control_loop(
 
         if dataset is not None:
             frame = {**observation, **action}
+            if "next.reward" in events:
+                frame["next.reward"] = events["next.reward"]
             dataset.add_frame(frame)
 
         if display_cameras and not is_headless():

lerobot/scripts/control_robot.py

@@ -224,7 +224,7 @@ def record(
     else:
         raise NotImplementedError("Only single-task recording is supported for now")
 
-        # Load pretrained policy
+    # Load pretrained policy
     if pretrained_policy_name_or_path is not None:
         policy, policy_fps, device, use_amp = init_policy(pretrained_policy_name_or_path, policy_overrides)
 

lerobot/scripts/control_sim_robot.py

@@ -35,7 +35,7 @@ python lerobot/scripts/visualize_dataset.py \
     --episode-index 0
 ```
 
-- Replay a sequence of test episodes: 
+- Replay a sequence of test episodes:
 ```bash
 python lerobot/scripts/control_sim_robot.py replay \
     --robot-path lerobot/configs/robot/your_robot_config.yaml \