Initial commit

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