Initial commit

s10_supermimic/core/upward.py

@@ -0,0 +1,336 @@
+"""V4-upward helpers: sign-aware upright reward, reward clamps and sustained fall.
+
+This module is used only by the ``s10-v4-upward`` task. V4 itself keeps its
+baseline exp-upright reward and immediate 70-degree fall termination.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from mjlab.envs import mdp
+from mjlab.managers import ManagerTermBase
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+from .world_model import GRAVITY
+
+# Raw-term caps (before weight and dt) for the recovery task.
+# Chosen from v4 random-orientation samples: p99/max were approximately
+# leg_posture .78/.86, standstill 5.37/7.90, joint_torques 4772/5016,
+# unnecessary_wheel_lift 750/826, rolling_leg_motion 2.38/4.47.
+UPWARD_REWARD_CLAMPS = {
+  "leg_posture": 2.0,
+  "standstill": 4.0,
+  "action_rate_l2": 4.0,
+  "joint_torques_l2": 5000.0,
+  "dof_pos_limits": 0.5,
+  "unnecessary_wheel_lift": 2.0,
+  "rolling_leg_motion": 5.0,
+}
+
+# IsaacLab upward: upward max 4 (raw 4 x weight 1), bad-orientation -80,
+# i.e. the bad penalty is 20x the upward maximum. Here upward max is 2
+# (raw 1 x weight 2), so the same ratio is -40.
+BAD_ORIENTATION_WEIGHT = -40.0
+BAD_ORIENTATION_XY = 0.7
+# The old max(gz, |gx|-0.7, |gy|-0.7) severity saturated for gz >= 0.3 and had
+# a local maximum near 105 deg, where the |gx|-0.7 term crossed gz. The reward
+# severity is now a function of the tilt angle only: 0 at 45 deg, 1 at 180 deg,
+# monotonic in between. The correction term steepens the near-threshold part
+# while keeping a positive slope at 180 deg and a finite derivative at 45 deg.
+BAD_ORIENTATION_START_DEG = 45.0
+BAD_ORIENTATION_CORRECTION = 2.5
+
+# Per-environment recovery curriculum, IsaacLab terrain-curriculum style.
+# Each env keeps the outcomes of its last HISTORY_LENGTH episodes and advances
+# when at least PROMOTE_SUCCESSES of them recovered (16/20 = 80% > 75%).
+# There is no global batch gate: terrain and tilt levels are per environment.
+TILT_STAGES_DEG = (60.0, 135.0, 180.0)
+HISTORY_LENGTH = 20
+PROMOTE_SUCCESSES = 16
+DEMOTE_SUCCESSES = 5
+
+
+def upright_upward_value(gravity: torch.Tensor) -> torch.Tensor:
+  """Normalized upward reward in [0, 1].
+
+  ``projected_gravity_b[:, 2]`` is -1 upright, 0 sideways and +1 upside down.
+  The old v4 exp-upright term is sign-blind: upside-down also scores 1. This
+  linear form is monotonic through the whole range and has a constant gradient
+  with respect to gravity_z, which is what the recovery task needs.
+  """
+  return (0.5 * (1.0 - gravity[:, 2])).clamp(0.0, 1.0)
+
+
+def upright_upward(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+  """Environment reward term replacing v4 ``upright`` for this task."""
+  gravity = env.scene[asset_cfg.name].data.projected_gravity_b
+  return upright_upward_value(gravity)
+
+
+def bad_orientation_value(gravity: torch.Tensor) -> torch.Tensor:
+  """Monotonic tilt-based bad-orientation severity in [0, 1].
+
+  ``projected_gravity_b[:, 2]`` is -1 upright and +1 upside down, so
+  ``tilt = acos(-gz)`` goes from 0 to pi. Severity is 0 below
+  ``BAD_ORIENTATION_START_DEG``, then increases monotonically to 1 at 180 deg:
+
+    x        = clamp((tilt - 45 deg) / 135 deg, 0, 1)
+    severity = x + 2.5 * x * (1 - x)^2
+
+  Unlike the previous max(gz, |gx|-0.7, |gy|-0.7) formula, this has no
+  saturation plateau for deep inversion and no local maximum near 105 deg.
+  The polynomial correction is used instead of a fractional power because
+  ``x ** (1/3)`` has an infinite derivative at x=0 (upright), which would
+  produce non-finite gradients during imagination.
+  """
+  gz = gravity[:, 2].clamp(-1.0 + 1e-4, 1.0 - 1e-4)
+  tilt = torch.acos(-gz)
+  start = math.radians(BAD_ORIENTATION_START_DEG)
+  x = ((tilt - start) / (math.pi - start)).clamp(0.0, 1.0)
+  return x + BAD_ORIENTATION_CORRECTION * x * (1.0 - x).square()
+
+
+def bad_orientation_mask(gravity: torch.Tensor) -> torch.Tensor:
+  """Binary bad set used by the 4-second termination.
+
+  This intentionally keeps the IsaacLab-style boundary: upside-down
+  hemisphere or |gx|/|gy| beyond 0.7. The reward severity above is a smooth
+  tilt-based shaping term and is not used to decide termination.
+  """
+  return ((gravity[:, 2] > 0.0)
+          | (gravity[:, :2].abs() > BAD_ORIENTATION_XY).any(-1))
+
+
+def bad_orientation_penalty(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+  """Continuous per-step bad-orientation penalty used by the environment."""
+  gravity = env.scene[asset_cfg.name].data.projected_gravity_b
+  return bad_orientation_value(gravity)
+
+
+def is_upward_task(env) -> bool:
+  """Identify the task from its upright reward function, not from a new flag."""
+  rewards = getattr(env.cfg, "rewards", {})
+  term = rewards.get("upright") if hasattr(rewards, "get") else None
+  return term is not None and getattr(term, "func", None) is upright_upward
+
+
+def _sample_offsets(range_dict, count: int, device):
+  """Sample x/y/z/roll/pitch/yaw offsets; missing keys are zero."""
+  offsets = []
+  for key in ("x", "y", "z", "roll", "pitch", "yaw"):
+    lo, hi = (range_dict or {}).get(key, (0.0, 0.0))
+    lo, hi = float(lo), float(hi)
+    if hi <= lo:
+      offsets.append(torch.full((count,), lo, device=device))
+    else:
+      offsets.append(torch.empty(count, device=device).uniform_(lo, hi))
+  return torch.stack(offsets, -1)
+
+
+def reset_root_state_upward(env, env_ids, pose_range, velocity_range=None,
+                            asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
+  """Reset with a per-environment roll/pitch limit.
+
+  ``env._upward_tilt_limits`` is a per-env tensor updated by the curriculum.
+  Play/evaluation has no such tensor and falls back to the full range in
+  ``pose_range`` (the task sets +/-pi there).
+  """
+  env_ids = mdp.resolve_env_ids(env, env_ids)
+  asset = env.scene[asset_cfg.name]
+  count = len(env_ids)
+  pose = _sample_offsets(pose_range, count, env.device)
+  limits = getattr(env, "_upward_tilt_limits", None)
+  if limits is None:
+    roll = pose_range.get("roll", (-math.pi, math.pi))
+    limit = max(abs(float(roll[0])), abs(float(roll[1])))
+    limits = torch.full((count,), limit, device=env.device)
+  else:
+    limits = limits[env_ids]
+  pose[:, 3] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
+  pose[:, 4] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
+  default = asset.data.default_root_state[env_ids].clone()
+  positions = default[:, 0:3] + pose[:, 0:3] + env.scene.env_origins[env_ids]
+  delta = mdp.quat_from_euler_xyz(pose[:, 3], pose[:, 4], pose[:, 5])
+  orientations = mdp.quat_mul(default[:, 3:7], delta)
+  velocities = default[:, 7:13] + _sample_offsets(velocity_range, count, env.device)
+  asset.write_root_link_pose_to_sim(
+    torch.cat((positions, orientations), dim=-1), env_ids=env_ids)
+  asset.write_root_link_velocity_to_sim(velocities, env_ids=env_ids)
+
+
+def init_upward_state(env):
+  """Allocate the per-env recovery ring and tilt levels once."""
+  if getattr(env, "_upward_history", None) is not None:
+    return
+  n, device = env.num_envs, env.device
+  env._upward_history = torch.zeros(
+    n, HISTORY_LENGTH, dtype=torch.bool, device=device)
+  env._upward_traverse_history = torch.zeros_like(env._upward_history)
+  env._upward_history_cursor = torch.zeros(n, dtype=torch.long, device=device)
+  env._upward_history_filled = torch.zeros(n, dtype=torch.long, device=device)
+  env._upward_tilt_levels = torch.zeros(n, dtype=torch.long, device=device)
+  env._upward_tilt_limit_table = torch.tensor(
+    [math.radians(d) for d in TILT_STAGES_DEG], device=device)
+  env._upward_tilt_limits = env._upward_tilt_limit_table[0].expand(n).clone()
+
+
+def upward_curriculum(env, env_ids):
+  """Per-env rolling recovery curriculum (IsaacLab terrain-curriculum style).
+
+  Each environment stores the outcomes of its last ``HISTORY_LENGTH``
+  episodes. ``recovered`` means the episode did not end in a bad-orientation
+  termination and the final projected gravity is not in the bad set. An env
+  advances when at least ``PROMOTE_SUCCESSES`` of the window recovered and
+  drops when at most ``DEMOTE_SUCCESSES`` did. Terrain and tilt levels are
+  updated independently per environment; no global batch rate is involved.
+  """
+  init_upward_state(env)
+  env_ids = mdp.resolve_env_ids(env, env_ids)
+  terrain = env.scene.terrain
+  data = env.scene["robot"].data
+  active = env.episode_length_buf[env_ids] > 0
+  fallen = env.termination_manager.terminated[env_ids]
+  gravity = data.projected_gravity_b[env_ids]
+  offset = data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
+  recovered = active & ~fallen & ~bad_orientation_mask(gravity)
+  traversed = (active & ~fallen & (offset[:, 0].abs() > 5.4)
+               & (offset[:, 1].abs() < 6.5))
+  ids = env_ids[active]
+  if len(ids):
+    cursor = env._upward_history_cursor[ids]
+    env._upward_history[ids, cursor] = recovered[active]
+    env._upward_traverse_history[ids, cursor] = traversed[active]
+    env._upward_history_cursor[ids] = (cursor + 1) % HISTORY_LENGTH
+    env._upward_history_filled[ids] = (
+      env._upward_history_filled[ids] + 1).clamp_max(HISTORY_LENGTH)
+  ready = env._upward_history_filled[env_ids] >= HISTORY_LENGTH
+  successes = env._upward_history[env_ids].sum(-1)
+  promote = ready & (successes >= PROMOTE_SUCCESSES)
+  demote = ready & (successes <= DEMOTE_SUCCESSES)
+  delta = promote.long() - demote.long()
+  old_levels = terrain.terrain_levels[env_ids].clone()
+  levels = (old_levels + delta).clamp(0, terrain.terrain_origins.shape[0] - 1)
+  old_tilt_levels = env._upward_tilt_levels[env_ids].clone()
+  tilt_levels = (old_tilt_levels + delta).clamp(0, len(TILT_STAGES_DEG) - 1)
+  # A changed environment must prove itself for a fresh 20-episode window at
+  # the new level; otherwise a long success streak could chain promotions in
+  # consecutive episodes. A clamped promote/demote at the min/max level is not
+  # a change and must keep its window so the logged rate stays meaningful.
+  changed = (levels != old_levels) | (tilt_levels != old_tilt_levels)
+  if changed.any():
+    changed_ids = env_ids[changed]
+    env._upward_history[changed_ids] = False
+    env._upward_traverse_history[changed_ids] = False
+    env._upward_history_filled[changed_ids] = 0
+    env._upward_history_cursor[changed_ids] = 0
+  terrain.terrain_levels[env_ids] = levels
+  terrain.env_origins[env_ids] = terrain.terrain_origins[
+    levels, terrain.terrain_types[env_ids]]
+  env._upward_tilt_levels[env_ids] = tilt_levels
+  env._upward_tilt_limits[env_ids] = env._upward_tilt_limit_table[tilt_levels]
+  return terrain.terrain_levels.float().mean()
+
+
+def upward_curriculum_stats(env):
+  """Summary for logging and for choosing the training-video tilt level."""
+  init_upward_state(env)
+  levels = env._upward_tilt_levels
+  terrain_levels = env.scene.terrain.terrain_levels
+  ready = env._upward_history_filled >= HISTORY_LENGTH
+  recovery = env._upward_history.sum(-1)
+  traverse = env._upward_traverse_history.sum(-1)
+  ready_count = int(ready.sum())
+  recovery_rate = (float(recovery[ready].float().mean().item()) / HISTORY_LENGTH
+                   if ready_count else 0.0)
+  traverse_rate = (float(traverse[ready].float().mean().item()) / HISTORY_LENGTH
+                   if ready_count else 0.0)
+  def fractions(values, count):
+    return [float((values == i).float().mean().item()) for i in range(count)]
+  q75 = int(torch.quantile(levels.float(), 0.75).round().item())
+  return {
+    "tilt_level_mean": float(levels.float().mean().item()),
+    "tilt_level_max": int(levels.max().item()),
+    "tilt_level_q75": q75,
+    "tilt_level_frac": fractions(levels, len(TILT_STAGES_DEG)),
+    "tilt_limit_mean_deg": float(
+      env._upward_tilt_limits.mean().item() * 180.0 / math.pi),
+    "terrain_level_frac": fractions(
+      terrain_levels, int(env.scene.terrain.terrain_origins.shape[0])),
+    "recovery_window_rate": recovery_rate,
+    "traverse_window_rate": traverse_rate,
+    "recovery_window_ready": float(ready.float().mean().item()),
+  }
+
+
+def upward_video_tilt_limit(env):
+  """Tilt limit reached by at least 75% of environments, for the video scene."""
+  q75 = upward_curriculum_stats(env)["tilt_level_q75"]
+  return math.radians(TILT_STAGES_DEG[q75])
+
+
+def _cap(value: torch.Tensor, max_value: float) -> torch.Tensor:
+  return value.clamp(max=float(max_value))
+
+
+def clamped_leg_posture(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
+  from .terrain_posture import env_posture_cost
+  return _cap(env_posture_cost(env, asset_cfg), max_value)
+
+
+def clamped_standstill(env, wheel_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
+  from s10_supermimic.env.flat import standstill_cost
+  return _cap(standstill_cost(env, wheel_cfg), max_value)
+
+
+def clamped_action_rate(env, max_value: float) -> torch.Tensor:
+  return _cap(mdp.action_rate_l2(env), max_value)
+
+
+def clamped_joint_torques(env, max_value: float) -> torch.Tensor:
+  return _cap(mdp.joint_torques_l2(env), max_value)
+
+
+def clamped_dof_pos_limits(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
+  return _cap(mdp.joint_pos_limits(env, asset_cfg), max_value)
+
+
+def clamped_rolling_cost(env, term: str, max_value: float) -> torch.Tensor:
+  from .rolling_rewards import env_rolling_cost
+  return _cap(env_rolling_cost(env, term), max_value)
+
+
+def apply_reward_clamps(terms: dict[str, torch.Tensor], constants: dict) -> None:
+  """Apply the same raw caps to imagined rewards used by the env task."""
+  for name, max_value in constants.get("reward_clamps", {}).items():
+    if name in terms:
+      terms[name] = terms[name].clamp(max=float(max_value))
+
+
+class BadOrientationSustained(ManagerTermBase):
+  """Terminate only after bad orientation persists for ``duration_s``.
+
+  This is the recovery-task state machine: a random roll/pitch spawn is allowed
+  to be bad for a few seconds, but the episode ends if the robot never rights
+  itself. Functionally it replaces v4's immediate 70-degree ``fell_over``.
+  """
+
+  def __init__(self, cfg, env):
+    super().__init__(env)
+    self.steps = torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
+
+  def __call__(self, env, duration_s: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+    gravity = env.scene[asset_cfg.name].data.projected_gravity_b
+    # Termination keeps the binary IsaacLab-style bad set; the reward severity
+    # above is only a smooth shaping term.
+    bad = bad_orientation_mask(gravity)
+    self.steps = torch.where(bad, self.steps + 1, torch.zeros_like(self.steps))
+    duration_steps = math.ceil(duration_s / env.step_dt)
+    return self.steps >= duration_steps
+
+  def reset(self, env_ids=None) -> None:
+    if env_ids is None:
+      self.steps.zero_()
+    else:
+      self.steps[env_ids] = 0