Adds basic architecture for the Lekiwi robot simulation in MuJoCo. (#11)

Signed-off-by: Ari Lowy <arilow@ekumenlabs.com>
Signed-off-by: Franco Cipollone <franco.c@ekumenlabs.com>
Co-authored-by: Franco Cipollone <franco.c@ekumenlabs.com>

Author: arilow

packages/lekiwi_sim/README.md

@@ -2,4 +2,6 @@ Create environment: `uv venv -p 3.11 --seed`
 
 Install: `uv pip install -e .`
 
-Run: `uv run lekiwi_sim`
+Run MuJoCo simulation host to be accessed via `lerobot.robot.LekiwiClient`: `uv run lekiwi_sim_host`
+
+Run standalone MuJoCo simulation: `uv run standalone_mujoco_sim`

packages/lekiwi_sim/lekiwi_sim/assets/mjcf_lcmm_robot.xml

@@ -0,0 +1,83 @@
+import numpy as np
+
+
+class LeKiwiMobileBase:
+    """Class representing the mobile base of the LeKiwi robot.
+
+    The mobile base holds the kinematics of a 3-wheeled omnidirectional robot.
+    Where the wheels are placed at 120 degrees from each other.
+
+    The wheel configuration is assumed to be:
+    - Wheel 1 is at the "left", forming 30 degrees to the forward direction.
+    - Wheel 2 is at the "right" clockwise 120 degrees from Wheel 1.
+    - Wheel 3 is at the "back", clockwise 240 degrees from Wheel 1.
+    """
+
+    def __init__(self, wheel_radius: float, robot_base_radius: float) -> None:
+        """Initialize the LeKiwiMobileBase.
+
+        Args:
+            wheel_radius (float): The radius of the wheels in meters.
+            robot_base_radius (float): The distance from the center of rotation to each wheel in meters.
+
+        """
+        self.wheel_radius = wheel_radius
+        self.robot_base_radius = robot_base_radius
+
+        self.F_matrix = self._compute_forward_kinematics_matrix(self.wheel_radius, self.robot_base_radius)
+        self.F_matrix_inv = np.linalg.inv(self.F_matrix)
+
+    def forward_kinematics(self, wheel_speeds: np.ndarray) -> np.ndarray:
+        """Compute the forward kinematics of the mobile base.
+
+        Args:
+            wheel_speeds (np.ndarray): Angular velocities of the wheels [left, right, back] in rad/s.
+
+        Returns:
+            np.ndarray: An array representing the robot's velocity [vx, vy, omega] in m/s and rad/s.
+
+        """
+        if wheel_speeds.shape != (3,):
+            raise ValueError("wheel_speeds must be a numpy array of shape (3,)")
+
+        # Compute the robot's velocity in the body frame
+        return self.F_matrix @ wheel_speeds
+
+    def inverse_kinematics(self, robot_velocity: np.ndarray) -> np.ndarray:
+        """Compute the inverse kinematics of the mobile base.
+
+        Args:
+            robot_velocity (np.ndarray): The desired robot velocity [vx, vy, omega] in m/s and rad/s.
+
+        Returns:
+            np.ndarray: An array representing the angular velocities of the wheels [left, right, back] in rad/s.
+
+        """
+        if robot_velocity.shape != (3,):
+            raise ValueError("robot_velocity must be a numpy array of shape (3,)")
+
+        # Compute the wheel speeds required to achieve the desired robot velocity
+        return self.F_matrix_inv @ robot_velocity
+
+    @staticmethod
+    def _compute_forward_kinematics_matrix(r: float, L: float) -> np.ndarray:
+        """Calculates the forward kinematics transformation matrix for a 3-wheeled omni-drive robot.
+
+        This matrix maps wheel speeds(angular velocities) to robot body speeds(linear and angular velocities)
+        following the conventions:
+        - The robot's forward direction is along the positive x-axis.
+        - The robot's left direction is along the positive y-axis.
+        - The robot's up direction is along the positive z-axis.
+        - The robot's angular velocity (omega) is positive for counter-clockwise rotation when viewed from above.
+
+        Args:
+            r (float): The radius of each wheel (in meters).
+            L (float): The radius of the robot base, from center to wheel (in meters).
+
+        Returns:
+            np.ndarray: A 3x3 transformation matrix that maps wheel speeds to robot body speeds.
+
+        """
+        return r * np.array(
+            [[np.sqrt(3) / 2, -np.sqrt(3) / 2, 0], [-1 / 2, -1 / 2, 1], [-1 / (3 * L), -1 / (3 * L), -1 / (3 * L)]]
+        )

packages/lekiwi_sim/lekiwi_sim/kinematics.py

@@ -0,0 +1,151 @@
+"""Lekiwi host but for simulation"""
+
+#!/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 argparse
+import json
+import logging
+import time
+
+# import cv2
+import zmq
+from lerobot.robots.lekiwi.config_lekiwi import LeKiwiHostConfig
+
+from .robot import LeKiwiMujoco, LeKiwiMujocoConfig
+
+
+class ZMQHandler:
+    """LeKiwi Host agent for simulation."""
+
+    def __init__(self, config: LeKiwiHostConfig):
+        """Initialize the LeKiwi Host agent for simulation."""
+        self.zmq_context = zmq.Context()
+        self.zmq_cmd_socket = self.zmq_context.socket(zmq.PULL)
+        self.zmq_cmd_socket.setsockopt(zmq.CONFLATE, 1)
+        self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")
+
+        self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
+        self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1)
+        self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}")
+
+        self.connection_time_s = config.connection_time_s
+        self.watchdog_timeout_ms = config.watchdog_timeout_ms
+        self.max_loop_freq_hz = config.max_loop_freq_hz
+
+    def disconnect(self) -> None:
+        """Disconnect the ZMQ sockets and context."""
+        self.zmq_observation_socket.close()
+        self.zmq_cmd_socket.close()
+        self.zmq_context.term()
+
+
+def main() -> None:
+    """Main function to run the LeKiwi simulation host."""
+    parser = argparse.ArgumentParser(
+        description="Run the LeKiwi simulation host to be accessed via lerobot.robot.LekiwiClient."
+    )
+
+    parser.add_argument(
+        "-l",
+        "--level",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        default="INFO",
+        help="Set the logging level (default: INFO). Case-insensitive.",
+    )
+    args = parser.parse_args()
+    log_level = args.level.upper()
+    logging.basicConfig(
+        level=log_level, format="%(asctime)s | %(levelname)-8s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+    )
+    logging.info("Configuring LeKiwi")
+    robot_config = LeKiwiMujocoConfig()
+    robot = LeKiwiMujoco(robot_config)
+
+    logging.info("Connecting LeKiwi")
+    robot.connect()
+
+    logging.info("Starting HostAgent")
+    host_config = LeKiwiHostConfig()
+    host = ZMQHandler(host_config)
+
+    last_cmd_time = time.time()
+    watchdog_active = False
+    logging.info("Waiting for commands...")
+    try:
+        # Business logic
+        start = time.perf_counter()
+        while robot.is_connected:
+            # while duration < host.connection_time_s:
+            loop_start_time = time.time()
+            try:
+                msg = host.zmq_cmd_socket.recv_string(zmq.NOBLOCK)
+                data = dict(json.loads(msg))
+                logging.debug("Received command: %s", data)
+                _action_sent = robot.send_action(data)
+                last_cmd_time = time.time()
+                watchdog_active = False
+            except zmq.Again:
+                if not watchdog_active:
+                    logging.warning("No command available")
+            except Exception as e:
+                logging.error("Message fetching failed: %s", e)
+
+            now = time.time()
+            if (now - last_cmd_time > host.watchdog_timeout_ms / 1000) and not watchdog_active:
+                logging.warning(
+                    f"Command not received for more than {host.watchdog_timeout_ms} milliseconds. Stopping the base."
+                )
+                watchdog_active = True
+                robot.stop_base()
+
+            last_observation = robot.get_observation()
+
+            # # Encode ndarrays to base64 strings
+            # for cam_key, _ in robot.cameras.items():
+            #     ret, buffer = cv2.imencode(
+            #         ".jpg", last_observation[cam_key], [int(cv2.IMWRITE_JPEG_QUALITY), 90]
+            #     )
+            #     if ret:
+            #         last_observation[cam_key] = base64.b64encode(buffer).decode("utf-8")
+            #     else:
+            #         last_observation[cam_key] = ""
+
+            # Send the observation to the remote agent
+            try:
+                host.zmq_observation_socket.send_string(json.dumps(last_observation), flags=zmq.NOBLOCK)
+            except zmq.Again:
+                logging.debug("Dropping observation, no client connected")
+
+            # Ensure a short sleep to avoid overloading the CPU.
+            elapsed = time.time() - loop_start_time
+
+            time.sleep(max(1 / host.max_loop_freq_hz - elapsed, 0))
+            time.perf_counter() - start
+        print("Cycle time reached.")
+
+    except KeyboardInterrupt:
+        print("Keyboard interrupt received. Exiting...")
+    finally:
+        print("Shutting down Lekiwi Host.")
+        robot.disconnect()
+        host.disconnect()
+
+    logging.info("Finished LeKiwi cleanly")
+
+
+if __name__ == "__main__":
+    main()