⚡ update add contact support for gmr pipline

jc-bao · jc-bao · commit ed8f46ab9efd · 2025-12-14T10:42:51.000-05:00
diff --git a/docs/guide/quick-start.md b/docs/guide/quick-start.md
@@ -21,7 +21,7 @@ This will:
 3. Display real-time visualization in MuJoCo viewer
 4. Save optimized trajectory and video
 
-## Full Workflow Example
+## Full Workflow Example For Dexterous Hand
 
 To process a task from scratch, follow these steps:
 
@@ -40,7 +40,7 @@ export DATASET_NAME=gigahand
 Convert raw human motion data to standardized format:
 
 ```bash
-uv run spider/process_datasets/oakink.py \
+uv run spider/process_datasets/${DATASET_NAME}.py \
   --task=${TASK} \
   --embodiment-type=${HAND_TYPE} \
   --data-id=${DATA_ID}
@@ -103,7 +103,52 @@ Optimize trajectory with physics constraints:
 
 ```bash
 uv run examples/run_mjwp.py \
-  +override=${DATASET_NAME} \
+  task=${TASK} \
+  dataset_name=${DATASET_NAME} \
+  data_id=${DATA_ID} \
+  robot_type=${ROBOT_TYPE} \
+  embodiment_type=${HAND_TYPE}
+```
+
+## Full Workflow Example For Humanoid Robot
+
+Set environment variables:
+
+```bash
+export TASK=dance
+export HAND_TYPE=humanoid
+export DATA_ID=0
+export ROBOT_TYPE=unitree_g1
+export DATASET_NAME=lafan
+```
+
+### Run IK
+
+```bash
+# with GMR (remember to generate GMR data trajectoru_gmr.pkl first with their official code. )
+uv run spider/process_datasets/gmr.py \
+  --task=${TASK} \
+  --dataset-name=${DATASET_NAME} \
+  --data-id=${DATA_ID} \
+  --robot-type=${ROBOT_TYPE} \
+  --embodiment-type=${HAND_TYPE} \
+  --contact-detection-mode=one
+
+# with locomujoco
+uv run spider/process_datasets/locomujoco.py \
+  --task=${TASK} \
+  --dataset-name=${DATASET_NAME} \
+  --data-id=${DATA_ID} \
+  --robot-type=${ROBOT_TYPE} \
+  --embodiment-type=${HAND_TYPE}
+```
+
+### Run Physics-Based Retargeting
+
+```bash
+uv run examples/run_mjwp.py \
+  +override=humanoid \
+  dataset_name=${DATASET_NAME} \
   task=${TASK} \
   data_id=${DATA_ID} \
   robot_type=${ROBOT_TYPE} \
diff --git a/examples/config/override/humanoid.yaml b/examples/config/override/humanoid.yaml
@@ -20,11 +20,12 @@ joint_noise_scale: 0.1
 # sampling
 improvement_threshold: 0.01
 max_num_iterations: 16
-temperature: 0.1
+temperature: 0.6
 num_samples: 2048
 # reward
 vel_rew_scale: 0.0
 joint_rew_scale: 1.0
 pos_rew_scale: 3.0
 rot_rew_scale: 3.0
+contact_rew_scale: 1.0
 # note: above parameter is optimized for speed. you can have better performance with horizon=1.0
diff --git a/spider/assets/robots/unitree_g1/scene.xml b/spider/assets/robots/unitree_g1/scene.xml
diff --git a/spider/assets/robots/unitree_g1/scene_soft.xml b/spider/assets/robots/unitree_g1/scene_soft.xml
diff --git a/spider/config.py b/spider/config.py
@@ -97,6 +97,7 @@ class Config:
     rot_rew_scale: float = 0.1
     vel_rew_scale: float = 0.0001
     terminal_rew_scale: float = 1.0
+    contact_rew_scale: float = 0.0
 
     # === VISUALIZATION CONFIGURATION ===
     show_viewer: bool = True
@@ -261,4 +262,16 @@ def process_config(config: Config):
         config.ref_dt = task_info["ref_dt"]
         loguru.logger.info(f"overriding ref_dt: {config.ref_dt} from task_info.json")
 
+    # override contact site ids
+    if config.contact_rew_scale > 0.0:
+        if "contact_site_ids" in task_info:
+            config.contact_site_ids = task_info["contact_site_ids"]
+            loguru.logger.info(
+                f"overriding contact_site_ids: {config.contact_site_ids} from task_info.json"
+            )
+        else:
+            raise ValueError(
+                "contact_site_ids not found in task_info.json while contact_rew_scale > 0.0"
+            )
+
     return config
diff --git a/spider/io.py b/spider/io.py
@@ -33,6 +33,12 @@ def load_data(
     raw_data = np.load(data_path)
     qpos_ref = raw_data["qpos"]
     qvel_ref = raw_data["qvel"]
+    # if contact_rew_scale > 0.0, we need to make sure both contact and contact_pos are provided.
+    if config.contact_rew_scale > 0.0:
+        if "contact" not in raw_data:
+            raise ValueError("contact data not found while contact_rew_scale > 0.0")
+        if "contact_pos" not in raw_data:
+            raise ValueError("contact_pos data not found while contact_rew_scale > 0.0")
     try:
         contact = raw_data["contact"]
     except:
diff --git a/spider/process_datasets/gmr.py b/spider/process_datasets/gmr.py
@@ -57,7 +57,28 @@ def main(
     enable_rate_limiter: bool = False,
     start_frame: int = 0,
     end_frame: int = -1,
+    contact_detection_mode: str = "one",
 ):
+    """
+    Process GMR data to create a SPIDER dataset.
+    Args:
+        dataset_dir: The directory containing the GMR data.
+        dataset_name: The name of the dataset.
+        robot_type: The type of robot.
+        embodiment_type: The type of embodiment.
+        task: The task to perform.
+        data_id: The id of the data.
+        show_viewer: Whether to show the viewer.
+        save_video: Whether to save the video.
+        overwrite: Whether to overwrite the existing data.
+        enable_rate_limiter: Whether to enable the rate limiter.
+        start_frame: The start frame of the data.
+        end_frame: The end frame of the data.
+        contact_detection_mode: The mode of contact detection.
+            "auto": Automatically detect contact based on mujoco contact detection.
+            "zero": Always disable contact.
+            "one": Always enable contact.
+    """
     dataset_dir = os.path.abspath(dataset_dir)
     processed_dir = get_processed_data_dir(
         dataset_dir=dataset_dir,
@@ -100,17 +121,26 @@ def main(
     shutil.copy(src_scene_file, tgt_scene_file)
     print(f"copy from {src_scene_file} to {tgt_scene_file}")
 
-    # create task info file
-    task_info_file = f"{scene_dir}/task_info.json"
-    with open(task_info_file, "w") as f:
-        json.dump({"ref_dt": 1.0 / fps}, f, indent=2)
-    print(f"Saved task info to {task_info_file}")
-
     # run mujoco
     mj_model = mujoco.MjModel.from_xml_path(tgt_scene_file)
     mj_data = mujoco.MjData(mj_model)
     run_viewer = get_viewer(show_viewer, mj_model, mj_data)
     rate_limiter = RateLimiter(fps)
+
+    # contact site id
+    contact_site_ids = []
+    for i in range(mj_model.nsite):
+        site_name = mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_SITE, i)
+        if site_name and "contact" in site_name:
+            contact_site_ids.append(i)
+    assert len(contact_site_ids) > 0 and contact_detection_mode != "zero", "No contact site found while you enable contact detection"
+
+    # create task info file
+    task_info_file = f"{scene_dir}/task_info.json"
+    with open(task_info_file, "w") as f:
+        json.dump({"ref_dt": 1.0 / fps, "contact_site_ids": contact_site_ids}, f, indent=2)
+    print(f"Saved task info to {task_info_file}")
+
     # log info
     info_list = []
     # log video
@@ -129,17 +159,27 @@ def main(
                 )
             else:
                 mj_data.qvel[:] = 0.0
-            # compute contact (currently it is a placeholder, will be implemented later)
-            contact = np.zeros(1)
+
             # compute ctrl
             mj_data.ctrl[:] = qpos[i][7:]
             mujoco.mj_forward(mj_model, mj_data)
+
+            # compute contact
+            contact_pos = mj_data.site_xpos[contact_site_ids, :]
+            if contact_detection_mode == "one":
+                contact = np.ones(len(contact_site_ids))
+            elif contact_detection_mode == "zero":
+                contact = np.zeros(len(contact_site_ids))
+            else:
+                contact = contact_pos[:, 2] < 0.001
+
             # log
             info = {
                 "qpos": mj_data.qpos.copy(),
                 "qvel": mj_data.qvel.copy(),
                 "ctrl": mj_data.ctrl.copy(),
                 "contact": contact,
+                "contact_pos": contact_pos.copy(),
             }
             info_list.append(info)
             # render
diff --git a/spider/simulators/mjwp.py b/spider/simulators/mjwp.py
@@ -270,9 +270,10 @@ def get_reward(
 
     TODO: move reward computation to task-specific module
     """
-    qpos_ref, qvel_ref, ctrl_ref, contact_ref, _contact_pos_ref = ref
+    qpos_ref, qvel_ref, ctrl_ref, contact_ref, contact_pos_ref = ref
     qpos_sim = wp.to_torch(env.data_wp.qpos)
     qvel_sim = wp.to_torch(env.data_wp.qvel)
+
     # weighted qpos tracking
     qpos_diff = _diff_qpos(
         config, qpos_sim, qpos_ref.unsqueeze(0).repeat(qpos_sim.shape[0], 1)
@@ -284,7 +285,18 @@ def get_reward(
 
     qpos_rew = -qpos_dist * 1.0
     qvel_rew = -config.vel_rew_scale * qvel_dist * 1.0
-    reward = qpos_rew + qvel_rew
+
+    # contact reward
+    if config.contact_rew_scale > 0.0 and len(config.contact_site_ids) > 0:
+        site_xpos_torch = wp.to_torch(env.data_wp.site_xpos)
+        contact_pos = site_xpos_torch[:, config.contact_site_ids]
+        contact_dist = torch.norm(contact_pos - contact_pos_ref, p=2, dim=-1)
+        contact_dist_masked = contact_dist * contact_ref.unsqueeze(0)
+        contact_rew = -contact_dist_masked.sum(dim=1)
+    else:
+        contact_rew = 0.0
+
+    reward = qpos_rew + qvel_rew + contact_rew
 
     info = {
         "qpos_dist": qpos_dist,
