From bc55713e7cd267a8dce06e85b3eaa74fc1e44948 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 13 Jul 2026 18:50:00 +0200 Subject: [PATCH] fix relative imports --- .../controllers/gr00t_locomotion.py | 14 ++++++++-- .../controllers/holosoma_locomotion.py | 21 ++++++++++++-- .../unitree_g1/controllers/sonic_pipeline.py | 19 +++++++++---- .../controllers/sonic_whole_body.py | 28 +++++++++++++------ src/lerobot/utils/import_utils.py | 2 ++ 5 files changed, 65 insertions(+), 19 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index 31166e123..0a33e401b 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py @@ -14,20 +14,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +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.robots.unitree_g1.g1_utils import ( +from lerobot.utils.import_utils import _onnxruntime_available, require_package + +from ..g1_utils import ( REMOTE_AXES, REMOTE_BUTTONS, G1_29_JointIndex, get_gravity_orientation, ) +if TYPE_CHECKING or _onnxruntime_available: + import onnxruntime as ort +else: + ort = None + logger = logging.getLogger(__name__) @@ -83,6 +92,7 @@ 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 857bb97bc..41675970b 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -14,21 +14,34 @@ # See the License for the specific language governing permissions and # limitations under the License. +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.robots.unitree_g1.g1_utils import ( +from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package + +from ..g1_utils import ( REMOTE_AXES, G1_29_JointArmIndex, G1_29_JointIndex, 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) @@ -101,6 +114,8 @@ 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_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 72de873fd..05e38e881 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -16,6 +16,8 @@ """SONIC planner pipeline: ONNX enc/dec/planner, movement state, and input helpers.""" +from __future__ import annotations + import math import queue import select @@ -26,15 +28,22 @@ import time import tty from dataclasses import dataclass from enum import IntEnum +from typing import TYPE_CHECKING import numpy as np -import onnxruntime as ort -from lerobot.robots.unitree_g1.g1_utils import ( +from lerobot.utils.import_utils import _onnxruntime_available + +from ..g1_utils import ( G1_29_JointIndex, get_gravity_orientation, ) +if TYPE_CHECKING or _onnxruntime_available: + import onnxruntime as ort +else: + ort = None + # ── Constants ──────────────────────────────────────────────────────────────── DEFAULT_ANGLES = np.array( @@ -419,7 +428,7 @@ def replan_interval(mode): return REPLAN_INTERVAL["default"] -def _ort_providers(force_cpu: bool = False) -> list[str]: +def ort_providers(force_cpu: bool = False) -> list[str]: """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed).""" avail = ort.get_available_providers() if not force_cpu and "CUDAExecutionProvider" in avail: @@ -471,7 +480,7 @@ class MovementSnapshot: facing: tuple[float, float, float] = (1.0, 0.0, 0.0) -def _snapshot_ms(ms: MovementState) -> MovementSnapshot: +def snapshot_ms(ms: MovementState) -> MovementSnapshot: md, fd = ms.movement_direction, ms.facing_direction return MovementSnapshot(ms.mode, ms.speed, ms.height, (md[0], md[1], md[2]), (fd[0], fd[1], fd[2])) @@ -738,7 +747,7 @@ def _build_planner_inputs(ctx, ms_dict, version, seed): def _planner_worker(path, req_q, res_q, stop_evt, version, seed, use_gpu): so = ort.SessionOptions() so.log_severity_level = 3 - providers = _ort_providers(force_cpu=not use_gpu) + providers = ort_providers(force_cpu=not use_gpu) sess = ort.InferenceSession(path, sess_options=so, providers=providers) while not stop_evt.is_set(): try: 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 d0b91b006..b5a0fb19f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -16,14 +16,19 @@ """SONIC full-body controller for Unitree G1.""" +from __future__ import annotations + import logging import os +from typing import TYPE_CHECKING import numpy as np -import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.robots.unitree_g1.controllers.sonic_pipeline import ( +from lerobot.utils.import_utils import _onnxruntime_available, require_package + +from ..g1_utils import lowstate_to_obs +from .sonic_pipeline import ( CONTROL_DT, DEBUG_PRINT_EVERY, DEFAULT_ANGLES, @@ -33,14 +38,18 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import ( MovementState, PlannerController, SonicPlanner, - _ort_providers, - _snapshot_ms, clamp_mode_params, compute_kp_kd, + ort_providers, process_joystick, should_replan_request, + snapshot_ms, ) -from lerobot.robots.unitree_g1.g1_utils import lowstate_to_obs + +if TYPE_CHECKING or _onnxruntime_available: + import onnxruntime as ort +else: + ort = None logger = logging.getLogger(__name__) @@ -89,11 +98,12 @@ class SonicRuntime: """Shared SONIC control loop state (standalone demo + locomotion controller).""" def __init__(self, force_cpu: bool = False): + require_package("onnxruntime", extra="unitree_g1") planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx") encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx") decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") - providers = _ort_providers(force_cpu=force_cpu) + providers = ort_providers(force_cpu=force_cpu) self.use_gpu = providers[0] == "CUDAExecutionProvider" so = ort.SessionOptions() so.log_severity_level = 3 @@ -113,7 +123,7 @@ class SonicRuntime: self.step = 0 self.replan_timer = 0.0 - self.last_ms = _snapshot_ms(self.ms) + self.last_ms = snapshot_ms(self.ms) @property def pipeline(self): @@ -134,7 +144,7 @@ class SonicRuntime: self.planner.request_replan(self.controller.ref_cursor, self.ms) self.replan_timer = 0.0 self.ms.needs_replan = False - self.last_ms = _snapshot_ms(self.ms) + self.last_ms = snapshot_ms(self.ms) do_enc = self.step % ENCODER_UPDATE_EVERY == 0 if debug is None: @@ -155,7 +165,7 @@ class SonicRuntime: self.controller.playing = True self.step = 0 self.replan_timer = 0.0 - self.last_ms = _snapshot_ms(self.ms) + self.last_ms = snapshot_ms(self.ms) def shutdown(self): self.planner.stop_subprocess() diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 112f76a96..974dd5b51 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -123,6 +123,8 @@ _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")