diff --git a/screenshot.png b/screenshot.png new file mode 100644 index 000000000..7a7b8016b Binary files /dev/null and b/screenshot.png differ diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py index 154e5c235..9ef6f95b4 100644 --- a/src/lerobot/cameras/zmq/camera_zmq.py +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -23,12 +23,14 @@ import time from pathlib import Path from threading import Event, Lock, Thread from typing import Any - +import base64 import cv2 import numpy as np import zmq from numpy.typing import NDArray - +import base64 +import msgpack +import msgpack_numpy as m from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..camera import Camera @@ -303,30 +305,19 @@ class ZMQCamera(Camera): """ Reads a single frame synchronously from the ZMQ camera. - This is a blocking call. It waits for the next available frame from the - ZMQ socket. - - Args: - color_mode: If specified, overrides the default color mode for this read. - - Returns: - np.ndarray: The captured frame as a NumPy array in the format - (height, width, channels), using the specified or default - color mode. - - Raises: - DeviceNotConnectedError: If the camera is not connected. - TimeoutError: If no frame is received within the timeout period. - RuntimeError: If reading the frame fails. + For the mujoco sim (127.0.0.1:5554), the server sends a msgpack-encoded + dict with base64-encoded JPEG images: + {"timestamps": {...}, "images": {camera_name: ""}} + For other ZMQ cameras, we assume a raw JPEG buffer. """ if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") - start_time = time.perf_counter() - if self.socket is None: raise DeviceNotConnectedError(f"{self} socket is not initialized") + start_time = time.perf_counter() + try: message = self.socket.recv() except zmq.Again: @@ -334,11 +325,49 @@ class ZMQCamera(Camera): except Exception as e: raise RuntimeError(f"{self} read failed: {e}") - # Decode JPEG - np_img = np.frombuffer(message, dtype=np.uint8) - frame = cv2.imdecode(np_img, cv2.IMREAD_COLOR) + frame = None - if frame is None: + # special-case: mujoco sim publisher (msgpack + base64) + if self.server_address in ("127.0.0.1", "localhost") and self.port == 5554: + try: + data = msgpack.unpackb(message, object_hook=m.decode) + + if not isinstance(data, dict) or "images" not in data: + raise RuntimeError(f"{self} received invalid message format from sim") + + images_dict = data["images"] + img_data = None + + # prefer named camera if present + if self.camera_name in images_dict: + img_data = images_dict[self.camera_name] + elif len(images_dict) > 0: + # fallback: first available camera + img_data = next(iter(images_dict.values())) + else: + raise RuntimeError(f"{self} no images found in sim message") + + # same logic as ImageUtils.decode_image + if isinstance(img_data, str): + color_bytes = base64.b64decode(img_data) + np_img = np.frombuffer(color_bytes, dtype=np.uint8) + frame = cv2.imdecode(np_img, cv2.IMREAD_COLOR) + elif isinstance(img_data, np.ndarray): + frame = img_data + else: + raise RuntimeError( + f"{self} unknown image payload type from sim: {type(img_data)}" + ) + + except Exception as e: + raise RuntimeError(f"{self} failed to decode sim image: {e}") + + # default path: raw jpeg over ZMQ (robot-side cameras etc.) + else: + np_img = np.frombuffer(message, dtype=np.uint8) + frame = cv2.imdecode(np_img, cv2.IMREAD_COLOR) + + if frame is None or not isinstance(frame, np.ndarray): raise RuntimeError(f"{self} failed to decode image") processed_frame = self._postprocess_image(frame, color_mode) diff --git a/src/lerobot/envs/factory.py b/src/lerobot/envs/factory.py index bda84fd78..110208c6a 100644 --- a/src/lerobot/envs/factory.py +++ b/src/lerobot/envs/factory.py @@ -111,7 +111,6 @@ def make_env( # import and surface clear import errors module = _import_hub_module(local_file, repo_id) - # call the hub-provided make_env raw_result = _call_make_env(module, n_envs=n_envs, use_async_envs=use_async_envs) diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 47e467f76..59e1ecb8b 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -26,7 +26,7 @@ from ..config import RobotConfig class UnitreeG1Config(RobotConfig): # id: str = "unitree_g1" motion_mode: bool = False - simulation_mode: bool = False + simulation_mode: bool = True kp_high = 40.0 kd_high = 3.0 kp_low = 80.0 diff --git a/src/lerobot/robots/unitree_g1/robot_kinematic_processor.py b/src/lerobot/robots/unitree_g1/robot_kinematic_processor.py index bb093bed9..7a2ff282e 100644 --- a/src/lerobot/robots/unitree_g1/robot_kinematic_processor.py +++ b/src/lerobot/robots/unitree_g1/robot_kinematic_processor.py @@ -91,7 +91,7 @@ def visualize_filter_comparison(filter_params, steps): plt.xlabel("Step") plt.ylabel("Value") plt.legend() - + # col13 should 2b filtered plt.subplot(len(filter_params), 2, idx * 2 + 2) plt.plot(filtered_data[:, 13], label=f"Filtered (Window {filter._window_size})") diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index d82969970..c8a181747 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -56,6 +56,7 @@ from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import ( MotionSwitcherClient, ) +from lerobot.envs.factory import make_env from scipy.spatial.transform import Rotation as R import struct @@ -169,7 +170,6 @@ class UnitreeG1(Robot): # Launch MuJoCo simulation environment logger_mp.info("Launching MuJoCo simulation environment...") - from lerobot.envs.factory import make_env self.mujoco_env = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) logger_mp.info("MuJoCo environment launched successfully!") else: @@ -573,7 +573,8 @@ class UnitreeG1(Robot): # Close MuJoCo environment if in simulation mode if self.simulation_mode and hasattr(self, 'mujoco_env'): logger_mp.info("Closing MuJoCo environment...") - self.mujoco_env.close() + print(self.mujoco_env) + self.mujoco_env["hub_env"][0].envs[0].kill_sim() logger_mp.info(f"{self} disconnected.")