Add leader support plus some fixes. (#40)

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
     ]
 }

README.md

@@ -18,13 +18,13 @@ Comprehensive monorepo for LeKiwi robot software development, combining simulati
 ## :package: Project Structure
 
 ```
-lekiwi-dora/
+lekiwi/
 ├── 📁 packages/                             # Python packages
 │   ├── 🎮 lekiwi_sim/                       # MuJoCo simulation environment
 │   ├── 🤖 lekiwi_lerobot/                   # LeRobot integration scripts
 │   └── 🕹️ lekiwi_teleoperate/               # Teleoperation interface
 ├── 📁 dora/                                 # Dora Integration
-│   └── 📁 graphs/                           # Dora dataflows
+│   └── 📁 lekiwi/graphs/                    # Dora dataflows
 │   └── 📁 node_hub/                         # Dora nodes
 │       ├── 🔗 dora_lekiwi_client/           # Robot interface node
 │       ├── 🧠 dora_run_policy/              # Policy execution node
@@ -166,9 +166,18 @@ Manual control interface using the LeRobot API:
 # Start simulation or real robot first
 uv run lekiwi_host_sim  # For simulation
 
-# Then teleoperate
+# Then teleoperate the simulated or real robot.
 uv run lekiwi_teleoperate
 ```
+By default it will allow you to teleoperate the Lekiwi completely using the keyboard.
+If you have a leader arm you can use it to teleoperate the arm part in the simulation as well:
+
+```bash
+uv run lekiwi_teleoperate --leader-arm
+```
+
+[lekiwi_sim_leader_LOW.webm](https://github.com/user-attachments/assets/76e565cd-93d2-42ae-976d-3d25091039a4)
+
 
 [lekiwi_sim_pick_cube.webm](https://github.com/user-attachments/assets/32af6eca-834b-4ba4-8609-33bc428cb75f)
 
@@ -182,7 +191,7 @@ uv run lekiwi_lerobot_record --repo-id your_username/dataset_name --episodes 50
 uv run lekiwi_lerobot_replay --repo-id your_username/dataset_name --episode 0
 
 # Train a policy (see lekiwi_lerobot README for full training options)
-python -m lerobot.scripts.train \
+uv run lerobot-train \
   --dataset.repo_id=your_username/dataset_name \
   --policy.type=act \
   --output_dir=outputs/my_policy
@@ -202,9 +211,9 @@ See [packages/lekiwi_lerobot/README.md](packages/lekiwi_lerobot/README.md) for d
 
 ### Available Dataflows
 
-The repository includes pre-configured dataflow graphs in `dora/lekiwi_sim/graphs/`:
+The repository includes pre-configured dataflow graphs in `dora/lekiwi/graphs/`:
 
-**1. Policy Execution Dataflow** (`mujoco_sim.yml`):
+**1. Policy Execution Dataflow** (`dataflow.yml`):
    - Complete pipeline for running trained policies on simulation
    - Connects robot observations → policy inference → robot actions
    - Includes camera feeds and state observations
@@ -214,21 +223,25 @@ The repository includes pre-configured dataflow graphs in `dora/lekiwi_sim/graph
 
 **Prerequisites:**
 ```bash
-# Start simulation in separate terminal
+
+# Start simulation (or alternatively the real robot).
 uv run lekiwi_host_sim
 ```
 
 **Run the policy execution dataflow:**
 ```bash
 # Navigate to dataflow directory
-cd dora/lekiwi_sim/graphs/
+cd dora/lekiwi/graphs/
+
+# Build if not built already
+dora build dataflow.yml
 
 # Start the dataflow
-dora run mujoco_sim.yml --uv
+dora run dataflow.yml --uv
 
 ```
 
-**Optional features** (uncomment in `mujoco_sim.yml`):
+**Optional features** (uncomment in `dataflow.yml`):
 - **Visualization**: Enable `rerun-viz` node for real-time 3D visualization
 - **Data Recording**: Enable `dora-record` node to save observations to Parquet files
 - **Testing Mode**: Use `dora_lekiwi_action_publisher` instead of policy for hardcoded actions

dora/lekiwi/.gitkeep

@@ -10,6 +10,8 @@ nodes:
       - observation_state
       - image_front
       - image_wrist
+    env:
+      LEKIWI_IP: 192.168.1.108 # Use the appropriate IP for real robot
 
   ######################################################
   # Example of a hard-coded action publisher.

dora/lekiwi_sim/graphs/mujoco_sim.yml dora/lekiwi/graphs/dataflow.ymldora/lekiwi_sim/graphs/mujoco_sim.yml renamed to dora/lekiwi/graphs/dataflow.yml

@@ -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/README.md

@@ -3,6 +3,7 @@
 
 from lekiwi_lerobot.utils import record_loop
 from lekiwi_teleoperate.teleoperate.arm import ArmTeleop
+from lerobot.cameras.configs import CameraConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
@@ -11,6 +12,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 +34,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 +68,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:
@@ -72,13 +93,28 @@ def main() -> None:
         level=log_level, format="%(asctime)s | %(levelname)-8s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
     )
 
+    # Camera config should match the one used in the robot config
+    # when starting the robot host or simulation.
+    #
+    # Based on: --robot.cameras="{ front: {type: opencv, index_or_path: /dev/video0, width: 640, height: 480, fps: 30},
+    # wrist: {type: opencv, index_or_path: /dev/video2, width: 640, height: 480, fps: 30}}"
+    camera_config: dict[str, CameraConfig] = {
+        "front": CameraConfig(width=640, height=480, fps=30),
+        "wrist": CameraConfig(width=640, height=480, fps=30),
+    }
+
     # 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", cameras=camera_config)
     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 +138,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 +151,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 +166,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 +181,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 +204,8 @@ def main() -> None:
 
     robot.disconnect()
     keyboard.disconnect()
+    if args.leader_arm:
+        leader_arm.disconnect()
     listener.stop()
 
 

packages/lekiwi_lerobot/lekiwi_lerobot/record.py

@@ -45,6 +45,13 @@ def main() -> None:
     """Main function to run the LeKiwi replay client."""
     parser = argparse.ArgumentParser(description="Run the LeKiwi replay client.")
 
+    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(
         "-l",
         "--level",
@@ -79,7 +86,7 @@ def main() -> None:
         level=log_level, format="%(asctime)s | %(levelname)-8s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
     )
 
-    robot_config = LeKiwiClientConfig(remote_ip="127.0.0.1", id="lekiwi")
+    robot_config = LeKiwiClientConfig(remote_ip=args.ip, id="lekiwi")
     robot = LeKiwiClient(robot_config)
 
     logging.info(f"Downloading dataset from {args.repo_id} into {args.directory}")

packages/lekiwi_lerobot/lekiwi_lerobot/replay.py

@@ -3,7 +3,6 @@
 import time
 
 from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
@@ -63,7 +62,13 @@ def main() -> None:
 
     logging.info(f"Loading policy from '{args.policy}'")
     if args.policy_type == "act":
+        from lerobot.policies.act.modeling_act import ACTPolicy
+
         policy = ACTPolicy.from_pretrained(args.policy)
+    elif args.policy_type == "smolvla":
+        from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+
+        policy = SmolVLAPolicy.from_pretrained(args.policy)
     else:
         raise ValueError(f"Policy type '{args.policy_type}' not supported.")
     policy.reset()

packages/lekiwi_lerobot/lekiwi_lerobot/run_policy.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.