From 503f3e57aef4bf1c0b1bb1b4d148d43dd17a70f2 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Wed, 29 Jul 2026 16:11:55 +0200 Subject: [PATCH] refactor(unitree_g1): minimize diff vs main (drop onnx guards, e-stop) Revert pyproject.toml and import_utils.py to main: the onnxruntime-gpu detection fallback and the _onnxruntime_available/_onnx_available flags are unnecessary. Controllers now import onnx/onnxruntime unconditionally (matching main) which also works with onnxruntime-gpu since the import name is stable. Drop the require_package() calls we had added. Remove the stdin e-stop listener from serve_onboard_controller and the now-unused os/sys imports; Ctrl-C still triggers graceful shutdown. --- pyproject.toml | 6 +----- src/lerobot/robots/unitree_g1/README.md | 2 +- .../unitree_g1/controllers/gr00t_locomotion.py | 10 +--------- .../controllers/holosoma_locomotion.py | 17 ++--------------- .../unitree_g1/controllers/sonic_whole_body.py | 10 +--------- src/lerobot/robots/unitree_g1/run_g1_server.py | 18 +----------------- src/lerobot/utils/import_utils.py | 16 ++-------------- 7 files changed, 9 insertions(+), 70 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 90c1653f7..d07080232 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -374,11 +374,7 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }] torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }] [tool.setuptools.package-data] -lerobot = [ - "envs/*.json", - "annotations/steerable_pipeline/prompts/*.txt", - "teleoperators/pico_headset/assets/*.npz", -] +lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"] [tool.setuptools.packages.find] where = ["src"] diff --git a/src/lerobot/robots/unitree_g1/README.md b/src/lerobot/robots/unitree_g1/README.md index 971684b0c..16a8f5381 100644 --- a/src/lerobot/robots/unitree_g1/README.md +++ b/src/lerobot/robots/unitree_g1/README.md @@ -60,7 +60,7 @@ lerobot-rollout \ --robot.is_simulation=false --robot.onboard=false \ --robot.robot_ip= \ --robot.controller=SonicWholeBodyController --robot.sonic_token_action=true \ - --robot.cameras='{ego_view: {type: zmq, server_address: , port: 5555, camera_name: ego_view}}' \ + --robot.cameras='{ego_view: {type: zmq, server_address: , port: 5555, camera_name: ego_view, width: 640, height: 480, fps: 30}}' \ --task="walk back and forth" --device=cuda ``` diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index 1313ad325..2cc39051a 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py @@ -18,13 +18,11 @@ from __future__ import annotations import logging from collections import deque -from typing import TYPE_CHECKING import numpy as np +import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.utils.import_utils import _onnxruntime_available, require_package - from ..g1_utils import ( REMOTE_AXES, REMOTE_BUTTONS, @@ -32,11 +30,6 @@ from ..g1_utils import ( get_gravity_orientation, ) -if TYPE_CHECKING or _onnxruntime_available: - import onnxruntime as ort -else: - ort = None - logger = logging.getLogger(__name__) @@ -98,7 +91,6 @@ class GrootLocomotionController: control_dt = CONTROL_DT # Expose for unitree_g1.py def __init__(self): - require_package("onnxruntime", extra="unitree_g1") # Load policies self.policy_balance, self.policy_walk = load_groot_policies() diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index 41675970b..c171c9179 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -18,13 +18,12 @@ from __future__ import annotations import json import logging -from typing import TYPE_CHECKING import numpy as np +import onnx +import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package - from ..g1_utils import ( REMOTE_AXES, G1_29_JointArmIndex, @@ -32,16 +31,6 @@ from ..g1_utils import ( get_gravity_orientation, ) -if TYPE_CHECKING or _onnxruntime_available: - import onnxruntime as ort -else: - ort = None - -if TYPE_CHECKING or _onnx_available: - import onnx -else: - onnx = None - logger = logging.getLogger(__name__) DEFAULT_ANGLES = np.zeros(29, dtype=np.float32) @@ -114,8 +103,6 @@ class HolosomaLocomotionController: control_dt = CONTROL_DT # Expose for unitree_g1.py def __init__(self): - require_package("onnxruntime", extra="unitree_g1") - require_package("onnx", extra="unitree_g1") # Load policy and gains self.policy, self.kp, self.kd = load_policy() diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index e1a187c5c..ccc4bda84 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -31,13 +31,11 @@ convert between them. Quaternions are scalar-first ``(w, x, y, z)``. from __future__ import annotations import logging -from typing import TYPE_CHECKING import numpy as np +import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.utils.import_utils import _onnxruntime_available, require_package - from ..g1_utils import ( ISAACLAB_TO_MUJOCO, MUJOCO_TO_ISAACLAB, @@ -48,11 +46,6 @@ from ..g1_utils import ( ort_providers, ) -if TYPE_CHECKING or _onnxruntime_available: - import onnxruntime as ort -else: - ort = None - logger = logging.getLogger(__name__) # ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ────── @@ -311,7 +304,6 @@ class SonicRuntime: """ def __init__(self, force_cpu: bool = False): - require_package("onnxruntime", extra="unitree_g1") decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") providers = ort_providers(force_cpu=force_cpu) diff --git a/src/lerobot/robots/unitree_g1/run_g1_server.py b/src/lerobot/robots/unitree_g1/run_g1_server.py index aadf81e11..95afeb172 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_server.py +++ b/src/lerobot/robots/unitree_g1/run_g1_server.py @@ -43,10 +43,8 @@ import argparse import base64 import contextlib import json -import os import re import signal -import sys import threading import time from typing import Any @@ -257,21 +255,7 @@ def serve_onboard_controller( sock.setsockopt(zmq.RCVTIMEO, 200) # keeps the loop responsive to the stop event sock.bind(f"tcp://0.0.0.0:{action_port}") print(f"Onboard controller live. Waiting for laptop actions on :{action_port} ...") - print("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).") - - def estop_listener() -> None: - for line in sys.stdin: - if line.strip().lower() == "e": - print("E-STOP ('e'): going passive NOW.") - try: - robot._shutdown_event.set() # stop the controller loop publishing - time.sleep(0.05) - robot._send_zero_torque() # motors limp; nothing overwrites it now - except Exception as e: # noqa: BLE001 - print(f"E-stop zero-torque failed: {e}") - os._exit(0) # immediate hard exit, no slow cleanup - - threading.Thread(target=estop_listener, daemon=True).start() + print("Ctrl-C for graceful shutdown.") state_sock = None if state_fps > 0: diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index be808787e..112f76a96 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -60,18 +60,8 @@ def is_package_available( # If the package can't be imported, it's not available package_exists = False else: - # The distribution may be published under a name that differs from the - # import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` / - # ``onnxruntime-silicon``). Resolve the import name to its actual - # distribution(s) and read the version from there before giving up. - try: - dists = importlib.metadata.packages_distributions().get(import_name, []) - if dists: - package_version = importlib.metadata.version(dists[0]) - else: - package_exists = False - except importlib.metadata.PackageNotFoundError: - package_exists = False + # For packages other than "torch", don't attempt the fallback and set as not available + package_exists = False logging.debug(f"Detected {pkg_name} version: {package_version}") if return_version: return package_exists, package_version @@ -133,8 +123,6 @@ _pyrealsense2_available = is_package_available("pyrealsense2") or is_package_ava "pyrealsense2-macosx", import_name="pyrealsense2" ) _zmq_available = is_package_available("pyzmq", import_name="zmq") -_onnxruntime_available = is_package_available("onnxruntime") -_onnx_available = is_package_available("onnx") _hebi_available = is_package_available("hebi-py", import_name="hebi") _teleop_available = is_package_available("teleop") _placo_available = is_package_available("placo")