Initial commit

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