feat: stage floor height + fix traj velocity export + numpyless pickle

RumblingTurtle · RumblingTurtle · commit d22f879de3da · 2025-03-07T20:20:27.000+03:00
diff --git a/.gitignore b/.gitignore
@@ -6,6 +6,7 @@ data/
 
 *.so
 *.npy
+*.pkl
 *.zip
 *.urdf
 *.xml
diff --git a/examples/play_trajectory.py b/examples/play_trajectory.py
@@ -1,3 +1,4 @@
+import pickle
 import numpy as np
 import glob
 import time
@@ -12,13 +13,14 @@
     "g1": G1_MODEL_PATH,
 }
 
-ROBOT_NAME = "g1"
-motion_files = glob.glob(os.path.join(os.path.dirname(__file__), f"./motions/{dataset[0]}/{ROBOT_NAME}/*.npy"))
+ROBOT_NAME = "h1"
+motion_files = glob.glob(os.path.join(os.path.dirname(__file__), f"./motions/{dataset[0]}/{ROBOT_NAME}/*.pkl"))
 with MujocoRenderer(robots[ROBOT_NAME]) as renderer:
     for motion_file in motion_files:
-        motion = np.load(motion_file, allow_pickle=True).item()
+        with open(motion_file, "rb") as f:
+            motion = pickle.load(f)
         data_dt = motion["dt"]
-        qs = motion["q"]
+        qs = np.array(motion["q"])
         for frame_idx in range(qs.shape[0]):
             renderer.set_configuration(qs[frame_idx], pin_notation=True)
             renderer.step()
diff --git a/examples/retarget_amass.py b/examples/retarget_amass.py
@@ -15,7 +15,7 @@
     "h1": H1_AMASS_CONFIG,
 }
 
-ROBOT_NAME = "g1"
+ROBOT_NAME = "h1"
 RENDER = True
 
 wbik_params = robots[ROBOT_NAME]
@@ -24,6 +24,7 @@
     datasets_path=os.path.join(os.path.dirname(__file__), "../data"),
     beta=wbik_params.beta,
     name_blacklist=["handrail", "jump", "box", "hop", "push"],
+    name_whitelist=["waltz10"],
     target_fps=30,
     template_scale=wbik_params.template_scale,
 )
@@ -51,6 +52,6 @@
         except NoSolutionException:
             print(f"Skipping {motion_name}")
             continue
-        out_filename = f"amass_{ROBOT_NAME}_{motion_name}.npy"
+        out_filename = f"amass_{ROBOT_NAME}_{motion_name}"
         trajectory.save(os.path.join(output_path, out_filename))
         renderer.flush_frames(fps=data_loader.target_fps)
diff --git a/examples/retarget_cmu.py b/examples/retarget_cmu.py
@@ -56,6 +56,6 @@
                 except NoSolutionException:
                     print(f"Skipping {subject} {motion_idx}")
                     continue
-                out_filename = f"cmu_{ROBOT_NAME}_{skill_name}_{subject}_{motion_idx}.npy"
+                out_filename = f"cmu_{ROBOT_NAME}_{skill_name}_{subject}_{motion_idx}"
                 trajectory.save(os.path.join(output_path, out_filename))
                 renderer.flush_frames(fps=CMU_DATASET_FPS)
diff --git a/examples/retarget_lafan.py b/examples/retarget_lafan.py
@@ -1,10 +1,7 @@
 import os
-from tqdm import tqdm
 from recap.wbik_solver import NoSolutionException
 from recap.trajectory import Trajectory
 from recap.lafan.retarget import LAFANRetarget, MOTION_PATHS
-import time
-import numpy as np
 from recap.config.robot.g1 import G1_LAFAN_CONFIG
 from recap.config.robot.h1 import H1_LAFAN_CONFIG
 
@@ -28,6 +25,13 @@
 
 with MujocoRenderer(robots[ROBOT_NAME].mjcf_path) as renderer:
     for path in MOTION_PATHS:
+        skip_motion = True
+        for motion_type in ["run", "walk", "sprint"]:
+            if motion_type in path:
+                skip_motion = False
+
+        if skip_motion:
+            continue
         retargetee.set_motion(path)
         motion_name = os.path.basename(path).split(".")[0]
         trajectory = Trajectory(1 / 30)
@@ -40,5 +44,5 @@
         except NoSolutionException:
             print(f"Skipping {path}")
             continue
-        out_filename = f"lafan_{ROBOT_NAME}_{motion_name}.npy"
+        out_filename = f"lafan_{ROBOT_NAME}_{motion_name}"
         trajectory.save(os.path.join(output_path, out_filename))
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,6 @@ dependencies = [
     "mujoco == 3.2.7",
     "torch == 2.6.0",
     "joblib == 1.4.2",
-    "scipy == 1.15.1",
     "mediapy == 1.2.2",
 ]
 
diff --git a/recap/config/robot/h1.py b/recap/config/robot/h1.py
@@ -81,6 +81,7 @@ class H1_AMASS_CONFIG(AMASSRetargetConfig):
         ],
     )
     template_scale = 1.0128
+    extra_bodies = ["imu"]
 
 
 class H1_CMU_CONFIG(CMURetargetConfig):
@@ -147,3 +148,4 @@ class H1_LAFAN_CONFIG(LAFANRetargetConfig):
         "left_shoulder": "left_shoulder_roll_link",
         "right_shoulder": "right_shoulder_roll_link",
     }
+    extra_bodies = ["imu"]
diff --git a/recap/config/wbik_config.py b/recap/config/wbik_config.py
@@ -13,6 +13,7 @@ class WBIKConfig:
     com_pos_weight (float): weight for keeping the com close to the foot support line
     body_to_model_map (dict): MJCF to task name mapping of the corresponding frames
     step_dt (float): integration step size
+    extra_bodies (list[string]): names of the body in the xml for which to export transform data aside from the task bodies
     """
 
     mjcf_path: str
@@ -57,3 +58,4 @@ class WBIKConfig:
         "left_shoulder": "",
         "right_shoulder": "",
     }
+    extra_bodies = []
diff --git a/recap/mujoco/renderer.py b/recap/mujoco/renderer.py
@@ -30,8 +30,8 @@
 
 
 class MujocoRenderer:
-    def __init__(self, mjcf_path: str, video_path: str = None):
-        self.setup_scene(mjcf_path)
+    def __init__(self, mjcf_path: str, video_path: str = None, floor_height=0.0):
+        self.setup_scene(mjcf_path, floor_height)
         self.reset()
         self.camera_tracking = True
 
@@ -78,15 +78,15 @@ def flush_frames(self, fps=None):
             )
             self.frames = []
 
-    def setup_scene(self, mjcf_path: str):
+    def setup_scene(self, mjcf_path: str, floor_height: float):
         self.spec = mujoco.MjSpec.from_file(mjcf_path)
         self.spec.visual.quality.shadowsize = 8192
         ground = self.spec.worldbody.add_geom()
 
         ground.name = "ground"
         ground.type = mujoco.mjtGeom.mjGEOM_PLANE
         ground.size = [0, 0, 0.05]
-        ground.pos = [0, 0, 0]
+        ground.pos = [0, 0, floor_height]
         ground.material = "floor_material"
 
         floor_texture = self.spec.add_texture()
diff --git a/recap/quat_utils.py b/recap/quat_utils.py
@@ -82,13 +82,17 @@ def quat_inv(q: torch.Tensor) -> torch.Tensor:
     return normalize(q * torch.tensor([[1, -1, -1, -1]], device=q.device))
 
 
-def angular_velocity(q1: torch.Tensor, q2: torch.Tensor, dt: float):
+def angular_velocity(q1, q2, dt: float):
     need_flatten = len(q1.shape) > 2
     batch_size = q1.shape[0]
     if need_flatten:
         q1 = q1.view(-1, 4)
         q2 = q2.view(-1, 4)
-    velocities = (2 / dt) * torch.column_stack(
+    if isinstance(q1, torch.Tensor):
+        stack_func = torch.column_stack
+    else:
+        stack_func = np.column_stack
+    velocities = (2 / dt) * stack_func(
         [
             q1[:, 0] * q2[:, 1] - q1[:, 1] * q2[:, 0] - q1[:, 2] * q2[:, 3] + q1[:, 3] * q2[:, 2],
             q1[:, 0] * q2[:, 2] + q1[:, 1] * q2[:, 3] - q1[:, 2] * q2[:, 0] - q1[:, 3] * q2[:, 1],
diff --git a/recap/trajectory.py b/recap/trajectory.py
@@ -1,11 +1,20 @@
 import numpy as np
 import quaternion
-from scipy.signal import savgol_filter
 from dataclasses import dataclass
+import pickle
+from recap.quat_utils import angular_velocity
 
 
 def derivative(data, sample_dt):
-    return savgol_filter(x=data, polyorder=3, deriv=1, window_length=9, delta=sample_dt, axis=0)
+    velocities = (
+        np.diff(
+            data,
+            axis=0,
+        )
+        / sample_dt
+    )
+
+    return np.concatenate([velocities, velocities[-1][None, :]])
 
 
 @dataclass
@@ -81,10 +90,9 @@ def linear_velocities(self):
     @property
     def angular_velocities(self):
         if self._angular_velocities is None:
-            quaternions = quaternion.from_float_array(self.quaternions)
-            self._angular_velocities = quaternion.angular_velocity(
-                quaternions, np.arange(quaternions.shape[0]) * self.sample_dt
-            )
+            ang_vel = angular_velocity(self.quaternions[:-1], self.quaternions[1:], self.sample_dt)
+            ang_vel = np.concatenate([ang_vel, ang_vel[-1][None, :]])
+            self._angular_velocities = ang_vel
 
         return self._angular_velocities
 
@@ -131,7 +139,7 @@ def add_sample(
             if "world" in transform.name:
                 continue
             if transform.name not in self.bodies.keys():
-                self.bodies[transform.name] = BodyTimeSeries(transform.name, pose_data.dt)
+                self.bodies[transform.name] = BodyTimeSeries(transform.name, self.dt)
             self.bodies[transform.name].add_sample(transform.position, transform.quaternion)
 
     @property
@@ -164,6 +172,26 @@ def to_dict(self):
                 out_dict["transforms"][name][k] = v
         return out_dict
 
+    def convert_ndarrays(self, data_dict):
+        for k, v in data_dict.items():
+            if isinstance(v, np.ndarray):
+                data_dict[k] = v.tolist()
+            if isinstance(v, dict):
+                data_dict[k] = self.convert_ndarrays(v)
+        return data_dict
+
     def save(self, path):
-        trajectory_dict = self.to_dict()
-        np.save(path, trajectory_dict, allow_pickle=True)
+        # Have to remove numpy arrays entirely
+        # to preserve compatibility between versions
+        trajectory_dict = self.convert_ndarrays(self.to_dict())
+        with open(path + ".pkl", "wb") as f:
+            pickle.dump(trajectory_dict, f)
+
+
+def to_numpy(data_dict):
+    for k, v in data_dict.items():
+        if isinstance(v, list):
+            data_dict[k] = np.array(v)
+        if isinstance(v, dict):
+            data_dict[k] = to_numpy(v)
+    return data_dict
diff --git a/recap/wbik_solver.py b/recap/wbik_solver.py
@@ -367,8 +367,10 @@ def build_pose_data(
         of all tracked frames
         """
         transforms = []
+        added_bodies = []
         for frame_name in self.body_to_pin_id:
             mjcf_name = self.body_to_model_map[frame_name]
+            added_bodies.append(mjcf_name)
             body_transform = self.body(frame_name)
             transforms.append(
                 BodyTransform(
@@ -378,6 +380,17 @@ def build_pose_data(
                 )
             )
 
+        for i, frame in enumerate(self.model.frames):
+            if frame.name not in added_bodies and frame.name in self.wbik_params.extra_bodies:
+                body_transform = self.configuration.data.oMf[i]
+                transforms.append(
+                    BodyTransform(
+                        name=frame.name,
+                        position=body_transform.translation.copy(),
+                        quaternion=quaternion.from_rotation_matrix(body_transform.rotation.copy()),
+                    )
+                )
+
         root_name = self.model.frames[list(self.model.parents).index(1)].name
 
         pose_data = PoseData(
diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -0,0 +1,49 @@
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+from recap.trajectory import to_numpy
+
+file_path = os.path.join(os.path.dirname(__file__), f"../examples/motions/amass/h1/amass_h1_dance_waltz10_poses.pkl")
+with open(file_path, "rb") as f:
+    motion = pickle.load(f)
+motion = to_numpy(motion)
+
+data_dt = motion["dt"]
+pos = motion["transforms"]["imu"]["position"]
+vel = motion["transforms"]["imu"]["linear_velocity"]
+integrated_pos = pos[0, :2] + np.cumsum(vel[:, :2] * data_dt, axis=0)
+
+plt.plot(pos[:, 0], pos[:, 1], label="positions")
+plt.plot(integrated_pos[:, 0], integrated_pos[:, 1], label="integrated")
+plt.legend()
+
+
+def quat_to_euler(quat):
+    w, x, y, z = (
+        quat[:, 0],
+        quat[:, 1],
+        quat[:, 2],
+        quat[:, 3],
+    )
+    ysqr = y * y
+    X = np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + ysqr))
+    Y = -np.pi / 2 + 2 * np.arctan2(np.sqrt(1 + 2.0 * (w * y - z * x)), np.sqrt(1 - 2.0 * (w * y - z * x)))
+    Z = np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (ysqr + z * z))
+    return np.vstack([X, Y, Z]).T
+
+
+quat = motion["transforms"]["imu"]["quaternion"]
+angvel = motion["transforms"]["imu"]["angular_velocity"]
+
+euler_angles = quat_to_euler(quat)
+integrated_angles = euler_angles[0] + np.cumsum(angvel * data_dt, axis=0)
+
+
+t = np.arange(0, euler_angles.shape[0] * data_dt, data_dt)
+for i, name in enumerate(["roll", "pitch", "yaw"]):
+    plt.figure()
+    plt.plot(t, euler_angles[:, i], label=name)
+    plt.plot(t, integrated_angles[:, i], label=f"{name} int")
+    plt.legend()
+plt.show()

-Original file line number
+Diff line change
 *.so
 *.npy
 +*.pkl
 *.zip
 *.urdf
 *.xml
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,6 @@ dependencies = [`
`14`	`14`	`"mujoco == 3.2.7",`
`15`	`15`	`"torch == 2.6.0",`
`16`	`16`	`"joblib == 1.4.2",`
`17`		`- "scipy == 1.15.1",`
`18`	`17`	`"mediapy == 1.2.2",`
`19`	`18`	`]`
`20`	`19`
Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,7 @@ class H1_AMASS_CONFIG(AMASSRetargetConfig):`
`81`	`81`	`],`
`82`	`82`	`)`
`83`	`83`	`template_scale = 1.0128`
	`84`	`+ extra_bodies = ["imu"]`
`84`	`85`
`85`	`86`
`86`	`87`	`class H1_CMU_CONFIG(CMURetargetConfig):`
`@@ -147,3 +148,4 @@ class H1_LAFAN_CONFIG(LAFANRetargetConfig):`
`147`	`148`	`"left_shoulder": "left_shoulder_roll_link",`
`148`	`149`	`"right_shoulder": "right_shoulder_roll_link",`
`149`	`150`	`}`
	`151`	`+ extra_bodies = ["imu"]`