Initial commit

tests/test_s10_model.py

@@ -0,0 +1,151 @@
+"""Standalone S10 MJCF/FK checks that do not require mjlab or torch.
+
+Run with the system Python:
+  python3 tests/test_s10_model.py
+
+The test validates:
+1. The S10 model in ``s10_supermimic/robot/model/s10.xml`` loads.
+2. Joint names match the training/deployment contract.
+3. The default pose puts the wheel bottom at ~0.434 m above the ground.
+4. The analytic wheel-bottom formula in ``core/rolling_rewards.py`` matches
+   MuJoCo forward kinematics for random joint angles and arbitrary base tilt.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import mujoco
+import numpy as np
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL = ROOT / "s10_supermimic" / "robot" / "model" / "s10.xml"
+LEGS = ("fl", "fr", "hl", "hr")
+LEG_JOINT_NAMES = tuple(
+    f"{leg}_{joint}_joint" for leg in LEGS for joint in ("hipx", "hipy", "knee")
+)
+WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in LEGS)
+DEFAULT_Q = {
+    "fl": (0.1, -0.35, 0.65),
+    "fr": (-0.1, -0.35, 0.65),
+    "hl": (0.1, 0.35, -0.65),
+    "hr": (-0.1, 0.35, -0.65),
+}
+# Geometry from S10.xml.
+X0 = {"fl": 0.2277, "fr": 0.2277, "hl": -0.2277, "hr": -0.2277}
+Y0 = {"fl": 0.0574, "fr": -0.0574, "hl": 0.0574, "hr": -0.0574}
+SIDE = {"fl": 1.0, "fr": -1.0, "hl": 1.0, "hr": -1.0}
+HIPY_Y = 0.0779
+WHEEL_Y = 0.045891
+LINK = 0.18
+WHEEL_RADIUS = 0.081
+WHEEL_HALF_WIDTH = 0.0185
+
+
+def analytic_wheel_bottom(q, leg, base_z=0.5, up=(0.0, 0.0, 1.0)) -> float:
+    """NumPy translation of the torch implementation in rolling_rewards.py."""
+    hipx, hipy, knee = (float(v) for v in q)
+    side = SIDE[leg]
+    x = LINK * math.sin(knee)
+    z = -LINK - LINK * math.cos(knee)
+    x3 = math.cos(hipy) * x - math.sin(hipy) * z
+    z3 = math.sin(hipy) * x + math.cos(hipy) * z
+    y3 = side * (HIPY_Y + WHEEL_Y)
+    y4 = math.cos(hipx) * y3 + math.sin(hipx) * z3
+    z4 = -math.sin(hipx) * y3 + math.cos(hipx) * z3
+    point = np.array([X0[leg] + x3, Y0[leg] + y4, z4])
+    up = np.asarray(up, dtype=float)
+    up /= np.linalg.norm(up)
+    axle_z = -up[1] * math.cos(hipx) + up[2] * math.sin(hipx)
+    radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
+    radius_z += WHEEL_HALF_WIDTH * abs(axle_z)
+    return base_z + float(point @ up) - radius_z
+
+
+def wheel_geom_id(model: mujoco.MjModel, leg: str) -> int:
+    body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel")
+    for geom_id in range(model.ngeom):
+        if model.geom_bodyid[geom_id] != body_id:
+            continue
+        if model.geom_type[geom_id] == mujoco.mjtGeom.mjGEOM_CYLINDER:
+            return geom_id
+    raise AssertionError(f"No cylinder collision geom for {leg}_wheel")
+
+
+def measured_wheel_bottom(model, data, geom_id: int) -> float:
+    axis = data.geom_xmat[geom_id].reshape(3, 3)[:, 2]
+    axle_z = abs(float(axis[2]))
+    radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
+    radius_z += WHEEL_HALF_WIDTH * axle_z
+    return float(data.geom_xpos[geom_id][2] - radius_z)
+
+
+def set_pose(model, data, q_map, base_z=0.5, quat=(1.0, 0.0, 0.0, 0.0)):
+    data.qpos[:] = 0.0
+    data.qpos[2] = base_z
+    data.qpos[3:7] = quat
+    for joint_name, value in q_map.items():
+        joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
+        data.qpos[model.jnt_qposadr[joint_id]] = value
+    mujoco.mj_forward(model, data)
+
+
+def main() -> None:
+    assert MODEL.is_file(), MODEL
+    model = mujoco.MjModel.from_xml_path(str(MODEL))
+    data = mujoco.MjData(model)
+
+    joint_names = [model.joint(i).name for i in range(model.njnt)]
+    expected = ("floating_base",) + LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
+    assert set(joint_names) == set(expected), joint_names
+    assert model.nq == 23 and model.nv == 22 and model.njnt == 17
+    assert model.nu == 0, "Robot XML must not contain actuators; mjlab adds them."
+
+    default_q = {
+        f"{leg}_{joint}_joint": value
+        for leg in LEGS
+        for joint, value in zip(("hipx", "hipy", "knee"), DEFAULT_Q[leg])
+    }
+    set_pose(model, data, default_q)
+    for leg in LEGS:
+        geom_id = wheel_geom_id(model, leg)
+        bottom = measured_wheel_bottom(model, data, geom_id)
+        pred = analytic_wheel_bottom(DEFAULT_Q[leg], leg)
+        assert abs(bottom - pred) < 1e-9, (leg, bottom, pred)
+        assert abs((0.5 - bottom) - 0.434145) < 1e-4, (leg, bottom)
+
+    rng = np.random.default_rng(0)
+    max_error = 0.0
+    for _ in range(200):
+        quat = rng.normal(size=4)
+        quat /= np.linalg.norm(quat)
+        q_map = {}
+        q_by_leg = {}
+        for leg in LEGS:
+            q = np.array([
+                rng.uniform(-0.6109, 0.6109),
+                rng.uniform(-2.5307, 2.5307),
+                rng.uniform(-2.7227, 2.7227),
+            ])
+            q_by_leg[leg] = q
+            for joint, value in zip(("hipx", "hipy", "knee"), q):
+                q_map[f"{leg}_{joint}_joint"] = value
+        set_pose(model, data, q_map, base_z=0.5, quat=quat)
+
+        rotation = data.xmat[
+            mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
+        ].reshape(3, 3)
+        up = rotation.T @ np.array([0.0, 0.0, 1.0])
+        for leg in LEGS:
+            geom_id = wheel_geom_id(model, leg)
+            measured = measured_wheel_bottom(model, data, geom_id)
+            pred = analytic_wheel_bottom(q_by_leg[leg], leg, base_z=0.5, up=up)
+            max_error = max(max_error, abs(measured - pred))
+    assert max_error < 1e-9, max_error
+    print(f"S10 model/FK checks passed (max error {max_error:.2e} m)")
+
+
+if __name__ == "__main__":
+    main()