Add leader support plus some fixes.

Signed-off-by: Franco Cipollone <franco.c@ekumenlabs.com>

Author: francocipollone

.devcontainer/lekiwi-dev/Dockerfile

@@ -9,7 +9,7 @@ FROM mcr.microsoft.com/devcontainers/base:${BASE_IMAGE}
 # Configuration
 ################################################################################
 
-ARG RUST_VERSION=1.86.0
+ARG RUST_VERSION=1.88.0
 
 ################################################################################
 # User 'dev'

.devcontainer/lekiwi-dev/devcontainer.json

@@ -40,6 +40,12 @@
         "--network=host",
         "--runtime=nvidia", // Use NVIDIA runtime for GPU access
         "--gpus=all", // Use all available GPUs
-        "--name=lekiwi-dora"
+        "--name=lekiwi-dev"
+    ],
+    // 9090: This is a web server that serves the frontend for the Rerun application.
+    // 9876: This is the grpc server to transmit info for the items to be drawn.
+    "forwardPorts": [
+        9090,
+        9876
     ]
 }

packages/lekiwi_lerobot/README.md

@@ -79,7 +79,7 @@ Once you have a dataset you can start training a model. For this, we can rely di
 Train imitation learning policies using collected data:
 
 ```bash
-uv run python -m lerobot.scripts.train \
+uv run lerobot-train \
   --dataset.repo_id=<username/my_dataset> \
   --policy.type=act \
   --output_dir=outputs/train/username/my_policy \

packages/lekiwi_lerobot/lekiwi_lerobot/record.py

@@ -11,6 +11,7 @@
     KeyboardTeleop,
     KeyboardTeleopConfig,
 )
+from lerobot.teleoperators.so101_leader import SO101Leader, SO101LeaderConfig
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import (
     init_keyboard_listener,
@@ -32,6 +33,13 @@ def main() -> None:
         default="INFO",
         help="Set the logging level (default: INFO). Case-insensitive.",
     )
+    parser.add_argument(
+        "-i",
+        "--ip",
+        type=str,
+        default="127.0.0.1",
+        help="IP address of the robot (default: 127.0.0.1).",
+    )
     parser.add_argument(
         "-r",
         "--repo-id",
@@ -59,6 +67,18 @@ def main() -> None:
         dest="visualize",
         help="Disable Rerun visualization during recording.",
     )
+    parser.add_argument(
+        "-la",
+        "--leader-arm",
+        action="store_true",
+        help="Use the leader arm for teleoperation (default: False).",
+    )
+    parser.add_argument(
+        "--leader-arm-port",
+        type=str,
+        default="/dev/ttyACM0",
+        help="Serial port for the leader arm (default: /dev/ttyACM0).",
+    )
 
     args = parser.parse_args()
     if args.repo_id is None:
@@ -73,12 +93,17 @@ def main() -> None:
     )
 
     # Create the robot and teleoperator configurations
-    robot_config = LeKiwiClientConfig(remote_ip="127.0.0.1", id="lekiwi")
+    robot_config = LeKiwiClientConfig(remote_ip=args.ip, id="lekiwi")
     keyboard_config = KeyboardTeleopConfig()
+    if args.leader_arm:
+        teleop_arm_config = SO101LeaderConfig(port=args.leader_arm_port, id="lekiwi_leader_arm")
 
     robot = LeKiwiClient(robot_config)
     keyboard = KeyboardTeleop(keyboard_config)
-    arm_keyboard_handler = ArmTeleop()
+    if args.leader_arm:
+        leader_arm = SO101Leader(teleop_arm_config)
+    else:
+        arm_keyboard_handler = ArmTeleop()
     # Configure the dataset features
     action_features = hw_to_dataset_features(robot.action_features, ACTION)
     obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
@@ -102,6 +127,8 @@ def main() -> None:
     #  - Sim robot: this script running on LeKiwi sim: `uv run lekiwi_sim --robot.id=my_awesome_kiwi`
     robot.connect()
     keyboard.connect()
+    if args.leader_arm:
+        leader_arm.connect()
 
     if args.visualize:
         logging.info("Initializing Rerun for visualization.")
@@ -113,10 +140,13 @@ def main() -> None:
 
     if not robot.is_connected or not keyboard.is_connected:
         raise ValueError("Robot or keyboard is not connected!")
+    if args.leader_arm and not leader_arm.is_connected:
+        raise ValueError("Leader arm is not connected!")
     logging.info("Robot and keyboard are connected.")
     recorded_episodes = 0
     while recorded_episodes < args.episodes and not events["stop_recording"]:
-        arm_keyboard_handler = ArmTeleop()
+        if not args.leader_arm:
+            arm_keyboard_handler = ArmTeleop()
         logging.info(f"Recording episode {recorded_episodes}")
         # Run the record loop
         record_loop(
@@ -125,7 +155,7 @@ def main() -> None:
             fps=FPS,
             dataset=dataset,
             keyboard_handler=keyboard,
-            arm_keyboard_handler=arm_keyboard_handler,
+            arm_keyboard_handler=leader_arm if args.leader_arm else arm_keyboard_handler,
             control_time_s=EPISODE_TIME_SEC,
             single_task=args.task,
             display_data=args.visualize,
@@ -140,7 +170,7 @@ def main() -> None:
                 fps=FPS,
                 dataset=None,  # Don't record during reset phase
                 keyboard_handler=keyboard,
-                arm_keyboard_handler=arm_keyboard_handler,
+                arm_keyboard_handler=leader_arm if args.leader_arm else arm_keyboard_handler,
                 control_time_s=RESET_TIME_SEC,
                 single_task=args.task,
                 display_data=args.visualize,
@@ -163,6 +193,8 @@ def main() -> None:
 
     robot.disconnect()
     keyboard.disconnect()
+    if args.leader_arm:
+        leader_arm.disconnect()
     listener.stop()
 
 

packages/lekiwi_lerobot/lekiwi_lerobot/utils.py

@@ -1,6 +1,6 @@
 import logging
 import time
-from typing import Any
+from typing import Any, Union
 
 from lekiwi_teleoperate.teleoperate.arm import ArmTeleop
 from lerobot.datasets.image_writer import safe_stop_image_writer
@@ -15,6 +15,7 @@
 from lerobot.teleoperators.keyboard import (
     KeyboardTeleop,
 )
+from lerobot.teleoperators.so101_leader import SO101Leader
 from lerobot.utils.control_utils import (
     predict_action,
 )
@@ -32,7 +33,7 @@ def record_loop(
     fps: int,
     dataset: LeRobotDataset | None = None,
     keyboard_handler: KeyboardTeleop | None = None,
-    arm_keyboard_handler: ArmTeleop | None = None,
+    arm_keyboard_handler: Union[ArmTeleop, SO101Leader, None] = None,
     policy: PreTrainedPolicy | None = None,
     preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None,
     postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None,
@@ -101,7 +102,13 @@ def record_loop(
         elif policy is None and keyboard_handler is not None and arm_keyboard_handler is not None:
             pressed_keys = keyboard_handler.get_action()
             base_action = robot._from_keyboard_to_base_action(pressed_keys)
-            arm_action = arm_keyboard_handler.from_keyboard_to_arm_action(pressed_keys)
+
+            # Handle both ArmTeleop (keyboard-based) and SO101Leader (physical arm)
+            if isinstance(arm_keyboard_handler, SO101Leader):
+                arm_action = arm_keyboard_handler.get_action()
+                arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+            else:
+                arm_action = arm_keyboard_handler.from_keyboard_to_arm_action(pressed_keys)
 
             action = {**base_action, **arm_action}  # Merge base and arm actions
             # TODO(francocipollone): We would probably want to use the teleop_action_processor here.

packages/lekiwi_sim/lekiwi_sim/assets/so_arm100/so_arm100.xml

@@ -33,7 +33,7 @@
       <joint frictionloss="0.1" armature="0.1"/>
       <position kp="50" dampratio="1" forcerange="-3.5 3.5"/>
       <default class="Rotation">
-        <joint axis="0 1 0" range="-1.92 1.92"/>
+        <joint axis="0 -1 0" range="-1.92 1.92"/>
       </default>
       <default class="Pitch">
         <joint axis="1 0 0" range="-1.747 1.747"/>

packages/lekiwi_teleoperate/lekiwi_teleoperate/main.py

@@ -4,6 +4,7 @@
 
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
+from lerobot.teleoperators.so101_leader import SO101Leader, SO101LeaderConfig
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
@@ -59,6 +60,18 @@ def main() -> None:
         default="127.0.0.1",
         help="IP address of the robot (default: 127.0.0.1).",
     )
+    parser.add_argument(
+        "-la",
+        "--leader-arm",
+        action="store_true",
+        help="Use the leader arm for teleoperation (default: False).",
+    )
+    parser.add_argument(
+        "--leader-arm-port",
+        type=str,
+        default="/dev/ttyACM0",
+        help="Serial port for the leader arm (default: /dev/ttyACM0).",
+    )
     args = parser.parse_args()
     log_level = args.level.upper()
     logging.basicConfig(
@@ -69,24 +82,33 @@ def main() -> None:
     # Create the robot and teleoperator configurations
     robot_config = LeKiwiClientConfig(remote_ip=args.ip, id="my_lekiwi")
     keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
+    if args.leader_arm:
+        teleop_arm_config = SO101LeaderConfig(port=args.leader_arm_port, id="lekiwi_leader_arm")
 
     robot = LeKiwiClient(robot_config)
     keyboard = KeyboardTeleop(keyboard_config)
 
     # To connect you already should have this script running on LeKiwi: `uv run lekiwi_sim`
     robot.connect()
     keyboard.connect()
+    if args.leader_arm:
+        leader_arm = SO101Leader(teleop_arm_config)
+        leader_arm.connect()
+    else:
+        arm_teleop = ArmTeleop()
 
     init_rerun(session_name="lekiwi_teleop")
 
-    if not robot.is_connected or not keyboard.is_connected:
-        raise ValueError("Robot, leader arm of keyboard is not connected!")
+    if not robot.is_connected:
+        raise ValueError("Robot is not connected!")
+    if not keyboard.is_connected:
+        raise ValueError("Keyboard is not connected!")
+    if args.leader_arm and not leader_arm.is_connected:
+        raise ValueError("Leader arm is not connected!")
 
     logging.info("Robot and keyboard are connected.")
     print(COMMANDS_STR.format(**robot_config.teleop_keys, **ArmTeleop.ARM_TELEOP_KEYS))
 
-    arm_teleop = ArmTeleop()
-
     try:
         while True:
             t0 = time.perf_counter()
@@ -95,7 +117,11 @@ def main() -> None:
 
             keyboard_keys = keyboard.get_action()
             base_action = robot._from_keyboard_to_base_action(keyboard_keys)
-            arm_action = arm_teleop.from_keyboard_to_arm_action(keyboard_keys)
+            if args.leader_arm:
+                arm_action = leader_arm.get_action()
+                arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+            else:
+                arm_action = arm_teleop.from_keyboard_to_arm_action(keyboard_keys)
 
             log_rerun_data(observation, {**arm_action, **base_action})
 

packages/lekiwi_teleoperate/lekiwi_teleoperate/teleoperate/arm.py

@@ -11,8 +11,8 @@ class ArmTeleop:
     """
 
     ARM_TELEOP_KEYS = {
-        "shoulder_pan_left": "t",
-        "shoulder_pan_right": "g",
+        "shoulder_pan_left": "g",
+        "shoulder_pan_right": "t",
         "shoulder_lift_up": "y",
         "shoulder_lift_down": "h",
         "elbow_flex_up": "u",