Refactor robot model handling and update dependencies

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

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:")

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":{

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]

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

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/MyoHub/myo_sim.git", branch = "dev" }

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)
 
 
 

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,

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"])]