Add H1 mimic env for mocap reference tracking

.gitignore

@@ -4,3 +4,6 @@ __pycache__*
 *.egg-info
 *.#~
 logs*
+
+# Mocap reference clips
+data/mocap_data/

envs/__init__.py

@@ -5,7 +5,7 @@
 """
 
 from envs.cartpole import CartpoleEnv
-from envs.h1 import H1Env, H1WalkEnv
+from envs.h1 import H1Env, H1MimicEnv, H1WalkEnv
 from envs.jvrc import JvrcStepEnv, JvrcWalkEnv
 
 # Registry of all available environments
@@ -15,6 +15,7 @@
     "jvrc_step": (JvrcStepEnv, "jvrc"),
     "h1": (H1Env, "h1"),
     "h1_walk": (H1WalkEnv, "h1"),
+    "h1_mimic": (H1MimicEnv, "h1"),
     "cartpole": (CartpoleEnv, "cartpole"),
 }
 
@@ -23,6 +24,7 @@
     "JvrcStepEnv",
     "H1Env",
     "H1WalkEnv",
+    "H1MimicEnv",
     "CartpoleEnv",
     "ENVIRONMENTS",
 ]

envs/common/utils.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+
+import pandas as pd
+
+
+def get_project_root() -> Path:
+    return Path(__file__).parent.parent.parent
+
+
+def parse_mocap_csv(filepath, usecols):
+    """Process a mocap CSV file exported from Motive/NatNet format.
+
+    Returns a DataFrame with columns:
+        Frame, Time, RotX, RotY, RotZ, RotW, PosX, PosY, PosZ
+    """
+    with open(filepath) as file:
+        lines = file.readlines()
+
+    for i, line in enumerate(lines):
+        if line.strip().startswith("Frame"):
+            header_index = i
+            break
+    else:
+        raise ValueError("Could not find header row starting with 'Frame'.")
+
+    df = pd.read_csv(filepath, skiprows=header_index, usecols=usecols)
+    df.columns = ["Frame", "Time", "RotX", "RotY", "RotZ", "RotW", "PosX", "PosY", "PosZ"]
+    return df

envs/h1/__init__.py

@@ -1,4 +1,5 @@
 from envs.h1.h1_env import H1Env
+from envs.h1.h1_mimic import H1MimicEnv
 from envs.h1.h1_walk import H1WalkEnv
 
-__all__ = ["H1Env", "H1WalkEnv"]
+__all__ = ["H1Env", "H1WalkEnv", "H1MimicEnv"]

envs/h1/configs/mimic.yaml

@@ -0,0 +1,72 @@
+sim_dt: 0.001
+control_dt: 0.025
+obs_history_len: 1
+action_smoothing: 0.5
+
+xml_export_path: /tmp/mjcf-export
+
+ctrllimited: false
+jointlimited: false
+reduced_xml: true
+
+init_noise: 3
+
+# Half-sitting pose for all 19 actuated joints (legs, waist, arms)
+half_sitting_pose: [0, 0, -0.2, 0.6, -0.4,
+                    0, 0, -0.2, 0.6, -0.4,
+                    0,
+                    0, 0, 0, 0,
+                    0, 0, 0, 0]
+
+pdgains:
+  "left_hip_yaw": [100, 10]
+  "left_hip_roll": [100, 10]
+  "left_hip_pitch": [100, 10]
+  "left_knee": [100, 10]
+  "left_ankle": [20, 4]
+  "right_hip_yaw": [100, 10]
+  "right_hip_roll": [100, 10]
+  "right_hip_pitch": [100, 10]
+  "right_knee": [100, 10]
+  "right_ankle": [20, 4]
+  "torso": [40, 4]
+  "left_shoulder_pitch": [20, 2]
+  "left_shoulder_roll": [20, 2]
+  "left_shoulder_yaw": [20, 2]
+  "left_elbow": [20, 2]
+  "right_shoulder_pitch": [20, 2]
+  "right_shoulder_roll": [20, 2]
+  "right_shoulder_yaw": [20, 2]
+  "right_elbow": [20, 2]
+
+# Mocap clip selection
+mimic:
+  # Path is resolved relative to the project root.
+  motion_pkl: data/mocap_data/rohan/ImitationBox_01/reference_motion.pkl
+  object_csv: data/mocap_data/rohan/ImitationBox_01/object_motion.csv
+  mode: RETURN_TO_ORIGIN  # one of BOX_PICKUP, BOX_DROPOFF, RETURN_TO_ORIGIN
+
+uneven_terrain:
+  enable: true
+
+observation_noise:
+  enabled: true
+  multiplier: 1.0
+  type: "uniform"
+  scales:
+    root_orient: 0.05
+    root_ang_vel: 0.05
+    motor_pos: 0.02
+    motor_vel: 0.05
+    motor_tau: 5.0
+
+perturbation:
+  enable: true
+  bodies: ["pelvis", "torso_link"]
+  force_magnitude: 10
+  torque_magnitude: 2
+  interval: 5
+
+dynamics_randomization:
+  enable: true
+  interval: 0.5

envs/h1/gen_xml.py

@@ -5,6 +5,7 @@
 import models
 
 H1_DESCRIPTION_PATH = os.path.join(os.path.dirname(models.__file__), "mujoco_menagerie/unitree_h1/scene.xml")
+TERRAINS_PATH = os.path.join(os.path.dirname(models.__file__), "terrains/hfield.png")
 
 LEG_JOINTS = [
     "left_hip_yaw",
@@ -47,6 +48,103 @@ def create_rangefinder_array(mjcf_model, num_rows=4, num_cols=4, spacing=0.4):
     return mjcf_model
 
 
+def create_hfield(mjcf_model):
+    mjcf_model.asset.add(
+        "hfield",
+        name="hf1",
+        size="2.5 2.5 0.08 0.01",
+        file=TERRAINS_PATH,
+    )
+    mjcf_model.worldbody.add("body", name="hfield", pos=[0, 0, -0.1])
+    mjcf_model.find("body", "hfield").add(
+        "geom",
+        name="hfield",
+        type="hfield",
+        dclass="collision",
+        condim="3",
+        conaffinity="15",
+        hfield="hf1",
+        friction=".8 .1 .1",
+    )
+    return mjcf_model
+
+
+def create_mocap_bodies(mjcf_model, body_names):
+    """Add mocap-target bodies for visualizing the reference motion.
+
+    Each target is a single sphere marker (not a geometry-faithful clone of
+    the source body). Geometry-faithful cloning is brittle when source geoms
+    inherit type/size from default classes (e.g. the H1 menagerie ankle).
+    """
+    for bn in body_names:
+        b = mjcf_model.find("body", bn)
+        if b is None:
+            continue
+        attr = b.get_attributes()
+        attr["name"] = "mocap_" + b.name
+        attr["mocap"] = "true"
+        mocap_b = mjcf_model.worldbody.add("body")
+        mocap_b.set_attributes(**attr)
+        mocap_b.add("inertial", pos="0 0 0", mass="0")
+        mocap_b.add("geom", dclass="mocap", type="sphere", size="0.04")
+    return mjcf_model
+
+
+def add_movable_box(mjcf_model, name):
+    body = mjcf_model.worldbody.add("body", name=name, pos=[1.5, -0.8, 0.81], euler=[0, 0, 1.57])
+    body.add("inertial", pos="0 0 0", mass="1.36", diaginertia="0.01655043 0.01566667 0.00995043")
+    body.add("joint", type="free")
+    for s in ["visual", "collision"]:
+        geom0 = body.add("geom", name=name + "-geom0-" + s, size=[0.20, 0.1, 0.01], pos=[0, 0, 0.1])
+        geom1 = body.add("geom", name=name + "-geom1-" + s, size=[0.18, 0.1, 0.1], pos=[0, 0, 0])
+        for geom in [geom0, geom1]:
+            geom.dclass = mjcf_model.default.default["h1"].default[s]
+            geom.type = "box"
+            geom.rgba = [1, 0, 0, 1]
+    body.add("site", name="box", size="0.01", pos=[0, 0, 0])
+
+    # Two fixed tables for pickup / dropoff
+    body = mjcf_model.worldbody.add("body", name="table1", pos=[1.5, -0.8, 0.65])
+    for s in ["visual", "collision"]:
+        body.add(
+            "geom",
+            name="table1-" + s,
+            dclass=s,
+            size="0.25 0.25 0.05",
+            type="box",
+            pos=[0, 0, 0],
+            rgba="0.18 0. 0.30 1",
+            density="0.001",
+        )
+    body.add("site", name="table1", size="0.01", pos=[0, 0, 0.05])
+
+    body = mjcf_model.worldbody.add("body", name="table2", pos=[0.47, -2.7, 0.65])
+    for s in ["visual", "collision"]:
+        body.add(
+            "geom",
+            name="table2-" + s,
+            dclass=s,
+            size="0.25 0.25 0.05",
+            type="box",
+            pos=[0, 0, 0],
+            rgba="0.18 0. 0.30 1",
+            density="0.001",
+        )
+    body.add("site", name="table2", size="0.01", pos=[0, 0, 0.05])
+    return mjcf_model
+
+
+def add_hand_sites(mjcf_model):
+    """Add hand sites at the end-effector of each elbow link.
+
+    The H1 menagerie model has no hand bodies; we approximate hand position
+    with sites offset from the elbow link, mirroring the internal mocap rig.
+    """
+    mjcf_model.find("body", "left_elbow_link").add("site", name="left_hand", pos="0.2605 0 -0.0185", size="0.01")
+    mjcf_model.find("body", "right_elbow_link").add("site", name="right_hand", pos="0.2605 0 -0.0185", size="0.01")
+    return mjcf_model
+
+
 def remove_joints_and_actuators(mjcf_model, config):
     # remove joints
     for limb in config["unused_joints"]:
@@ -98,6 +196,28 @@ def builder(export_path, config):
     # set name of freejoint
     mjcf_model.find("body", "pelvis").freejoint.name = "root"
 
+    # add hand sites (for tasks that need end-effector targets)
+    if config.get("hand_sites", False):
+        mjcf_model = add_hand_sites(mjcf_model)
+
+    # add a movable box and surrounding tables
+    if "movable_box" in config and config["movable_box"]:
+        mjcf_model = add_movable_box(mjcf_model, config["movable_box"])
+
+    # add mocap bodies (visualization of reference motion)
+    if config.get("mocap_bodies"):
+        mocap_default = mjcf_model.default.default["h1"].add("default", dclass="mocap")
+        mocap_default.geom.contype = "0"
+        mocap_default.geom.conaffinity = "0"
+        mocap_default.geom.group = "3"
+        mocap_default.geom.rgba = "0 1 0 0.5"
+        mocap_default.geom.density = "0"
+        mjcf_model = create_mocap_bodies(mjcf_model, config["mocap_bodies"])
+
+    # add an uneven heightfield underneath the robot
+    if config.get("hfield", False):
+        mjcf_model = create_hfield(mjcf_model)
+
     # add rangefinder
     if "rangefinder" in config:
         if config["rangefinder"]: