92 lines
4.3 KiB
Python
92 lines
4.3 KiB
Python
"""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()
|