feat(lekiwi): de-couple classes + make it single-threaded
This commit is contained in:
Steven Palma
parent
66325b5a42
commit
2b100122f5
|
|
@ -1,224 +0,0 @@
|
|||
# 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 base64
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import zmq
|
||||
|
||||
from lerobot.common.robots.mobile_manipulator import LeKiwi
|
||||
|
||||
|
||||
def setup_zmq_sockets(config):
|
||||
context = zmq.Context()
|
||||
cmd_socket = context.socket(zmq.PULL)
|
||||
cmd_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
cmd_socket.bind(f"tcp://*:{config.port}")
|
||||
|
||||
video_socket = context.socket(zmq.PUSH)
|
||||
video_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
video_socket.bind(f"tcp://*:{config.video_port}")
|
||||
|
||||
return context, cmd_socket, video_socket
|
||||
|
||||
|
||||
def run_camera_capture(cameras, images_lock, latest_images_dict, stop_event):
|
||||
while not stop_event.is_set():
|
||||
local_dict = {}
|
||||
for name, cam in cameras.items():
|
||||
frame = cam.async_read()
|
||||
ret, buffer = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
||||
if ret:
|
||||
local_dict[name] = base64.b64encode(buffer).decode("utf-8")
|
||||
else:
|
||||
local_dict[name] = ""
|
||||
with images_lock:
|
||||
latest_images_dict.update(local_dict)
|
||||
time.sleep(0.01)
|
||||
|
||||
|
||||
def calibrate_follower_arm(motors_bus, calib_dir_str):
|
||||
"""
|
||||
Calibrates the follower arm. Attempts to load an existing calibration file;
|
||||
if not found, runs manual calibration and saves the result.
|
||||
"""
|
||||
calib_dir = Path(calib_dir_str)
|
||||
calib_dir.mkdir(parents=True, exist_ok=True)
|
||||
calib_file = calib_dir / "main_follower.json"
|
||||
try:
|
||||
from lerobot.common.motors.feetech.feetech_calibration import run_full_arm_calibration
|
||||
except ImportError:
|
||||
print("[WARNING] Calibration function not available. Skipping calibration.")
|
||||
return
|
||||
|
||||
if calib_file.exists():
|
||||
with open(calib_file) as f:
|
||||
calibration = json.load(f)
|
||||
print(f"[INFO] Loaded calibration from {calib_file}")
|
||||
else:
|
||||
print("[INFO] Calibration file not found. Running manual calibration...")
|
||||
calibration = run_full_arm_calibration(motors_bus, "lekiwi", "follower_arm", "follower")
|
||||
print(f"[INFO] Calibration complete. Saving to {calib_file}")
|
||||
with open(calib_file, "w") as f:
|
||||
json.dump(calibration, f)
|
||||
try:
|
||||
motors_bus.set_calibration(calibration)
|
||||
print("[INFO] Applied calibration for follower arm.")
|
||||
except Exception as e:
|
||||
print(f"[WARNING] Could not apply calibration: {e}")
|
||||
|
||||
|
||||
def run_lekiwi(robot_config):
|
||||
"""
|
||||
Runs the LeKiwi robot:
|
||||
- Sets up cameras and connects them.
|
||||
- Initializes the follower arm motors.
|
||||
- Calibrates the follower arm if necessary.
|
||||
- Creates ZeroMQ sockets for receiving commands and streaming observations.
|
||||
- Processes incoming commands (arm and wheel commands) and sends back sensor and camera data.
|
||||
"""
|
||||
# Import helper functions and classes
|
||||
from lerobot.common.cameras.utils import make_cameras_from_configs
|
||||
from lerobot.common.motors.feetech.feetech import FeetechMotorsBus, TorqueMode
|
||||
|
||||
# Initialize cameras from the robot configuration.
|
||||
cameras = make_cameras_from_configs(robot_config.cameras)
|
||||
for cam in cameras.values():
|
||||
cam.connect()
|
||||
|
||||
# Initialize the motors bus using the follower arm configuration.
|
||||
motor_config = robot_config.follower_arms.get("main")
|
||||
if motor_config is None:
|
||||
print("[ERROR] Follower arm 'main' configuration not found.")
|
||||
return
|
||||
motors_bus = FeetechMotorsBus(motor_config)
|
||||
motors_bus.connect()
|
||||
|
||||
# Calibrate the follower arm.
|
||||
calibrate_follower_arm(motors_bus, robot_config.calibration_dir)
|
||||
|
||||
# Create the LeKiwi robot instance.
|
||||
robot = LeKiwi(motors_bus)
|
||||
|
||||
# Define the expected arm motor IDs.
|
||||
arm_motor_ids = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
||||
|
||||
# Disable torque for each arm motor.
|
||||
for motor in arm_motor_ids:
|
||||
motors_bus.write("Torque_Enable", TorqueMode.DISABLED.value, motor)
|
||||
|
||||
# Set up ZeroMQ sockets.
|
||||
context, cmd_socket, video_socket = setup_zmq_sockets(robot_config)
|
||||
|
||||
# Start the camera capture thread.
|
||||
latest_images_dict = {}
|
||||
images_lock = threading.Lock()
|
||||
stop_event = threading.Event()
|
||||
cam_thread = threading.Thread(
|
||||
target=run_camera_capture, args=(cameras, images_lock, latest_images_dict, stop_event), daemon=True
|
||||
)
|
||||
cam_thread.start()
|
||||
|
||||
last_cmd_time = time.time()
|
||||
print("LeKiwi robot server started. Waiting for commands...")
|
||||
|
||||
try:
|
||||
while True:
|
||||
loop_start_time = time.time()
|
||||
|
||||
# Process incoming commands (non-blocking).
|
||||
while True:
|
||||
try:
|
||||
msg = cmd_socket.recv_string(zmq.NOBLOCK)
|
||||
except zmq.Again:
|
||||
break
|
||||
try:
|
||||
data = json.loads(msg)
|
||||
# Process arm position commands.
|
||||
if "arm_positions" in data:
|
||||
arm_positions = data["arm_positions"]
|
||||
if not isinstance(arm_positions, list):
|
||||
print(f"[ERROR] Invalid arm_positions: {arm_positions}")
|
||||
elif len(arm_positions) < len(arm_motor_ids):
|
||||
print(
|
||||
f"[WARNING] Received {len(arm_positions)} arm positions, expected {len(arm_motor_ids)}"
|
||||
)
|
||||
else:
|
||||
for motor, pos in zip(arm_motor_ids, arm_positions, strict=False):
|
||||
motors_bus.write("Goal_Position", pos, motor)
|
||||
# Process wheel (base) commands.
|
||||
if "raw_velocity" in data:
|
||||
raw_command = data["raw_velocity"]
|
||||
# Expect keys: "left_wheel", "back_wheel", "right_wheel".
|
||||
command_speeds = [
|
||||
int(raw_command.get("left_wheel", 0)),
|
||||
int(raw_command.get("back_wheel", 0)),
|
||||
int(raw_command.get("right_wheel", 0)),
|
||||
]
|
||||
robot.set_velocity(command_speeds)
|
||||
last_cmd_time = time.time()
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Parsing message failed: {e}")
|
||||
|
||||
# Watchdog: stop the robot if no command is received for over 0.5 seconds.
|
||||
now = time.time()
|
||||
if now - last_cmd_time > 0.5:
|
||||
robot.stop()
|
||||
last_cmd_time = now
|
||||
|
||||
# Read current wheel speeds from the robot.
|
||||
current_velocity = robot.read_velocity()
|
||||
|
||||
# Read the follower arm state from the motors bus.
|
||||
follower_arm_state = []
|
||||
for motor in arm_motor_ids:
|
||||
try:
|
||||
pos = motors_bus.read("Present_Position", motor)
|
||||
# Convert the position to a float (or use as is if already numeric).
|
||||
follower_arm_state.append(float(pos) if not isinstance(pos, (int, float)) else pos)
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Reading motor {motor} failed: {e}")
|
||||
|
||||
# Get the latest camera images.
|
||||
with images_lock:
|
||||
images_dict_copy = dict(latest_images_dict)
|
||||
|
||||
# Build the observation dictionary.
|
||||
observation = {
|
||||
"images": images_dict_copy,
|
||||
"present_speed": current_velocity,
|
||||
"follower_arm_state": follower_arm_state,
|
||||
}
|
||||
# Send the observation over the video socket.
|
||||
video_socket.send_string(json.dumps(observation))
|
||||
|
||||
# Ensure a short sleep to avoid overloading the CPU.
|
||||
elapsed = time.time() - loop_start_time
|
||||
time.sleep(
|
||||
max(0.033 - elapsed, 0)
|
||||
) # If robot jitters increase the sleep and monitor cpu load with `top` in cmd
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down LeKiwi server.")
|
||||
finally:
|
||||
stop_event.set()
|
||||
cam_thread.join()
|
||||
robot.stop()
|
||||
motors_bus.disconnect()
|
||||
cmd_socket.close()
|
||||
video_socket.close()
|
||||
context.term()
|
||||
|
|
@ -1,692 +0,0 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import zmq
|
||||
|
||||
from lerobot.common.cameras.utils import make_cameras_from_configs
|
||||
from lerobot.common.errors import DeviceNotConnectedError
|
||||
from lerobot.common.motors.feetech.feetech import TorqueMode
|
||||
from lerobot.common.motors.feetech.feetech_calibration import run_full_arm_calibration
|
||||
from lerobot.common.motors.motors_bus import MotorsBus
|
||||
from lerobot.common.motors.utils import make_motors_buses_from_configs
|
||||
from lerobot.common.robots.lekiwi.configuration_lekiwi import LeKiwiRobotConfig
|
||||
from lerobot.common.robots.utils import get_arm_id
|
||||
|
||||
PYNPUT_AVAILABLE = True
|
||||
try:
|
||||
# Only import if there's a valid X server or if we're not on a Pi
|
||||
if ("DISPLAY" not in os.environ) and ("linux" in sys.platform):
|
||||
print("No DISPLAY set. Skipping pynput import.")
|
||||
raise ImportError("pynput blocked intentionally due to no display.")
|
||||
|
||||
from pynput import keyboard
|
||||
except ImportError:
|
||||
keyboard = None
|
||||
PYNPUT_AVAILABLE = False
|
||||
except Exception as e:
|
||||
keyboard = None
|
||||
PYNPUT_AVAILABLE = False
|
||||
print(f"Could not import pynput: {e}")
|
||||
|
||||
|
||||
class MobileManipulator:
|
||||
"""
|
||||
MobileManipulator is a class for connecting to and controlling a remote mobile manipulator robot.
|
||||
The robot includes a three omniwheel mobile base and a remote follower arm.
|
||||
The leader arm is connected locally (on the laptop) and its joint positions are recorded and then
|
||||
forwarded to the remote follower arm (after applying a safety clamp).
|
||||
In parallel, keyboard teleoperation is used to generate raw velocity commands for the wheels.
|
||||
"""
|
||||
|
||||
def __init__(self, config: LeKiwiRobotConfig):
|
||||
"""
|
||||
Expected keys in config:
|
||||
- ip, port, video_port for the remote connection.
|
||||
- calibration_dir, leader_arms, follower_arms, max_relative_target, etc.
|
||||
"""
|
||||
self.robot_type = config.type
|
||||
self.config = config
|
||||
self.remote_ip = config.ip
|
||||
self.remote_port = config.port
|
||||
self.remote_port_video = config.video_port
|
||||
self.calibration_dir = Path(self.config.calibration_dir)
|
||||
self.logs = {}
|
||||
|
||||
self.teleop_keys = self.config.teleop_keys
|
||||
|
||||
# For teleoperation, the leader arm (local) is used to record the desired arm pose.
|
||||
self.leader_arms = make_motors_buses_from_configs(self.config.leader_arms)
|
||||
|
||||
self.follower_arms = make_motors_buses_from_configs(self.config.follower_arms)
|
||||
|
||||
self.cameras = make_cameras_from_configs(self.config.cameras)
|
||||
|
||||
self.is_connected = False
|
||||
|
||||
self.last_frames = {}
|
||||
self.last_present_speed = {}
|
||||
self.last_remote_arm_state = torch.zeros(6, dtype=torch.float32)
|
||||
|
||||
# Define three speed levels and a current index
|
||||
self.speed_levels = [
|
||||
{"xy": 0.1, "theta": 30}, # slow
|
||||
{"xy": 0.2, "theta": 60}, # medium
|
||||
{"xy": 0.3, "theta": 90}, # fast
|
||||
]
|
||||
self.speed_index = 0 # Start at slow
|
||||
|
||||
# ZeroMQ context and sockets.
|
||||
self.context = None
|
||||
self.cmd_socket = None
|
||||
self.video_socket = None
|
||||
|
||||
# Keyboard state for base teleoperation.
|
||||
self.running = True
|
||||
self.pressed_keys = {
|
||||
"forward": False,
|
||||
"backward": False,
|
||||
"left": False,
|
||||
"right": False,
|
||||
"rotate_left": False,
|
||||
"rotate_right": False,
|
||||
}
|
||||
|
||||
if PYNPUT_AVAILABLE:
|
||||
print("pynput is available - enabling local keyboard listener.")
|
||||
self.listener = keyboard.Listener(
|
||||
on_press=self.on_press,
|
||||
on_release=self.on_release,
|
||||
)
|
||||
self.listener.start()
|
||||
else:
|
||||
print("pynput not available - skipping local keyboard listener.")
|
||||
self.listener = None
|
||||
|
||||
def get_motor_names(self, arms: dict[str, MotorsBus]) -> list:
|
||||
return [f"{arm}_{motor}" for arm, bus in arms.items() for motor in bus.motors]
|
||||
|
||||
@property
|
||||
def camera_features(self) -> dict:
|
||||
cam_ft = {}
|
||||
for cam_key, cam in self.cameras.items():
|
||||
key = f"observation.images.{cam_key}"
|
||||
cam_ft[key] = {
|
||||
"shape": (cam.height, cam.width, cam.channels),
|
||||
"names": ["height", "width", "channels"],
|
||||
"info": None,
|
||||
}
|
||||
return cam_ft
|
||||
|
||||
@property
|
||||
def motor_features(self) -> dict:
|
||||
follower_arm_names = [
|
||||
"shoulder_pan",
|
||||
"shoulder_lift",
|
||||
"elbow_flex",
|
||||
"wrist_flex",
|
||||
"wrist_roll",
|
||||
"gripper",
|
||||
]
|
||||
observations = ["x_mm", "y_mm", "theta"]
|
||||
combined_names = follower_arm_names + observations
|
||||
return {
|
||||
"action": {
|
||||
"dtype": "float32",
|
||||
"shape": (len(combined_names),),
|
||||
"names": combined_names,
|
||||
},
|
||||
"observation.state": {
|
||||
"dtype": "float32",
|
||||
"shape": (len(combined_names),),
|
||||
"names": combined_names,
|
||||
},
|
||||
}
|
||||
|
||||
@property
|
||||
def features(self):
|
||||
return {**self.motor_features, **self.camera_features}
|
||||
|
||||
@property
|
||||
def has_camera(self):
|
||||
return len(self.cameras) > 0
|
||||
|
||||
@property
|
||||
def num_cameras(self):
|
||||
return len(self.cameras)
|
||||
|
||||
@property
|
||||
def available_arms(self):
|
||||
available = []
|
||||
for name in self.leader_arms:
|
||||
available.append(get_arm_id(name, "leader"))
|
||||
for name in self.follower_arms:
|
||||
available.append(get_arm_id(name, "follower"))
|
||||
return available
|
||||
|
||||
def on_press(self, key):
|
||||
try:
|
||||
# Movement
|
||||
if key.char == self.teleop_keys["forward"]:
|
||||
self.pressed_keys["forward"] = True
|
||||
elif key.char == self.teleop_keys["backward"]:
|
||||
self.pressed_keys["backward"] = True
|
||||
elif key.char == self.teleop_keys["left"]:
|
||||
self.pressed_keys["left"] = True
|
||||
elif key.char == self.teleop_keys["right"]:
|
||||
self.pressed_keys["right"] = True
|
||||
elif key.char == self.teleop_keys["rotate_left"]:
|
||||
self.pressed_keys["rotate_left"] = True
|
||||
elif key.char == self.teleop_keys["rotate_right"]:
|
||||
self.pressed_keys["rotate_right"] = True
|
||||
|
||||
# Quit teleoperation
|
||||
elif key.char == self.teleop_keys["quit"]:
|
||||
self.running = False
|
||||
return False
|
||||
|
||||
# Speed control
|
||||
elif key.char == self.teleop_keys["speed_up"]:
|
||||
self.speed_index = min(self.speed_index + 1, 2)
|
||||
print(f"Speed index increased to {self.speed_index}")
|
||||
elif key.char == self.teleop_keys["speed_down"]:
|
||||
self.speed_index = max(self.speed_index - 1, 0)
|
||||
print(f"Speed index decreased to {self.speed_index}")
|
||||
|
||||
except AttributeError:
|
||||
# e.g., if key is special like Key.esc
|
||||
if key == keyboard.Key.esc:
|
||||
self.running = False
|
||||
return False
|
||||
|
||||
def on_release(self, key):
|
||||
try:
|
||||
if hasattr(key, "char"):
|
||||
if key.char == self.teleop_keys["forward"]:
|
||||
self.pressed_keys["forward"] = False
|
||||
elif key.char == self.teleop_keys["backward"]:
|
||||
self.pressed_keys["backward"] = False
|
||||
elif key.char == self.teleop_keys["left"]:
|
||||
self.pressed_keys["left"] = False
|
||||
elif key.char == self.teleop_keys["right"]:
|
||||
self.pressed_keys["right"] = False
|
||||
elif key.char == self.teleop_keys["rotate_left"]:
|
||||
self.pressed_keys["rotate_left"] = False
|
||||
elif key.char == self.teleop_keys["rotate_right"]:
|
||||
self.pressed_keys["rotate_right"] = False
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
def connect(self):
|
||||
if not self.leader_arms:
|
||||
raise ValueError("MobileManipulator has no leader arm to connect.")
|
||||
for name in self.leader_arms:
|
||||
print(f"Connecting {name} leader arm.")
|
||||
self.calibrate_leader()
|
||||
|
||||
# Set up ZeroMQ sockets to communicate with the remote mobile robot.
|
||||
self.context = zmq.Context()
|
||||
self.cmd_socket = self.context.socket(zmq.PUSH)
|
||||
connection_string = f"tcp://{self.remote_ip}:{self.remote_port}"
|
||||
self.cmd_socket.connect(connection_string)
|
||||
self.cmd_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
self.video_socket = self.context.socket(zmq.PULL)
|
||||
video_connection = f"tcp://{self.remote_ip}:{self.remote_port_video}"
|
||||
self.video_socket.connect(video_connection)
|
||||
self.video_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
print(
|
||||
f"[INFO] Connected to remote robot at {connection_string} and video stream at {video_connection}."
|
||||
)
|
||||
self.is_connected = True
|
||||
|
||||
def load_or_run_calibration_(self, name, arm, arm_type):
|
||||
arm_id = get_arm_id(name, arm_type)
|
||||
arm_calib_path = self.calibration_dir / f"{arm_id}.json"
|
||||
|
||||
if arm_calib_path.exists():
|
||||
with open(arm_calib_path) as f:
|
||||
calibration = json.load(f)
|
||||
else:
|
||||
print(f"Missing calibration file '{arm_calib_path}'")
|
||||
calibration = run_full_arm_calibration(arm, self.robot_type, name, arm_type)
|
||||
print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
|
||||
arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(arm_calib_path, "w") as f:
|
||||
json.dump(calibration, f)
|
||||
|
||||
return calibration
|
||||
|
||||
def calibrate_leader(self):
|
||||
for name, arm in self.leader_arms.items():
|
||||
# Connect the bus
|
||||
arm.connect()
|
||||
|
||||
# Disable torque on all motors
|
||||
for motor_id in arm.motors:
|
||||
arm.write("Torque_Enable", TorqueMode.DISABLED.value, motor_id)
|
||||
|
||||
# Now run calibration
|
||||
calibration = self.load_or_run_calibration_(name, arm, "leader")
|
||||
arm.set_calibration(calibration)
|
||||
|
||||
def calibrate_follower(self):
|
||||
for name, bus in self.follower_arms.items():
|
||||
bus.connect()
|
||||
|
||||
# Disable torque on all motors
|
||||
for motor_id in bus.motors:
|
||||
bus.write("Torque_Enable", 0, motor_id)
|
||||
|
||||
# Then filter out wheels
|
||||
arm_only_dict = {k: v for k, v in bus.motors.items() if not k.startswith("wheel_")}
|
||||
if not arm_only_dict:
|
||||
continue
|
||||
|
||||
original_motors = bus.motors
|
||||
bus.motors = arm_only_dict
|
||||
|
||||
calibration = self.load_or_run_calibration_(name, bus, "follower")
|
||||
bus.set_calibration(calibration)
|
||||
|
||||
bus.motors = original_motors
|
||||
|
||||
def _get_data(self):
|
||||
"""
|
||||
Polls the video socket for up to 15 ms. If data arrives, decode only
|
||||
the *latest* message, returning frames, speed, and arm state. If
|
||||
nothing arrives for any field, use the last known values.
|
||||
"""
|
||||
frames = {}
|
||||
present_speed = {}
|
||||
remote_arm_state_tensor = torch.zeros(6, dtype=torch.float32)
|
||||
|
||||
# Poll up to 15 ms
|
||||
poller = zmq.Poller()
|
||||
poller.register(self.video_socket, zmq.POLLIN)
|
||||
socks = dict(poller.poll(15))
|
||||
if self.video_socket not in socks or socks[self.video_socket] != zmq.POLLIN:
|
||||
# No new data arrived → reuse ALL old data
|
||||
return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
|
||||
|
||||
# Drain all messages, keep only the last
|
||||
last_msg = None
|
||||
while True:
|
||||
try:
|
||||
obs_string = self.video_socket.recv_string(zmq.NOBLOCK)
|
||||
last_msg = obs_string
|
||||
except zmq.Again:
|
||||
break
|
||||
|
||||
if not last_msg:
|
||||
# No new message → also reuse old
|
||||
return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
|
||||
|
||||
# Decode only the final message
|
||||
try:
|
||||
observation = json.loads(last_msg)
|
||||
|
||||
images_dict = observation.get("images", {})
|
||||
new_speed = observation.get("present_speed", {})
|
||||
new_arm_state = observation.get("follower_arm_state", None)
|
||||
|
||||
# Convert images
|
||||
for cam_name, image_b64 in images_dict.items():
|
||||
if image_b64:
|
||||
jpg_data = base64.b64decode(image_b64)
|
||||
np_arr = np.frombuffer(jpg_data, dtype=np.uint8)
|
||||
frame_candidate = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
||||
if frame_candidate is not None:
|
||||
frames[cam_name] = frame_candidate
|
||||
|
||||
# If remote_arm_state is None and frames is None there is no message then use the previous message
|
||||
if new_arm_state is not None and frames is not None:
|
||||
self.last_frames = frames
|
||||
|
||||
remote_arm_state_tensor = torch.tensor(new_arm_state, dtype=torch.float32)
|
||||
self.last_remote_arm_state = remote_arm_state_tensor
|
||||
|
||||
present_speed = new_speed
|
||||
self.last_present_speed = new_speed
|
||||
else:
|
||||
frames = self.last_frames
|
||||
|
||||
remote_arm_state_tensor = self.last_remote_arm_state
|
||||
|
||||
present_speed = self.last_present_speed
|
||||
|
||||
except Exception as e:
|
||||
print(f"[DEBUG] Error decoding video message: {e}")
|
||||
# If decode fails, fall back to old data
|
||||
return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
|
||||
|
||||
return frames, present_speed, remote_arm_state_tensor
|
||||
|
||||
def _process_present_speed(self, present_speed: dict) -> torch.Tensor:
|
||||
state_tensor = torch.zeros(3, dtype=torch.int32)
|
||||
if present_speed:
|
||||
decoded = {key: MobileManipulator.raw_to_degps(value) for key, value in present_speed.items()}
|
||||
if "1" in decoded:
|
||||
state_tensor[0] = decoded["1"]
|
||||
if "2" in decoded:
|
||||
state_tensor[1] = decoded["2"]
|
||||
if "3" in decoded:
|
||||
state_tensor[2] = decoded["3"]
|
||||
return state_tensor
|
||||
|
||||
def teleop_step(
|
||||
self, record_data: bool = False
|
||||
) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError("MobileManipulator is not connected. Run `connect()` first.")
|
||||
|
||||
speed_setting = self.speed_levels[self.speed_index]
|
||||
xy_speed = speed_setting["xy"] # e.g. 0.1, 0.25, or 0.4
|
||||
theta_speed = speed_setting["theta"] # e.g. 30, 60, or 90
|
||||
|
||||
# Prepare to assign the position of the leader to the follower
|
||||
arm_positions = []
|
||||
for name in self.leader_arms:
|
||||
pos = self.leader_arms[name].read("Present_Position")
|
||||
pos_tensor = torch.from_numpy(pos).float()
|
||||
# Instead of pos_tensor.item(), use tolist() to convert the entire tensor to a list
|
||||
arm_positions.extend(pos_tensor.tolist())
|
||||
|
||||
# (The rest of your code for generating wheel commands remains unchanged)
|
||||
x_cmd = 0.0 # m/s forward/backward
|
||||
y_cmd = 0.0 # m/s lateral
|
||||
theta_cmd = 0.0 # deg/s rotation
|
||||
if self.pressed_keys["forward"]:
|
||||
x_cmd += xy_speed
|
||||
if self.pressed_keys["backward"]:
|
||||
x_cmd -= xy_speed
|
||||
if self.pressed_keys["left"]:
|
||||
y_cmd += xy_speed
|
||||
if self.pressed_keys["right"]:
|
||||
y_cmd -= xy_speed
|
||||
if self.pressed_keys["rotate_left"]:
|
||||
theta_cmd += theta_speed
|
||||
if self.pressed_keys["rotate_right"]:
|
||||
theta_cmd -= theta_speed
|
||||
|
||||
wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
|
||||
|
||||
message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions}
|
||||
self.cmd_socket.send_string(json.dumps(message))
|
||||
|
||||
if not record_data:
|
||||
return
|
||||
|
||||
obs_dict = self.capture_observation()
|
||||
|
||||
arm_state_tensor = torch.tensor(arm_positions, dtype=torch.float32)
|
||||
|
||||
wheel_velocity_tuple = self.wheel_raw_to_body(wheel_commands)
|
||||
wheel_velocity_mm = (
|
||||
wheel_velocity_tuple[0] * 1000.0,
|
||||
wheel_velocity_tuple[1] * 1000.0,
|
||||
wheel_velocity_tuple[2],
|
||||
)
|
||||
wheel_tensor = torch.tensor(wheel_velocity_mm, dtype=torch.float32)
|
||||
action_tensor = torch.cat([arm_state_tensor, wheel_tensor])
|
||||
action_dict = {"action": action_tensor}
|
||||
|
||||
return obs_dict, action_dict
|
||||
|
||||
def capture_observation(self) -> dict:
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError("Not connected. Run `connect()` first.")
|
||||
|
||||
frames, present_speed, remote_arm_state_tensor = self._get_data()
|
||||
|
||||
body_state = self.wheel_raw_to_body(present_speed)
|
||||
|
||||
body_state_mm = (body_state[0] * 1000.0, body_state[1] * 1000.0, body_state[2]) # Convert x,y to mm/s
|
||||
wheel_state_tensor = torch.tensor(body_state_mm, dtype=torch.float32)
|
||||
combined_state_tensor = torch.cat((remote_arm_state_tensor, wheel_state_tensor), dim=0)
|
||||
|
||||
obs_dict = {"observation.state": combined_state_tensor}
|
||||
|
||||
# Loop over each configured camera
|
||||
for cam_name, cam in self.cameras.items():
|
||||
frame = frames.get(cam_name, None)
|
||||
if frame is None:
|
||||
# Create a black image using the camera's configured width, height, and channels
|
||||
frame = np.zeros((cam.height, cam.width, cam.channels), dtype=np.uint8)
|
||||
obs_dict[f"observation.images.{cam_name}"] = torch.from_numpy(frame)
|
||||
|
||||
return obs_dict
|
||||
|
||||
def send_action(self, action: torch.Tensor) -> torch.Tensor:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError("Not connected. Run `connect()` first.")
|
||||
|
||||
# Ensure the action tensor has at least 9 elements:
|
||||
# - First 6: arm positions.
|
||||
# - Last 3: base commands.
|
||||
if action.numel() < 9:
|
||||
# Pad with zeros if there are not enough elements.
|
||||
padded = torch.zeros(9, dtype=action.dtype)
|
||||
padded[: action.numel()] = action
|
||||
action = padded
|
||||
|
||||
# Extract arm and base actions.
|
||||
arm_actions = action[:6].flatten()
|
||||
base_actions = action[6:].flatten()
|
||||
|
||||
x_cmd_mm = base_actions[0].item() # mm/s
|
||||
y_cmd_mm = base_actions[1].item() # mm/s
|
||||
theta_cmd = base_actions[2].item() # deg/s
|
||||
|
||||
# Convert mm/s to m/s for the kinematics calculations.
|
||||
x_cmd = x_cmd_mm / 1000.0 # m/s
|
||||
y_cmd = y_cmd_mm / 1000.0 # m/s
|
||||
|
||||
# Compute wheel commands from body commands.
|
||||
wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
|
||||
|
||||
arm_positions_list = arm_actions.tolist()
|
||||
|
||||
message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions_list}
|
||||
self.cmd_socket.send_string(json.dumps(message))
|
||||
|
||||
return action
|
||||
|
||||
def print_logs(self):
|
||||
pass
|
||||
|
||||
def disconnect(self):
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError("Not connected.")
|
||||
if self.cmd_socket:
|
||||
stop_cmd = {
|
||||
"raw_velocity": {"left_wheel": 0, "back_wheel": 0, "right_wheel": 0},
|
||||
"arm_positions": {},
|
||||
}
|
||||
self.cmd_socket.send_string(json.dumps(stop_cmd))
|
||||
self.cmd_socket.close()
|
||||
if self.video_socket:
|
||||
self.video_socket.close()
|
||||
if self.context:
|
||||
self.context.term()
|
||||
if PYNPUT_AVAILABLE:
|
||||
self.listener.stop()
|
||||
self.is_connected = False
|
||||
print("[INFO] Disconnected from remote robot.")
|
||||
|
||||
def __del__(self):
|
||||
if getattr(self, "is_connected", False):
|
||||
self.disconnect()
|
||||
if PYNPUT_AVAILABLE:
|
||||
self.listener.stop()
|
||||
|
||||
@staticmethod
|
||||
def degps_to_raw(degps: float) -> int:
|
||||
steps_per_deg = 4096.0 / 360.0
|
||||
speed_in_steps = abs(degps) * steps_per_deg
|
||||
speed_int = int(round(speed_in_steps))
|
||||
if speed_int > 0x7FFF:
|
||||
speed_int = 0x7FFF
|
||||
if degps < 0:
|
||||
return speed_int | 0x8000
|
||||
else:
|
||||
return speed_int & 0x7FFF
|
||||
|
||||
@staticmethod
|
||||
def raw_to_degps(raw_speed: int) -> float:
|
||||
steps_per_deg = 4096.0 / 360.0
|
||||
magnitude = raw_speed & 0x7FFF
|
||||
degps = magnitude / steps_per_deg
|
||||
if raw_speed & 0x8000:
|
||||
degps = -degps
|
||||
return degps
|
||||
|
||||
def body_to_wheel_raw(
|
||||
self,
|
||||
x_cmd: float,
|
||||
y_cmd: float,
|
||||
theta_cmd: float,
|
||||
wheel_radius: float = 0.05,
|
||||
base_radius: float = 0.125,
|
||||
max_raw: int = 3000,
|
||||
) -> dict:
|
||||
"""
|
||||
Convert desired body-frame velocities into wheel raw commands.
|
||||
|
||||
Parameters:
|
||||
x_cmd : Linear velocity in x (m/s).
|
||||
y_cmd : Linear velocity in y (m/s).
|
||||
theta_cmd : Rotational velocity (deg/s).
|
||||
wheel_radius: Radius of each wheel (meters).
|
||||
base_radius : Distance from the center of rotation to each wheel (meters).
|
||||
max_raw : Maximum allowed raw command (ticks) per wheel.
|
||||
|
||||
Returns:
|
||||
A dictionary with wheel raw commands:
|
||||
{"left_wheel": value, "back_wheel": value, "right_wheel": value}.
|
||||
|
||||
Notes:
|
||||
- Internally, the method converts theta_cmd to rad/s for the kinematics.
|
||||
- The raw command is computed from the wheels angular speed in deg/s
|
||||
using degps_to_raw(). If any command exceeds max_raw, all commands
|
||||
are scaled down proportionally.
|
||||
"""
|
||||
# Convert rotational velocity from deg/s to rad/s.
|
||||
theta_rad = theta_cmd * (np.pi / 180.0)
|
||||
# Create the body velocity vector [x, y, theta_rad].
|
||||
velocity_vector = np.array([x_cmd, y_cmd, theta_rad])
|
||||
|
||||
# Define the wheel mounting angles with a -90° offset.
|
||||
angles = np.radians(np.array([240, 120, 0]) - 90)
|
||||
# Build the kinematic matrix: each row maps body velocities to a wheel’s linear speed.
|
||||
# The third column (base_radius) accounts for the effect of rotation.
|
||||
m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
|
||||
|
||||
# Compute each wheel’s linear speed (m/s) and then its angular speed (rad/s).
|
||||
wheel_linear_speeds = m.dot(velocity_vector)
|
||||
wheel_angular_speeds = wheel_linear_speeds / wheel_radius
|
||||
|
||||
# Convert wheel angular speeds from rad/s to deg/s.
|
||||
wheel_degps = wheel_angular_speeds * (180.0 / np.pi)
|
||||
|
||||
# Scaling
|
||||
steps_per_deg = 4096.0 / 360.0
|
||||
raw_floats = [abs(degps) * steps_per_deg for degps in wheel_degps]
|
||||
max_raw_computed = max(raw_floats)
|
||||
if max_raw_computed > max_raw:
|
||||
scale = max_raw / max_raw_computed
|
||||
wheel_degps = wheel_degps * scale
|
||||
|
||||
# Convert each wheel’s angular speed (deg/s) to a raw integer.
|
||||
wheel_raw = [MobileManipulator.degps_to_raw(deg) for deg in wheel_degps]
|
||||
|
||||
return {"left_wheel": wheel_raw[0], "back_wheel": wheel_raw[1], "right_wheel": wheel_raw[2]}
|
||||
|
||||
def wheel_raw_to_body(
|
||||
self, wheel_raw: dict, wheel_radius: float = 0.05, base_radius: float = 0.125
|
||||
) -> tuple:
|
||||
"""
|
||||
Convert wheel raw command feedback back into body-frame velocities.
|
||||
|
||||
Parameters:
|
||||
wheel_raw : Dictionary with raw wheel commands (keys: "left_wheel", "back_wheel", "right_wheel").
|
||||
wheel_radius: Radius of each wheel (meters).
|
||||
base_radius : Distance from the robot center to each wheel (meters).
|
||||
|
||||
Returns:
|
||||
A tuple (x_cmd, y_cmd, theta_cmd) where:
|
||||
x_cmd : Linear velocity in x (m/s).
|
||||
y_cmd : Linear velocity in y (m/s).
|
||||
theta_cmd : Rotational velocity in deg/s.
|
||||
"""
|
||||
# Extract the raw values in order.
|
||||
raw_list = [
|
||||
int(wheel_raw.get("left_wheel", 0)),
|
||||
int(wheel_raw.get("back_wheel", 0)),
|
||||
int(wheel_raw.get("right_wheel", 0)),
|
||||
]
|
||||
|
||||
# Convert each raw command back to an angular speed in deg/s.
|
||||
wheel_degps = np.array([MobileManipulator.raw_to_degps(r) for r in raw_list])
|
||||
# Convert from deg/s to rad/s.
|
||||
wheel_radps = wheel_degps * (np.pi / 180.0)
|
||||
# Compute each wheel’s linear speed (m/s) from its angular speed.
|
||||
wheel_linear_speeds = wheel_radps * wheel_radius
|
||||
|
||||
# Define the wheel mounting angles with a -90° offset.
|
||||
angles = np.radians(np.array([240, 120, 0]) - 90)
|
||||
m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
|
||||
|
||||
# Solve the inverse kinematics: body_velocity = M⁻¹ · wheel_linear_speeds.
|
||||
m_inv = np.linalg.inv(m)
|
||||
velocity_vector = m_inv.dot(wheel_linear_speeds)
|
||||
x_cmd, y_cmd, theta_rad = velocity_vector
|
||||
theta_cmd = theta_rad * (180.0 / np.pi)
|
||||
return (x_cmd, y_cmd, theta_cmd)
|
||||
|
||||
|
||||
class LeKiwi:
|
||||
def __init__(self, motor_bus):
|
||||
"""
|
||||
Initializes the LeKiwi with Feetech motors bus.
|
||||
"""
|
||||
self.motor_bus = motor_bus
|
||||
self.motor_ids = ["left_wheel", "back_wheel", "right_wheel"]
|
||||
|
||||
# Initialize motors in velocity mode.
|
||||
self.motor_bus.write("Lock", 0)
|
||||
self.motor_bus.write("Mode", [1, 1, 1], self.motor_ids)
|
||||
self.motor_bus.write("Lock", 1)
|
||||
print("Motors set to velocity mode.")
|
||||
|
||||
def read_velocity(self):
|
||||
"""
|
||||
Reads the raw speeds for all wheels. Returns a dictionary with motor names:
|
||||
"""
|
||||
raw_speeds = self.motor_bus.read("Present_Velocity", self.motor_ids)
|
||||
return {
|
||||
"left_wheel": int(raw_speeds[0]),
|
||||
"back_wheel": int(raw_speeds[1]),
|
||||
"right_wheel": int(raw_speeds[2]),
|
||||
}
|
||||
|
||||
def set_velocity(self, command_speeds):
|
||||
"""
|
||||
Sends raw velocity commands (16-bit encoded values) directly to the motor bus.
|
||||
The order of speeds must correspond to self.motor_ids.
|
||||
"""
|
||||
self.motor_bus.write("Goal_Velocity", command_speeds, self.motor_ids)
|
||||
|
||||
def stop(self):
|
||||
"""Stops the robot by setting all motor speeds to zero."""
|
||||
self.motor_bus.write("Goal_Velocity", [0, 0, 0], self.motor_ids)
|
||||
print("Motors stopped.")
|
||||
|
|
@ -8,9 +8,6 @@ from lerobot.common.robots.config import RobotConfig
|
|||
@RobotConfig.register_subclass("lekiwi")
|
||||
@dataclass
|
||||
class LeKiwiRobotConfig(RobotConfig):
|
||||
port_zmq_cmd: int = 5555
|
||||
port_zmq_observations: int = 5556
|
||||
|
||||
id = "lekiwi"
|
||||
|
||||
cameras: dict[str, CameraConfig] = field(
|
||||
|
|
|
|||
|
|
@ -138,7 +138,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
# Consider moving these static functions out of the class
|
||||
# Copied from robot_lekiwi MobileManipulator class
|
||||
@staticmethod
|
||||
def degps_to_raw(degps: float) -> int:
|
||||
def _degps_to_raw(degps: float) -> int:
|
||||
steps_per_deg = 4096.0 / 360.0
|
||||
speed_in_steps = abs(degps) * steps_per_deg
|
||||
speed_int = int(round(speed_in_steps))
|
||||
|
|
@ -151,7 +151,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
|
||||
# Copied from robot_lekiwi MobileManipulator class
|
||||
@staticmethod
|
||||
def raw_to_degps(raw_speed: int) -> float:
|
||||
def _raw_to_degps(raw_speed: int) -> float:
|
||||
steps_per_deg = 4096.0 / 360.0
|
||||
magnitude = raw_speed & 0x7FFF
|
||||
degps = magnitude / steps_per_deg
|
||||
|
|
@ -160,7 +160,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
return degps
|
||||
|
||||
# Copied from robot_lekiwi MobileManipulator class
|
||||
def body_to_wheel_raw(
|
||||
def _body_to_wheel_raw(
|
||||
self,
|
||||
x_cmd: float,
|
||||
y_cmd: float,
|
||||
|
|
@ -187,7 +187,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
Notes:
|
||||
- Internally, the method converts theta_cmd to rad/s for the kinematics.
|
||||
- The raw command is computed from the wheels angular speed in deg/s
|
||||
using degps_to_raw(). If any command exceeds max_raw, all commands
|
||||
using _degps_to_raw(). If any command exceeds max_raw, all commands
|
||||
are scaled down proportionally.
|
||||
"""
|
||||
# Convert rotational velocity from deg/s to rad/s.
|
||||
|
|
@ -217,12 +217,12 @@ class DaemonLeKiwiRobot(Robot):
|
|||
wheel_degps = wheel_degps * scale
|
||||
|
||||
# Convert each wheel’s angular speed (deg/s) to a raw integer.
|
||||
wheel_raw = [DaemonLeKiwiRobot.degps_to_raw(deg) for deg in wheel_degps]
|
||||
wheel_raw = [DaemonLeKiwiRobot._degps_to_raw(deg) for deg in wheel_degps]
|
||||
|
||||
return {"left_wheel": wheel_raw[0], "back_wheel": wheel_raw[1], "right_wheel": wheel_raw[2]}
|
||||
|
||||
# Copied from robot_lekiwi MobileManipulator class
|
||||
def wheel_raw_to_body(
|
||||
def _wheel_raw_to_body(
|
||||
self, wheel_raw: np.array, wheel_radius: float = 0.05, base_radius: float = 0.125
|
||||
) -> tuple:
|
||||
"""
|
||||
|
|
@ -241,7 +241,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
"""
|
||||
|
||||
# Convert each raw command back to an angular speed in deg/s.
|
||||
wheel_degps = np.array([DaemonLeKiwiRobot.raw_to_degps(int(r)) for r in wheel_raw])
|
||||
wheel_degps = np.array([DaemonLeKiwiRobot._raw_to_degps(int(r)) for r in wheel_raw])
|
||||
# Convert from deg/s to rad/s.
|
||||
wheel_radps = wheel_degps * (np.pi / 180.0)
|
||||
# Compute each wheel’s linear speed (m/s) from its angular speed.
|
||||
|
|
@ -259,7 +259,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
return (x_cmd, y_cmd, theta_cmd)
|
||||
|
||||
# TODO(Steven): This is flaky, for example, if we received a state but failed decoding the image, we will not update any value
|
||||
def get_data(self):
|
||||
def _get_data(self):
|
||||
# Copied from robot_lekiwi.py
|
||||
"""Polls the video socket for up to 15 ms. If data arrives, decode only
|
||||
the *latest* message, returning frames, speed, and arm state. If
|
||||
|
|
@ -348,8 +348,8 @@ class DaemonLeKiwiRobot(Robot):
|
|||
|
||||
obs_dict = {}
|
||||
|
||||
frames, present_speed, remote_arm_state_tensor = self.get_data()
|
||||
body_state = self.wheel_raw_to_body(present_speed)
|
||||
frames, present_speed, remote_arm_state_tensor = self._get_data()
|
||||
body_state = self._wheel_raw_to_body(present_speed)
|
||||
body_state_mm = (body_state[0] * 1000.0, body_state[1] * 1000.0, body_state[2]) # Convert x,y to mm/s
|
||||
wheel_state_tensor = torch.tensor(body_state_mm, dtype=torch.float32)
|
||||
combined_state_tensor = torch.cat((remote_arm_state_tensor, wheel_state_tensor), dim=0)
|
||||
|
|
@ -366,7 +366,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
|
||||
return obs_dict
|
||||
|
||||
def from_keyboard_to_wheel_action(self, pressed_keys: np.ndarray):
|
||||
def _from_keyboard_to_wheel_action(self, pressed_keys: np.ndarray):
|
||||
# Speed control
|
||||
if self.teleop_keys["speed_up"] in pressed_keys:
|
||||
self.speed_index = min(self.speed_index + 1, 2)
|
||||
|
|
@ -393,7 +393,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
if self.teleop_keys["rotate_right"] in pressed_keys:
|
||||
theta_cmd -= theta_speed
|
||||
|
||||
return self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
|
||||
return self._body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
|
||||
|
||||
# TODO(Steven): This assumes this call is always called from a keyboard teleop command
|
||||
# TODO(Steven): Doing this mapping in here adds latecy between send_action and movement from the user perspective.
|
||||
|
|
@ -439,7 +439,7 @@ class DaemonLeKiwiRobot(Robot):
|
|||
goal_pos[:6] = action[:6]
|
||||
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()]
|
||||
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
|
||||
|
|
|
|||
|
|
@ -1,8 +1,35 @@
|
|||
import json
|
||||
import logging
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
|
||||
from lerobot.common.constants import OBS_STATE
|
||||
from lerobot.common.robots.lekiwi.configuration_lekiwi import LeKiwiRobotConfig
|
||||
from lerobot.common.robots.lekiwi.lekiwi_robot import LeKiwiRobot
|
||||
|
||||
# Network Configuration
|
||||
PORT_ZMQ_CMD: int = 5555
|
||||
PORT_ZMQ_OBSERVATIONS: int = 5556
|
||||
|
||||
|
||||
class RemoteAgent:
|
||||
def __init__(self):
|
||||
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://*:{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://*:{PORT_ZMQ_OBSERVATIONS}")
|
||||
|
||||
def disconnect(self):
|
||||
self.zmq_observation_socket.close()
|
||||
self.zmq_cmd_socket.close()
|
||||
self.zmq_context.term()
|
||||
|
||||
|
||||
def main():
|
||||
logging.info("Configuring LeKiwiRobot")
|
||||
|
|
@ -12,9 +39,51 @@ def main():
|
|||
logging.info("Connecting LeKiwiRobot")
|
||||
robot.connect()
|
||||
|
||||
# Remotely teleoperated
|
||||
logging.info("Starting LeKiwiRobot teleoperation")
|
||||
robot.run()
|
||||
logging.info("Starting RemoteAgent")
|
||||
remote_agent = RemoteAgent()
|
||||
|
||||
last_cmd_time = time.time()
|
||||
logging.info("Waiting for commands...")
|
||||
try:
|
||||
# Business logic
|
||||
start = time.perf_counter()
|
||||
duration = 0
|
||||
while duration < 100:
|
||||
loop_start_time = time.time()
|
||||
try:
|
||||
msg = remote_agent.zmq_cmd_socket.recv_string(zmq.NOBLOCK)
|
||||
data = np.array(json.loads(msg))
|
||||
_action_sent = robot.send_action(data)
|
||||
last_cmd_time = time.time()
|
||||
except zmq.Again:
|
||||
logging.warning("No command available")
|
||||
except Exception as e:
|
||||
logging.error("Message fetching failed: %s", e)
|
||||
|
||||
# TODO(Steven): Check this value
|
||||
# Watchdog: stop the robot if no command is received for over 0.5 seconds.
|
||||
now = time.time()
|
||||
if now - last_cmd_time > 0.5:
|
||||
robot.stop_base()
|
||||
|
||||
last_observation = robot.get_observation()
|
||||
last_observation[OBS_STATE] = last_observation[OBS_STATE].tolist()
|
||||
remote_agent.zmq_observation_socket.send_string(json.dumps(last_observation))
|
||||
|
||||
# Ensure a short sleep to avoid overloading the CPU.
|
||||
elapsed = time.time() - loop_start_time
|
||||
|
||||
# TODO(Steven): Check this value
|
||||
time.sleep(
|
||||
max(0.033 - elapsed, 0)
|
||||
) # If robot jitters increase the sleep and monitor cpu load with `top` in cmd
|
||||
duration = time.perf_counter() - start
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down LeKiwi server.")
|
||||
finally:
|
||||
robot.disconnect()
|
||||
remote_agent.disconnect()
|
||||
|
||||
logging.info("Finished LeKiwiRobot cleanly")
|
||||
|
||||
|
|
|
|||
|
|
@ -17,12 +17,10 @@
|
|||
import base64
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import zmq
|
||||
|
||||
from lerobot.common.cameras.utils import make_cameras_from_configs
|
||||
from lerobot.common.constants import OBS_IMAGES, OBS_STATE
|
||||
|
|
@ -55,9 +53,6 @@ class LeKiwiRobot(Robot):
|
|||
self.config = config
|
||||
self.id = config.id
|
||||
|
||||
self.port_zmq_cmd = config.port_zmq_cmd
|
||||
self.port_zmq_observations = config.port_zmq_observations
|
||||
|
||||
# TODO(Steven): Consider in the future using S100 robot class
|
||||
# TODO(Steven): Another option is to use the motorbus factory, but in this case we assume that
|
||||
# what we consider 'lekiwi robot' always uses the FeetechMotorsBus
|
||||
|
|
@ -85,13 +80,6 @@ class LeKiwiRobot(Robot):
|
|||
# TODO(Steven): Consider removing cameras from configs
|
||||
self.cameras = make_cameras_from_configs(config.cameras)
|
||||
|
||||
self.observation_lock = threading.Lock()
|
||||
self.last_observation = None
|
||||
|
||||
self.zmq_context = None
|
||||
self.zmq_cmd_socket = None
|
||||
self.zmq_observation_socket = None
|
||||
|
||||
self.is_connected = False
|
||||
self.logs = {}
|
||||
|
||||
|
|
@ -118,18 +106,6 @@ class LeKiwiRobot(Robot):
|
|||
}
|
||||
return cam_ft
|
||||
|
||||
def setup_zmq_sockets(self):
|
||||
context = zmq.Context()
|
||||
cmd_socket = context.socket(zmq.PULL)
|
||||
cmd_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
cmd_socket.bind(f"tcp://*:{self.port_zmq_cmd}")
|
||||
|
||||
observation_socket = context.socket(zmq.PUSH)
|
||||
observation_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
observation_socket.bind(f"tcp://*:{self.port_zmq_observations}")
|
||||
|
||||
return context, cmd_socket, observation_socket
|
||||
|
||||
def setup_actuators(self):
|
||||
# Set-up arm actuators (position mode)
|
||||
# We assume that at connection time, arm is in a rest position,
|
||||
|
|
@ -177,9 +153,6 @@ class LeKiwiRobot(Robot):
|
|||
for cam in self.cameras.values():
|
||||
cam.connect()
|
||||
|
||||
logging.info("Connecting ZMQ sockets.")
|
||||
self.zmq_context, self.zmq_cmd_socket, self.zmq_observation_socket = self.setup_zmq_sockets()
|
||||
|
||||
self.is_connected = True
|
||||
|
||||
def calibrate(self) -> None:
|
||||
|
|
@ -274,80 +247,13 @@ class LeKiwiRobot(Robot):
|
|||
|
||||
return goal_pos
|
||||
|
||||
def update_last_observation(self, stop_event):
|
||||
while not stop_event.is_set():
|
||||
obs = self.get_observation()
|
||||
obs[OBS_STATE] = obs[OBS_STATE].tolist() # Needed for np.array be serializable
|
||||
with self.observation_lock:
|
||||
self.last_observation = obs
|
||||
# TODO(Steven): Consider adding a delay to not starve the CPU
|
||||
# TODO(Steven): Check this value
|
||||
time.sleep(0.5)
|
||||
|
||||
def stop(self):
|
||||
def stop_base(self):
|
||||
# TODO(Steven): Assumes there's only 3 motors for base
|
||||
logging.info("Stopping base")
|
||||
# TODO(Steven): Check if these operations are thread safe!
|
||||
self.actuators_bus.write("Goal_Speed", [0, 0, 0], self.base_actuators)
|
||||
logging.info("Base motors stopped")
|
||||
|
||||
def run(self):
|
||||
# Copied logic from run_lekiwi in lekiwi_remote.py
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError("LeKiwiRobot is not connected. You need to run `robot.connect()`.")
|
||||
|
||||
stop_event = threading.Event()
|
||||
observation_thread = threading.Thread(
|
||||
target=self.update_last_observation, args=[stop_event], daemon=True
|
||||
)
|
||||
observation_thread.start()
|
||||
|
||||
last_cmd_time = time.time()
|
||||
logging.info("LeKiwi robot server started. Waiting for commands...")
|
||||
|
||||
try:
|
||||
start = time.perf_counter()
|
||||
duration = 0
|
||||
while duration < 100:
|
||||
loop_start_time = time.time()
|
||||
|
||||
try:
|
||||
msg = self.zmq_cmd_socket.recv_string(zmq.NOBLOCK)
|
||||
data = np.array(json.loads(msg))
|
||||
_action_sent = self.send_action(data)
|
||||
last_cmd_time = time.time()
|
||||
except zmq.Again:
|
||||
logging.warning("ZMQ again")
|
||||
except Exception as e:
|
||||
logging.error("Message fetching failed: %s", e)
|
||||
|
||||
# TODO(Steven): Check this value
|
||||
# Watchdog: stop the robot if no command is received for over 0.5 seconds.
|
||||
now = time.time()
|
||||
if now - last_cmd_time > 0.5:
|
||||
# TODO(Steven): This doesn't seem to be thread safe!
|
||||
# self.stop()
|
||||
pass
|
||||
|
||||
with self.observation_lock:
|
||||
# TODO(Steven): This operation is blocking if no listener is available
|
||||
self.zmq_observation_socket.send_string(json.dumps(self.last_observation))
|
||||
|
||||
# Ensure a short sleep to avoid overloading the CPU.
|
||||
elapsed = time.time() - loop_start_time
|
||||
|
||||
# TODO(Steven): Check this value
|
||||
time.sleep(
|
||||
max(0.033 - elapsed, 0)
|
||||
) # If robot jitters increase the sleep and monitor cpu load with `top` in cmd
|
||||
duration = time.perf_counter() - start
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down LeKiwi server.")
|
||||
finally:
|
||||
stop_event.set()
|
||||
observation_thread.join()
|
||||
self.disconnect()
|
||||
|
||||
def print_logs(self):
|
||||
# TODO(Steven): Refactor logger
|
||||
pass
|
||||
|
|
@ -358,13 +264,10 @@ class LeKiwiRobot(Robot):
|
|||
"LeKiwi is not connected. You need to run `robot.connect()` before disconnecting."
|
||||
)
|
||||
|
||||
self.stop()
|
||||
self.stop_base()
|
||||
self.actuators_bus.disconnect()
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
self.zmq_observation_socket.close()
|
||||
self.zmq_cmd_socket.close()
|
||||
self.zmq_context.term()
|
||||
self.is_connected = False
|
||||
|
||||
def __del__(self):
|
||||
|
|
|
|||
Loading…
Reference in New Issue