From 0b866eadff8bf4f8349140ee61d9777e05f6ae02 Mon Sep 17 00:00:00 2001 From: cherylwang20 Date: Mon, 29 Jun 2026 14:36:00 +0800 Subject: [PATCH 1/2] Refactor robot model handling and update dependencies - Replace musclemimic_models with myo-sim for robot model integration. - Introduce ROBOT_MODEL_DICT for managing robot models and their paths. - Update get_robot_model function to handle both XML paths and model objects. - Modify relevant scripts and modules to utilize the new model handling. - Add unit tests for myo_sim model integration and XML loading functionality. --- general_motion_retargeting/__init__.py | 11 ++- general_motion_retargeting/motion_retarget.py | 10 ++- general_motion_retargeting/params.py | 28 ++++++- .../robot_motion_viewer.py | 8 +- .../utils/shape_fitting.py | 9 ++- pyproject.toml | 3 +- scripts/bvh_to_robot_dataset.py | 3 - scripts/fit_smpl_shape.py | 13 ++-- tests/test_robot_model_factory.py | 75 +++++++++++++++++++ 9 files changed, 136 insertions(+), 24 deletions(-) create mode 100644 tests/test_robot_model_factory.py diff --git a/general_motion_retargeting/__init__.py b/general_motion_retargeting/__init__.py index 2c7c312..67b6f7a 100644 --- a/general_motion_retargeting/__init__.py +++ b/general_motion_retargeting/__init__.py @@ -1,5 +1,14 @@ from rich import print -from .params import IK_CONFIG_ROOT, ASSET_ROOT, ROBOT_XML_DICT, IK_CONFIG_DICT, ROBOT_BASE_DICT, VIEWER_CAM_DISTANCE_DICT +from .params import ( + IK_CONFIG_ROOT, + ASSET_ROOT, + ROBOT_XML_DICT, + ROBOT_MODEL_DICT, + IK_CONFIG_DICT, + ROBOT_BASE_DICT, + VIEWER_CAM_DISTANCE_DICT, + get_robot_model, +) from .motion_retarget import GeneralMotionRetargeting from .robot_motion_viewer import RobotMotionViewer from .data_loader import load_robot_motion diff --git a/general_motion_retargeting/motion_retarget.py b/general_motion_retargeting/motion_retarget.py index 3f868a4..8ec161f 100644 --- a/general_motion_retargeting/motion_retarget.py +++ b/general_motion_retargeting/motion_retarget.py @@ -4,7 +4,7 @@ import numpy as np import json from scipy.spatial.transform import Rotation as R -from .params import ROBOT_XML_DICT, IK_CONFIG_DICT +from .params import ROBOT_XML_DICT, ROBOT_MODEL_DICT, IK_CONFIG_DICT, get_robot_model from .utils.shape_fitting import load_fitted_shape from rich import print from mink.tasks.equality_constraint_task import EqualityConstraintTask @@ -27,10 +27,12 @@ def __init__( ) -> None: # load the robot model - self.xml_file = str(ROBOT_XML_DICT[tgt_robot]) + robot_model = ROBOT_MODEL_DICT[tgt_robot] + self.model_source = robot_model if tgt_robot in ROBOT_XML_DICT else tgt_robot + self.xml_file = str(self.model_source) if verbose: - print("Use robot model: ", self.xml_file) - self.model = mj.MjModel.from_xml_path(self.xml_file) + print("Use robot model: ", self.model_source) + self.model = get_robot_model(tgt_robot) # Print DoF names in order print("[GMR] Robot Degrees of Freedom (DoF) names and their order:") diff --git a/general_motion_retargeting/params.py b/general_motion_retargeting/params.py index 56be6a4..59d9ff3 100644 --- a/general_motion_retargeting/params.py +++ b/general_motion_retargeting/params.py @@ -1,11 +1,17 @@ import pathlib -from musclemimic_models import get_xml_path HERE = pathlib.Path(__file__).parent IK_CONFIG_ROOT = HERE / "ik_configs" ASSET_ROOT = HERE / ".." / "assets" -ROBOT_XML_DICT = { + +def get_myo_sim_model(name: str): + from myo_sim.build.compose import build_model + + return build_model(name) + + +ROBOT_MODEL_DICT = { "unitree_g1": ASSET_ROOT / "unitree_g1" / "g1_mocap_29dof.xml", "unitree_g1_with_hands": ASSET_ROOT / "unitree_g1" / "g1_mocap_29dof_with_hands.xml", "unitree_h1": ASSET_ROOT / "unitree_h1" / "h1.xml", @@ -23,10 +29,26 @@ "pnd_adam_lite": ASSET_ROOT / "pnd_adam_lite" / "scene.xml", "tienkung": ASSET_ROOT / "tienkung" / "mjcf" / "tienkung.xml", "pal_talos": ASSET_ROOT / "pal_talos" / "talos.xml", - "myofullbody": get_xml_path("myofullbody"), + "myofullbody": get_myo_sim_model("myofullbody"), +} +ROBOT_XML_DICT = { + robot_name: robot_model + for robot_name, robot_model in ROBOT_MODEL_DICT.items() + if isinstance(robot_model, pathlib.Path) } + +def get_robot_model(tgt_robot: str): + robot_model = ROBOT_MODEL_DICT[tgt_robot] + if not isinstance(robot_model, pathlib.Path): + return robot_model + + import mujoco as mj + + xml_path = robot_model + return mj.MjModel.from_xml_path(str(xml_path)) + IK_CONFIG_DICT = { # offline data "smplx":{ diff --git a/general_motion_retargeting/robot_motion_viewer.py b/general_motion_retargeting/robot_motion_viewer.py index b91af80..c6d8fe3 100644 --- a/general_motion_retargeting/robot_motion_viewer.py +++ b/general_motion_retargeting/robot_motion_viewer.py @@ -9,7 +9,7 @@ import mujoco.viewer as mjv import imageio from scipy.spatial.transform import Rotation as R -from general_motion_retargeting import ROBOT_XML_DICT, ROBOT_BASE_DICT, VIEWER_CAM_DISTANCE_DICT +from general_motion_retargeting import ROBOT_XML_DICT, ROBOT_MODEL_DICT, ROBOT_BASE_DICT, VIEWER_CAM_DISTANCE_DICT, get_robot_model from loop_rate_limiters import RateLimiter import numpy as np from rich import print @@ -62,8 +62,10 @@ def __init__(self, ): self.robot_type = robot_type - self.xml_path = ROBOT_XML_DICT[robot_type] - self.model = mj.MjModel.from_xml_path(str(self.xml_path)) + robot_model = ROBOT_MODEL_DICT[robot_type] + self.model_source = robot_model if robot_type in ROBOT_XML_DICT else robot_type + self.xml_path = self.model_source + self.model = get_robot_model(robot_type) self.data = mj.MjData(self.model) self.robot_base = ROBOT_BASE_DICT[robot_type] self.viewer_cam_distance = VIEWER_CAM_DISTANCE_DICT[robot_type] diff --git a/general_motion_retargeting/utils/shape_fitting.py b/general_motion_retargeting/utils/shape_fitting.py index 1931910..c44b630 100644 --- a/general_motion_retargeting/utils/shape_fitting.py +++ b/general_motion_retargeting/utils/shape_fitting.py @@ -131,7 +131,7 @@ def set_robot_to_tpose( def get_robot_tpose_targets( - robot_xml_path: str, + robot_model_or_path, robot_type: str, ik_config: Dict, include_rotations: bool = False, @@ -140,7 +140,7 @@ def get_robot_tpose_targets( Extract target joint positions from robot in T-pose. Args: - robot_xml_path: Path to robot MuJoCo XML + robot_model_or_path: Robot MuJoCo model or path to robot MuJoCo XML robot_type: Robot type (e.g., 'unitree_g1') ik_config: IK configuration with site-to-joint mappings @@ -149,7 +149,10 @@ def get_robot_tpose_targets( smpl_joint_names: List of corresponding SMPL-H joint names """ # Load robot model - model = mj.MjModel.from_xml_path(robot_xml_path) + if isinstance(robot_model_or_path, mj.MjModel): + model = robot_model_or_path + else: + model = mj.MjModel.from_xml_path(str(robot_model_or_path)) data = mj.MjData(model) # Set robot to T-pose diff --git a/pyproject.toml b/pyproject.toml index 781793c..690d6be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [ "torch>=2.0.0", "joblib>=1.3.0", "pytest>=9.0.1", - "musclemimic_models>=1.0.2" + "myo-sim" ] [project.urls] @@ -49,3 +49,4 @@ packages = {find = {}} [tool.uv.sources] smplx = { git = "https://github.com/amathislab/smplx" } +myo-sim = { git = "https://github.com/amathislab/gmr_plus.git", branch = "dev", subdirectory = "external/myo_sim" } diff --git a/scripts/bvh_to_robot_dataset.py b/scripts/bvh_to_robot_dataset.py index d87b598..23caa8f 100644 --- a/scripts/bvh_to_robot_dataset.py +++ b/scripts/bvh_to_robot_dataset.py @@ -1,7 +1,6 @@ import argparse import pathlib import os -import mujoco as mj import numpy as np from tqdm import tqdm import torch @@ -86,8 +85,6 @@ tgt_robot=args.robot, actual_human_height=actual_human_height, ) - model = mj.MjModel.from_xml_path(retarget.xml_file) - data = mj.MjData(model) diff --git a/scripts/fit_smpl_shape.py b/scripts/fit_smpl_shape.py index 616b781..e369577 100644 --- a/scripts/fit_smpl_shape.py +++ b/scripts/fit_smpl_shape.py @@ -20,7 +20,7 @@ import json import numpy as np -from general_motion_retargeting import ROBOT_XML_DICT, IK_CONFIG_DICT +from general_motion_retargeting import ROBOT_XML_DICT, ROBOT_MODEL_DICT, IK_CONFIG_DICT, get_robot_model from general_motion_retargeting.utils.smpl import SMPLH_Parser from general_motion_retargeting.utils.shape_fitting import ( get_robot_tpose_targets, @@ -47,7 +47,7 @@ def main(): "--robot", type=str, required=True, - choices=list(ROBOT_XML_DICT.keys()), + choices=list(ROBOT_MODEL_DICT.keys()), help="Robot type to fit", ) parser.add_argument( @@ -104,19 +104,20 @@ def main(): console.print(f"\n[bold cyan]Step 1: Loading robot and IK configuration[/bold cyan]") console.print(f" Robot: [green]{args.robot}[/green]") - # Load robot XML and IK config (use SMPL-H configs for shape fitting) - robot_xml_path = str(ROBOT_XML_DICT[args.robot]) + # Load robot model and IK config (use SMPL-H configs for shape fitting) + robot_model = get_robot_model(args.robot) + robot_model_source = ROBOT_MODEL_DICT[args.robot] if args.robot in ROBOT_XML_DICT else args.robot ik_config_path = IK_CONFIG_DICT["smplh"][args.robot] with open(ik_config_path, 'r') as f: ik_config = json.load(f) - console.print(f" Robot XML: {robot_xml_path}") + console.print(f" Robot model: {robot_model_source}") console.print(f" IK Config: {ik_config_path}") console.print(f"\n[bold cyan]Step 2: Extracting robot T-pose targets[/bold cyan]") target_positions, smpl_joint_names, target_rotations, human_joint_names = get_robot_tpose_targets( - robot_xml_path, + robot_model, args.robot, ik_config, include_rotations=True, diff --git a/tests/test_robot_model_factory.py b/tests/test_robot_model_factory.py new file mode 100644 index 0000000..fccf80d --- /dev/null +++ b/tests/test_robot_model_factory.py @@ -0,0 +1,75 @@ +import importlib.util +import sys +import types +from pathlib import Path + + +def load_params_module(monkeypatch): + if "myo_sim" not in sys.modules: + myo_sim = types.ModuleType("myo_sim") + compose = types.ModuleType("myo_sim.build.compose") + build_package = types.ModuleType("myo_sim.build") + + myo_sim.get_xml_path = lambda name: Path(f"/fake/{name}.xml") + compose.GENERATE_XML_TARGETS = {} + compose.build_model = lambda name: {"model": name} + compose.build_generated_model_spec = lambda name: None + compose.write_spec_xml = lambda spec, output_path: None + monkeypatch.setitem(sys.modules, "myo_sim", myo_sim) + monkeypatch.setitem(sys.modules, "myo_sim.build", build_package) + monkeypatch.setitem(sys.modules, "myo_sim.build.compose", compose) + + module_path = Path(__file__).resolve().parents[1] / "general_motion_retargeting" / "params.py" + module_name = "gmr_params_under_test" + sys.modules.pop(module_name, None) + + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_myofullbody_model_builds_from_myo_sim_mjspec(monkeypatch): + built_models = [] + + myo_sim = types.ModuleType("myo_sim") + compose = types.ModuleType("myo_sim.build.compose") + build_package = types.ModuleType("myo_sim.build") + + def build_model(name): + built_models.append(name) + return {"model": name} + + myo_sim.get_xml_path = lambda name: (_ for _ in ()).throw(ValueError(name)) + compose.build_model = build_model + monkeypatch.setitem(sys.modules, "myo_sim", myo_sim) + monkeypatch.setitem(sys.modules, "myo_sim.build", build_package) + monkeypatch.setitem(sys.modules, "myo_sim.build.compose", compose) + + params = load_params_module(monkeypatch) + + assert params.ROBOT_MODEL_DICT["myofullbody"] == {"model": "myofullbody"} + assert params.get_robot_model("myofullbody") == {"model": "myofullbody"} + assert built_models == ["myofullbody"] + + +def test_xml_robot_model_loads_from_xml_path(monkeypatch): + loaded_paths = [] + + mujoco = types.ModuleType("mujoco") + + class MjModel: + @staticmethod + def from_xml_path(path): + loaded_paths.append(path) + return {"xml_path": path} + + mujoco.MjModel = MjModel + monkeypatch.setitem(sys.modules, "mujoco", mujoco) + + params = load_params_module(monkeypatch) + + model = params.get_robot_model("unitree_g1") + + assert model == {"xml_path": str(params.ROBOT_XML_DICT["unitree_g1"])} + assert loaded_paths == [str(params.ROBOT_XML_DICT["unitree_g1"])] From bbcbafeb2365faebcef98b20e0b78ff9f535768a Mon Sep 17 00:00:00 2001 From: cherylwang20 Date: Wed, 15 Jul 2026 20:04:13 +0800 Subject: [PATCH 2/2] update myo_sim pointing --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 690d6be..cc3db3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,4 +49,4 @@ packages = {find = {}} [tool.uv.sources] smplx = { git = "https://github.com/amathislab/smplx" } -myo-sim = { git = "https://github.com/amathislab/gmr_plus.git", branch = "dev", subdirectory = "external/myo_sim" } +myo-sim = { git = "https://github.com/MyoHub/myo_sim.git", branch = "dev" }