From 9a25c41f91ed6310a15c40d17fd47ad0e7d20bb7 Mon Sep 17 00:00:00 2001
From: Steven Palma <steven.palma@huggingface.co>
Date: Mon, 17 Mar 2025 17:38:08 +0100
Subject: [PATCH] fix(lekiwi): HW fixes v0.2

---
 lerobot/common/robots/lekiwi/daemon_lekiwi.py | 20 ++++++++++++-------
 .../common/robots/lekiwi/daemon_lekiwi_app.py | 12 +++++------
 lerobot/common/robots/lekiwi/lekiwi_robot.py  |  2 +-
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/lerobot/common/robots/lekiwi/daemon_lekiwi.py b/lerobot/common/robots/lekiwi/daemon_lekiwi.py
index 3878e08e..3a76f809 100644
--- a/lerobot/common/robots/lekiwi/daemon_lekiwi.py
+++ b/lerobot/common/robots/lekiwi/daemon_lekiwi.py
@@ -108,6 +108,8 @@ class DaemonLeKiwiRobot(Robot):
         pass
 
     def connect(self) -> None:
+        """Establishes ZMQ sockets with the remote mobile robot"""
+
         if self.is_connected:
             raise DeviceAlreadyConnectedError(
                 "LeKiwi Daemon is already connected. Do not run `robot.connect()` twice."
@@ -340,7 +342,7 @@ class DaemonLeKiwiRobot(Robot):
         """
         Capture observations from the remote robot: current follower arm positions,
         present wheel speeds (converted to body-frame velocities: x, y, theta),
-        and a camera frame.
+        and a camera frame. Receives over ZMQ, translate to body-frame vel
         """
         if not self.is_connected:
             raise DeviceNotConnectedError(
@@ -409,7 +411,7 @@ class DaemonLeKiwiRobot(Robot):
     # t3': execute motor_cmd
     # t3'-t2' << t3-t1
     def send_action(self, action: np.ndarray) -> np.ndarray:
-        """Command lekiwi to move to a target joint configuration.
+        """Command lekiwi to move to a target joint configuration. Translates to motor space + sends over ZMQ
 
         Args:
             action (np.ndarray): array containing the goal positions for the motors.
@@ -425,7 +427,7 @@ class DaemonLeKiwiRobot(Robot):
                 "ManipulatorRobot is not connected. You need to run `robot.connect()`."
             )
 
-        goal_pos: np.array = np.empty(9)
+        goal_pos: np.array = np.zeros(9)
 
         if self.robot_mode is RobotMode.AUTO:
             # TODO(Steven): Not yet implemented. The policy outputs might need a different conversion
@@ -437,11 +439,13 @@ class DaemonLeKiwiRobot(Robot):
                 logging.error("Action should include at least the 6 states of the leader arm")
                 raise InvalidActionError
 
-            # TODO(Steven): Assumes size and order is respected
-            wheel_actions = [v for _, v in self.from_keyboard_to_wheel_action(action[6:])]
             goal_pos[:6] = action[:6]
-            goal_pos[6:] = wheel_actions
-            self.zmq_cmd_socket.send_string(json.dumps(goal_pos))  # action is in motor space
+            if action.size > 6:
+                # TODO(Steven): Assumes size and order is respected
+                wheel_actions = [v for _, v in self.from_keyboard_to_wheel_action(action[6:]).items()]
+                goal_pos[6:] = wheel_actions
+
+            self.zmq_cmd_socket.send_string(json.dumps(goal_pos.tolist()))  # action is in motor space
 
         return goal_pos
 
@@ -450,6 +454,8 @@ class DaemonLeKiwiRobot(Robot):
         pass
 
     def disconnect(self):
+        """Cleans ZMQ comms"""
+
         if not self.is_connected:
             raise DeviceNotConnectedError(
                 "LeKiwi is not connected. You need to run `robot.connect()` before disconnecting."
diff --git a/lerobot/common/robots/lekiwi/daemon_lekiwi_app.py b/lerobot/common/robots/lekiwi/daemon_lekiwi_app.py
index b459f21c..a52f534b 100644
--- a/lerobot/common/robots/lekiwi/daemon_lekiwi_app.py
+++ b/lerobot/common/robots/lekiwi/daemon_lekiwi_app.py
@@ -11,7 +11,7 @@ from lerobot.common.teleoperators.so100 import SO100Teleop, SO100TeleopConfig
 
 def main():
     logging.info("Configuring Teleop Devices")
-    leader_arm_config = SO100TeleopConfig(port="/dev/tty.usbmodem585A0085511")
+    leader_arm_config = SO100TeleopConfig(port="/dev/tty.usbmodem58760429271")
     leader_arm = SO100Teleop(leader_arm_config)
 
     keyboard_config = KeyboardTeleopConfig()
@@ -26,7 +26,7 @@ def main():
     robot = DaemonLeKiwiRobot(robot_config)
 
     logging.info("Connecting remote LeKiwiRobot")
-    robot.connect()  # Establishes ZMQ sockets with the remote mobile robot
+    robot.connect()
     robot.robot_mode = RobotMode.TELEOP
 
     logging.info("Starting LeKiwiRobot teleoperation")
@@ -35,9 +35,9 @@ def main():
     while duration < 20:
         arm_action = leader_arm.get_action()
         base_action = keyboard.get_action()
-        action = np.concatenate((arm_action, base_action))
-        _action_sent = robot.send_action(action)  # Translates to motor space + sends over ZMQ
-        _observation = robot.get_observation()  # Receives over ZMQ, translate to body-frame vel
+        action = np.append(arm_action, base_action) if base_action.size > 0 else arm_action
+        _action_sent = robot.send_action(action)
+        _observation = robot.get_observation()
 
         # dataset.save(action_sent, obs)
 
@@ -48,7 +48,7 @@ def main():
         duration = time.perf_counter() - start
 
     logging.info("Disconnecting Teleop Devices and LeKiwiRobot Daemon")
-    robot.disconnect()  # Cleans ZMQ comms
+    robot.disconnect()
     leader_arm.disconnect()
     keyboard.disconnect()
 
diff --git a/lerobot/common/robots/lekiwi/lekiwi_robot.py b/lerobot/common/robots/lekiwi/lekiwi_robot.py
index f5204ffd..ed0e7af4 100644
--- a/lerobot/common/robots/lekiwi/lekiwi_robot.py
+++ b/lerobot/common/robots/lekiwi/lekiwi_robot.py
@@ -304,7 +304,7 @@ class LeKiwiRobot(Robot):
 
                 try:
                     msg = self.zmq_cmd_socket.recv_string(zmq.NOBLOCK)
-                    data = json.loads(msg)
+                    data = np.array(json.loads(msg))
                     self.send_action(data)
                     last_cmd_time = time.time()
                 except zmq.Again: