admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 18:26:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(unitree_g1): load GR00T policies from LEROBOT_GROOT_POLICY_DIR

Browse Source 
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 <cursoragent@cursor.com>
This commit is contained in:
Martino Russi
2026-07-24 12:02:51 +02:00
parent d5cf538ffb
commit 39fe436248
1 changed files with 21 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/robots/unitree_g1/gr00t_locomotion.py
+21 -12
View File
@@ -15,6 +15,7 @@
# limitations under the License. # limitations under the License.
import logging import logging
import os
from collections import deque from collections import deque
import numpy as np import numpy as np
@@ -63,22 +64,30 @@ DEFAULT_GROOT_REPO_ID = "nepyope/GR00T-WholeBodyControl_g1"
def load_groot_policies( def load_groot_policies(
repo_id: str = DEFAULT_GROOT_REPO_ID, repo_id: str = DEFAULT_GROOT_REPO_ID,
) -> tuple[ort.InferenceSession, ort.InferenceSession]: ) -> 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: Args:
repo_id: Hugging Face Hub repository ID containing the ONNX policies. repo_id: Hugging Face Hub repository ID containing the ONNX policies.
""" """
logger.info(f"Loading GR00T dual-policy system from the hub ({repo_id})...") local_dir = os.environ.get("LEROBOT_GROOT_POLICY_DIR")
if local_dir:
# Download ONNX policies from Hugging Face Hub balance_path = os.path.join(local_dir, "GR00T-WholeBodyControl-Balance.onnx")
balance_path = hf_hub_download( walk_path = os.path.join(local_dir, "GR00T-WholeBodyControl-Walk.onnx")
repo_id=repo_id, logger.info(f"Loading GR00T dual-policy system from local dir ({local_dir})...")
filename="GR00T-WholeBodyControl-Balance.onnx", else:
) logger.info(f"Loading GR00T dual-policy system from the hub ({repo_id})...")
walk_path = hf_hub_download( balance_path = hf_hub_download(
repo_id=repo_id, repo_id=repo_id,
filename="GR00T-WholeBodyControl-Walk.onnx", 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 # 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 # policies run in the real-time control loop, and letting ORT grab one intra-op