From 39fe4362480f2a903f994dbee483c1c119600036 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 24 Jul 2026 12:02:51 +0200 Subject: [PATCH] feat(unitree_g1): load GR00T policies from LEROBOT_GROOT_POLICY_DIR Allow load_groot_policies() to read the Balance/Walk ONNX from a local directory (e.g. finetuned checkpoints) via the LEROBOT_GROOT_POLICY_DIR env var, falling back to the Hub download when unset. Co-authored-by: Cursor --- .../robots/unitree_g1/gr00t_locomotion.py | 33 ++++++++++++------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/gr00t_locomotion.py index 51a873e71..29f53f802 100644 --- a/src/lerobot/robots/unitree_g1/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/gr00t_locomotion.py @@ -15,6 +15,7 @@ # limitations under the License. import logging +import os from collections import deque import numpy as np @@ -63,22 +64,30 @@ DEFAULT_GROOT_REPO_ID = "nepyope/GR00T-WholeBodyControl_g1" def load_groot_policies( repo_id: str = DEFAULT_GROOT_REPO_ID, ) -> tuple[ort.InferenceSession, ort.InferenceSession]: - """Load GR00T dual-policy system (Balance + Walk) from the hub. + """Load GR00T dual-policy system (Balance + Walk). + + If the env var ``LEROBOT_GROOT_POLICY_DIR`` is set, the two ONNX files are + loaded from that local directory (e.g. finetuned checkpoints) instead of the + Hub. Otherwise they are downloaded from ``repo_id``. Args: repo_id: Hugging Face Hub repository ID containing the ONNX policies. """ - logger.info(f"Loading GR00T dual-policy system from the hub ({repo_id})...") - - # Download ONNX policies from Hugging Face Hub - balance_path = hf_hub_download( - repo_id=repo_id, - filename="GR00T-WholeBodyControl-Balance.onnx", - ) - walk_path = hf_hub_download( - repo_id=repo_id, - filename="GR00T-WholeBodyControl-Walk.onnx", - ) + local_dir = os.environ.get("LEROBOT_GROOT_POLICY_DIR") + if local_dir: + balance_path = os.path.join(local_dir, "GR00T-WholeBodyControl-Balance.onnx") + walk_path = os.path.join(local_dir, "GR00T-WholeBodyControl-Walk.onnx") + logger.info(f"Loading GR00T dual-policy system from local dir ({local_dir})...") + else: + logger.info(f"Loading GR00T dual-policy system from the hub ({repo_id})...") + balance_path = hf_hub_download( + repo_id=repo_id, + filename="GR00T-WholeBodyControl-Balance.onnx", + ) + walk_path = hf_hub_download( + repo_id=repo_id, + filename="GR00T-WholeBodyControl-Walk.onnx", + ) # Load ONNX policies. Cap onnxruntime to a single thread: these are tiny MLP # policies run in the real-time control loop, and letting ORT grab one intra-op