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()

tests/test_tasks.py

@@ -0,0 +1,91 @@
+"""Task registry and dual Actor observation schema checks (no mjlab needed)."""
+
+import math
+
+import torch
+
+from s10_supermimic.core.command_curriculum import stages_for
+from s10_supermimic.core.policy import Actor, LEGACY_FRAME_DIM, FRAME_DIM, frame, push_history
+from s10_supermimic.core.wheel_intrusion import pair_body_intrusion, wheel_body_intrusion
+from s10_supermimic.tasks import get_task, task_names
+
+
+def main():
+  assert task_names() == ("s10-v4", "s10-v4-upward", "s10-v7"), task_names()
+  v4, v7, upward = get_task("s10-v4"), get_task("s10-v7"), get_task("s10-v4-upward")
+  assert (v4.terrain_revision, v7.terrain_revision) == (4, 7)
+  assert upward.terrain_revision == 4
+  assert upward.actor.frame_dim == LEGACY_FRAME_DIM
+  assert upward.actor.wheel_vel
+  assert upward.courses is False and upward.collisions is False
+  assert upward.defaults["iterations"] == 30000
+  assert upward.defaults["action_noise"] == 0.22
+  assert upward.defaults["horizon"] == 64
+  assert (v4.actor.frame_dim, v7.actor.frame_dim) == (LEGACY_FRAME_DIM, LEGACY_FRAME_DIM)
+  assert (v4.actor.history_dim(5), v7.actor.history_dim(5)) == (265, 265)
+  assert v4.actor.wheel_vel and v7.actor.wheel_vel
+  assert v7.actor.name == "s10-v7-53"
+  assert v4.courses is False and v4.collisions is False
+  assert v7.courses and v7.collisions
+  assert v7.defaults["imag_batch"] == 2048, "v7 keeps the 16 GB batch-size exception"
+  assert v7.defaults["iterations"] == 40000
+  assert v7.defaults["command_curriculum"] is True, "v7 uses stair-gated stages capped at 1.5/0.8/1.2"
+  assert stages_for(7)[-1][1:] == ((-1.5, 1.5), (-.8, .8), (-1.2, 1.2))
+  assert stages_for(4)[-1][1:] == ((-2., 2.), (-1., 1.), (-3., 3.))
+  assert v4.defaults["iterations"] == 20000
+
+  state = torch.zeros(2, 38)
+  action = torch.zeros(2, 16)
+  command = torch.zeros(2, 3)
+  bias = torch.zeros(2, 12)
+  for history, terms, expected in ((5, v4.actor.terms, 265), (5, v7.actor.terms, 265)):
+    actor = Actor(history, actor_wheel_vel=True)
+    latest = frame(state, action, command, bias,
+                   wheel_vel=actor.actor_wheel_vel)
+    assert latest.shape == (2, actor.frame_dim)
+    assert push_history(torch.zeros(2, expected), latest, history, terms).shape == (2, expected)
+  legacy = Actor(5, actor_wheel_vel=False)
+  assert legacy.frame_dim == FRAME_DIM, "legacy 49-value checkpoints remain loadable"
+
+  # v7 safe stance: default hipx 0.05 / hipy 25 deg / knee 1.0 has zero
+  # intrusion; a wheel tucked under the body crosses the 1/4 boundary.
+  state = torch.zeros(1, 38)
+  state[0, 0], state[0, 3] = 0.44, -1.0
+  def default_q(hipy, knee=1.0):
+    return torch.tensor([0.05, -hipy, knee, -0.05, -hipy, knee,
+                         0.05, hipy, -knee, -0.05, hipy, -knee])
+  assert wheel_body_intrusion(state, default_q(math.radians(25.0))).item() == 0.0
+  # 0.8 rad splay puts the wheel near x=0.134 m: inside the new outer-quarter
+  # boundary (0.15 m), outside the old outer-third boundary (0.10 m).
+  assert wheel_body_intrusion(state, default_q(.8)).item() > 0.0
+  normal = wheel_body_intrusion(state, default_q(.5)).item()
+  tucked = wheel_body_intrusion(state, default_q(1.0)).item()
+  assert normal == 0.0 < tucked, (normal, tucked)
+
+  # Pair 5/6 restriction: zero during the 0.5 s grace, then active only when
+  # both same-side legs are inside. A single leg never triggers it.
+  timers = torch.zeros(1, 4)
+  for _ in range(24):
+    cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
+    assert cost.item() == 0.0
+  cost, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
+  assert cost.item() > 0.0, cost
+  # Front-left alone stays inside; the pair timer stays unarmed.
+  one_leg = default_q(math.radians(25.0))
+  one_leg[1] = -1.0
+  timers = torch.zeros(1, 4)
+  for _ in range(40):
+    cost, timers = pair_body_intrusion(state, one_leg, timers, .02)
+    assert cost.item() == 0.0
+  assert timers[0, 1].item() == 0.0, timers
+  # Leaving the region resets that wheel's timer.
+  timers = torch.zeros(1, 4)
+  for _ in range(30):
+    _, timers = pair_body_intrusion(state, default_q(1.0), timers, .02)
+  cost, timers = pair_body_intrusion(state, default_q(math.radians(25.0)), timers, .02)
+  assert cost.item() == 0.0 and timers[0, 0].item() == 0.0, (cost, timers)
+  print("Task registry / Actor schema / v7 safety checks passed")
+
+
+if __name__ == "__main__":
+  main()

tests/test_upward.py

@@ -0,0 +1,236 @@
+"""Pure-torch checks for the upward recovery reward and per-env curriculum."""
+
+import math
+
+import torch
+
+from s10_supermimic.core.imagine import continuation_value
+from s10_supermimic.core.upward import (
+  BAD_ORIENTATION_START_DEG,
+  HISTORY_LENGTH,
+  PROMOTE_SUCCESSES,
+  TILT_STAGES_DEG,
+  bad_orientation_mask,
+  bad_orientation_value,
+  init_upward_state,
+  reset_root_state_upward,
+  upward_curriculum,
+  upward_curriculum_stats,
+  upward_video_tilt_limit,
+)
+from s10_supermimic.core.world_model import GRAVITY
+
+
+def gravity_at(deg: float, direction: float = 0.0) -> torch.Tensor:
+  """Projected gravity for a tilt ``deg`` in a fixed horizontal direction."""
+  t = math.radians(deg)
+  return torch.tensor([[
+    math.sin(t) * math.cos(direction),
+    math.sin(t) * math.sin(direction),
+    -math.cos(t),
+  ]])
+
+
+def test_severity_is_monotonic() -> None:
+  values = torch.cat([bad_orientation_value(gravity_at(d)) for d in
+                      torch.arange(0.0, 180.5, 0.5).tolist()])
+  assert (values[1:] >= values[:-1] - 1e-6).all(), "severity must be monotonic"
+  assert values.max() <= 1.0 + 1e-5, values.max()
+  assert values.min() >= 0.0, values.min()
+  assert bad_orientation_value(gravity_at(BAD_ORIENTATION_START_DEG)).item() < 1e-6
+  assert bad_orientation_value(gravity_at(180.0)).item() > 0.99
+  v105 = bad_orientation_value(gravity_at(105.0)).item()
+  v110 = bad_orientation_value(gravity_at(110.0)).item()
+  v120 = bad_orientation_value(gravity_at(120.0)).item()
+  assert v105 < v110 < v120, (v105, v110, v120)
+
+
+def test_severity_gradient_is_finite() -> None:
+  for g in (gravity_at(0.0), gravity_at(90.0), gravity_at(180.0)):
+    x = g.clone().requires_grad_(True)
+    bad_orientation_value(x).sum().backward()
+    assert torch.isfinite(x.grad).all(), x.grad
+
+
+def test_mask_keeps_bad_set() -> None:
+  assert bad_orientation_mask(gravity_at(60.0)).item()
+  assert not bad_orientation_mask(gravity_at(30.0)).item()
+  assert not bad_orientation_mask(gravity_at(50.0, math.pi / 4)).item()
+  assert bad_orientation_mask(gravity_at(180.0)).item()
+
+
+def test_upward_continuation_does_not_use_v4_fall_proxy() -> None:
+  state = torch.zeros(3, 38)
+  state[0, GRAVITY] = torch.tensor([0.0, 0.0, -1.0])
+  state[1, GRAVITY] = torch.tensor([0.0, 0.0, 0.0])
+  state[2, GRAVITY] = torch.tensor([0.0, 0.0, 1.0])
+  upward = continuation_value(state, {"upright_mode": "upward"})
+  v4 = continuation_value(state, {})
+  assert torch.allclose(upward, torch.ones(3)), upward
+  assert v4[0] > 0.99 and v4[2] < 1e-4, v4
+
+
+class _Data:
+  def __init__(self, n: int):
+    self.root_link_pos_w = torch.zeros(n, 3)
+    self.projected_gravity_b = torch.tensor([[0.0, 0.0, -1.0]] * n)
+
+
+class _Robot:
+  def __init__(self, n: int):
+    self.data = _Data(n)
+
+
+class _Terrain:
+  def __init__(self, n: int):
+    self.terrain_levels = torch.zeros(n, dtype=torch.long)
+    self.terrain_types = torch.zeros(n, dtype=torch.long)
+    self.terrain_origins = torch.zeros(6, 1, 3)
+    self.env_origins = torch.zeros(n, 3)
+
+
+class _Scene:
+  def __init__(self, n: int):
+    self.robot = _Robot(n)
+    self.terrain = _Terrain(n)
+    self.env_origins = torch.zeros(n, 3)
+
+  def __getitem__(self, key):
+    return self.robot
+
+
+class _Termination:
+  def __init__(self, n: int):
+    self.terminated = torch.zeros(n, dtype=torch.bool)
+
+
+class _Env:
+  def __init__(self, n: int):
+    self.num_envs = n
+    self.device = "cpu"
+    self.scene = _Scene(n)
+    self.episode_length_buf = torch.ones(n, dtype=torch.long)
+    self.termination_manager = _Termination(n)
+
+
+def _episode(env: _Env, recovered: torch.Tensor) -> None:
+  """Advance one finished episode for every env and run the curriculum."""
+  env.termination_manager.terminated[:] = ~recovered
+  env.scene.robot.data.projected_gravity_b[:] = torch.where(
+    recovered[:, None],
+    torch.tensor([0.0, 0.0, -1.0]),
+    torch.tensor([0.0, 0.0, 1.0]))
+  upward_curriculum(env, torch.arange(env.num_envs))
+
+
+def test_per_env_rolling_window_advances_and_demotes() -> None:
+  env = _Env(4)
+  outcomes = torch.zeros(4, dtype=torch.bool)
+  for episode in range(HISTORY_LENGTH):
+    outcomes[0] = True          # 20/20 -> advance
+    outcomes[1] = False         # 0/20 -> stays at level 0
+    outcomes[2] = episode < 10  # 10/20 -> neutral
+    outcomes[3] = episode < PROMOTE_SUCCESSES  # exactly 16/20 -> advance
+    _episode(env, outcomes)
+  stats = upward_curriculum_stats(env)
+  assert env._upward_tilt_levels.tolist() == [1, 0, 0, 1]
+  assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 1]
+  assert stats["tilt_level_q75"] == 1
+  # Advanced envs reset their window and need 20 new episodes; only envs 1
+  # (0/20) and 2 (10/20) are ready after the first block.
+  assert abs(stats["recovery_window_rate"] - 0.25) < 1e-6
+  assert abs(stats["recovery_window_ready"] - 0.5) < 1e-6
+  assert math.isclose(upward_video_tilt_limit(env), math.radians(135.0))
+
+  # The history is per-env: env 1 now gets 20 successes and advances alone;
+  # env 2 stays neutral with 10 successes.
+  for episode in range(HISTORY_LENGTH):
+    outcomes[0] = True
+    outcomes[1] = True
+    outcomes[2] = episode < 10
+    outcomes[3] = True
+    _episode(env, outcomes)
+  assert env._upward_tilt_levels.tolist() == [2, 1, 0, 2]
+  assert env.scene.terrain.terrain_levels.tolist() == [2, 1, 0, 2]
+
+  # A full 20-episode failure window demotes env 0 from level 2 to level 1.
+  for _ in range(HISTORY_LENGTH):
+    outcomes[:] = False
+    outcomes[3] = True  # keep env 3 at level 2
+    _episode(env, outcomes)
+  assert env._upward_tilt_levels.tolist() == [1, 0, 0, 2]
+  # Terrain has six levels and is independent of the three tilt levels.
+  assert env.scene.terrain.terrain_levels.tolist() == [1, 0, 0, 3]
+
+  # Before 20 episodes no env is eligible to move.
+  fresh = _Env(4)
+  stats = upward_curriculum_stats(fresh)
+  assert stats["recovery_window_ready"] == 0.0
+  assert stats["tilt_level_q75"] == 0
+  outcomes = torch.ones(4, dtype=torch.bool)
+  for _ in range(HISTORY_LENGTH - 1):
+    _episode(fresh, outcomes)
+  assert fresh._upward_tilt_levels.tolist() == [0, 0, 0, 0]
+  _episode(fresh, outcomes)
+  assert fresh._upward_tilt_levels.tolist() == [1, 1, 1, 1]
+
+
+class _Asset:
+  def __init__(self, n: int):
+    self.data = type("Data", (), {})()
+    self.data.default_root_state = torch.zeros(n, 13)
+    self.data.default_root_state[:, 3] = 1.0  # identity wxyz quaternion
+    self.pose = None
+    self.velocity = None
+
+  def write_root_link_pose_to_sim(self, pose, env_ids):
+    self.pose = pose.clone()
+
+  def write_root_link_velocity_to_sim(self, velocity, env_ids):
+    self.velocity = velocity.clone()
+
+
+class _ResetScene:
+  def __init__(self, n: int):
+    self.asset = _Asset(n)
+    self.env_origins = torch.zeros(n, 3)
+
+  def __getitem__(self, key):
+    return self.asset
+
+
+class _ResetEnv:
+  def __init__(self, n: int):
+    self.num_envs = n
+    self.device = "cpu"
+    self.scene = _ResetScene(n)
+
+
+def test_reset_supports_per_env_and_play_fallback_limits() -> None:
+  env = _ResetEnv(2)
+  pose_range = {"x": (-0.1, 0.1), "y": (-0.1, 0.1), "z": (0.0, 0.2),
+                "roll": (-math.pi, math.pi), "pitch": (-math.pi, math.pi),
+                "yaw": (-math.pi, math.pi)}
+  ids = torch.arange(2)
+  reset_root_state_upward(env, ids, pose_range)
+  assert env.scene.asset.pose.shape == (2, 7)
+  assert env.scene.asset.velocity.shape == (2, 6)
+  assert torch.isfinite(env.scene.asset.pose).all()
+  env._upward_tilt_limits = torch.tensor(
+    [math.radians(60.0), math.radians(180.0)])
+  reset_root_state_upward(env, ids, pose_range)
+  assert torch.isfinite(env.scene.asset.pose).all()
+
+
+def main() -> None:
+  test_severity_is_monotonic()
+  test_severity_gradient_is_finite()
+  test_mask_keeps_bad_set()
+  test_upward_continuation_does_not_use_v4_fall_proxy()
+  test_per_env_rolling_window_advances_and_demotes()
+  test_reset_supports_per_env_and_play_fallback_limits()
+  print("Upward reward / continuation / per-env curriculum checks passed")
+
+
+if __name__ == "__main__":
+  main()