Initial commit

s10_supermimic/core/imagine.py

@@ -0,0 +1,198 @@
+"""Differentiable rewards, full-return actor objective and TD(lambda) critic targets."""
+
+import math
+from contextlib import contextmanager
+import torch
+from .world_model import GRAVITY, LIN_VEL, ANG_VEL, LEG_POS, LEG_VEL, WHEEL_VEL
+from .policy import frame, push_history, critic_input
+
+
+@contextmanager
+def frozen(*modules):
+  """Freeze weights without disabling gradients through inputs/actions."""
+  parameters = [(p, p.requires_grad) for m in modules for p in m.parameters()]
+  try:
+    for p, _ in parameters:
+      p.requires_grad_(False)
+    yield
+  finally:
+    for p, enabled in parameters:
+      p.requires_grad_(enabled)
+
+
+def survival(state):
+  gravity = torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
+  # Smooth proxy for the environment's 70-degree fall cutoff, width ~5 deg.
+  return torch.sigmoid((-gravity[:, 2] - math.cos(math.radians(70))) / 0.08)
+
+
+def continuation_value(state, constants):
+  """Imagined continuation probability for the current task.
+
+  The v4 proxy treats a 70-degree tilt as an immediate fall. The upward task
+  instead terminates only after 4 s of bad orientation, which is 200 control
+  steps, far beyond the 32/64-step imagination horizon, and the counter is not
+  part of the state. Applying the v4 proxy there truncates the return after
+  the first step whenever the robot is inverted, so the actor cannot learn to
+  flip back.
+  """
+  if constants.get("upright_mode") == "upward":
+    return torch.ones(state.shape[0], dtype=state.dtype, device=state.device)
+  return survival(state)
+
+
+def reward_terms(state, action, previous, command, constants, survival_prob, scan=None):
+  v, w = state[:, LIN_VEL], state[:, ANG_VEL]
+  q, qd, wd = state[:, LEG_POS], state[:, LEG_VEL], state[:, WHEEL_VEL]
+  standing = (command.square().sum(-1) < 1e-6).float()
+  # Endpoint PD torque surrogate: actual MuJoCo forces also reflect implicit
+  # integration over substeps. The low-weight effort term is approximate.
+  leg_torque = (constants["kp"] * (constants.get("leg_action_scale", 0.25) * action[:, :12] - q)
+                - constants["kd"] * qd).clamp(-constants["effort"], constants["effort"])
+  wheel_limit = constants.get("wheel_velocity_limit", 35.)
+  wheel_effort = constants.get("wheel_effort", 45.)
+  wheel_torque = (constants.get("wheel_kd", 2.) * ((constants.get("wheel_action_scale", 15.) * action[:, 12:]).clamp(-wheel_limit, wheel_limit) - wd)).clamp(-wheel_effort, wheel_effort)
+  absolute_q = q + constants["default_q"]
+  limits = ((constants["limits"][..., 0] - absolute_q).clamp_min(0)
+            + (absolute_q - constants["limits"][..., 1]).clamp_min(0)).sum(-1)
+  posture = q.square().sum(-1)
+  if constants.get("terrain_posture_v2", False) and not constants.get("wheel_local_rewards", False):
+    if scan is None:
+      raise ValueError("Terrain v2 posture reward requires the privileged scan")
+    from .terrain_posture import posture_cost
+    posture = posture_cost(q, scan, constants.get("posture_relief_floor", 0.125))
+  terrain = constants.get("terrain_rewards", False)
+  gravity = state[:, GRAVITY]
+  if constants.get("upright_mode") == "upward":
+    # Sign-aware recovery reward: 1 upright, 0.5 sideways, 0 upside down.
+    upright = (0.5 * (1. - gravity[:, 2])).clamp(0., 1.)
+  else:
+    upright = torch.exp(
+      -gravity[:, :2].square().sum(-1) / constants.get("upright_std", 0.45)**2)
+  terms = {
+    "track_linear_velocity": torch.exp(-((v[:, :2] - command[:, :2]).square().sum(-1)
+                                          + (0. if terrain else v[:, 2].square())) / 0.25),
+    "track_angular_velocity": torch.exp(-((w[:, 2] - command[:, 2]).square()
+                                           + (0. if terrain else w[:, :2].square().sum(-1))) / 0.5),
+    "upright": upright,
+    "leg_posture": posture,
+    "standstill": standing * (v.square().sum(-1) + w.square().sum(-1)
+                              + 0.01 * wd.square().sum(-1)),
+    "action_rate_l2": (action - previous).square().sum(-1),
+    "joint_torques_l2": leg_torque.square().sum(-1) + wheel_torque.square().sum(-1),
+    "dof_pos_limits": limits,
+    "terminated": 1. - survival_prob,
+  }
+
+  if constants.get("rolling_rewards", False) and not constants.get("wheel_local_rewards", False):
+    if scan is None:
+      raise ValueError("Rolling rewards require a terrain scan")
+    from .rolling_rewards import rolling_costs
+    terms.update(rolling_costs(state, command, scan, constants["default_q"]))
+  if constants.get("obstacle_progress", False):
+    from .obstacle_rewards import obstacle_progress
+    terms["obstacle_progress"] = obstacle_progress(state, command, scan)
+  if constants.get("wheel_body_intrusion", False):
+    from .wheel_intrusion import wheel_body_intrusion
+    terms["wheel_body_intrusion"] = wheel_body_intrusion(
+      state, constants["default_q"])
+  if constants.get("wheel_local_rewards", False):
+    from .wheel_terrain_rewards import wheel_terms
+    terms.update(wheel_terms(state, command, scan, constants["default_q"]))
+  if constants.get("bad_orientation_penalty", False):
+    from .upward import bad_orientation_value
+    terms["bad_orientation_penalty"] = bad_orientation_value(gravity)
+  from .upward import apply_reward_clamps
+  apply_reward_clamps(terms, constants)
+  return terms
+
+
+def lambda_returns(rewards, continuation, next_values, gamma=0.99, lam=0.95):
+  carry = next_values[-1]
+  result = []
+  for t in reversed(range(len(rewards))):
+    carry = rewards[t] + gamma * continuation[t] * ((1 - lam) * next_values[t] + lam * carry)
+    result.append(carry)
+  return torch.stack(result[::-1])
+
+
+def discounted_objective(rewards, continuation, terminal_value, gamma=.99):
+  alive = torch.ones_like(terminal_value)
+  total = torch.zeros_like(terminal_value)
+  discount = 1.
+  for reward, cont in zip(rewards, continuation):
+    total = total + discount * alive * reward
+    alive = alive * cont
+    discount *= gamma
+  return total + discount * alive * terminal_value
+
+
+def make_imagined_step(model, actor, constants, weights, env_action_order,
+                       dt=.02, noise=True, field=None):
+  """Pure single-step graph, shared by eager and compiled execution."""
+  def step(state, previous, history, pose, scan, command, context, bias, timers):
+    action = actor(history)
+    model_context = context if field is None else torch.cat((context, scan), -1)
+    collision_prob = None
+    if constants.get('nonwheel_collision', False):
+      nxt, delta, _, collision_logit = model.forward_all(state, action, previous, model_context)
+      collision_prob = collision_logit.squeeze(-1).sigmoid()
+    else:
+      nxt, delta = model(state, action, previous, model_context)
+    if field is not None:
+      from .terrain import advance_pose
+      pose = advance_pose(pose, delta)
+    next_scan = None if field is None else field.scan(nxt, pose)
+    cont = continuation_value(nxt, constants)
+    if collision_prob is not None:
+      cont = cont * (1. - collision_prob)
+    terms = reward_terms(nxt, action, previous, command, constants, cont,
+                         scan=next_scan)
+    if constants.get('pair_body_intrusion', False):
+      from .wheel_intrusion import pair_body_intrusion
+      terms['pair_body_intrusion'], timers = pair_body_intrusion(
+        nxt, constants['default_q'], timers, dt)
+    reward = dt * sum(weights[name] * value for name, value in terms.items())
+    latest = frame(nxt, action[:, env_action_order], command, bias, noise=noise,
+                   wheel_vel=actor.actor_wheel_vel)
+    history = push_history(history, latest, actor.history, actor.terms)
+    return nxt, action, history, pose, next_scan, reward, cont, timers
+  return step
+
+
+def rollout(model, actor, target_critic, seed, constants, weights, env_action_order,
+            horizon=16, dt=0.02, gamma=0.99, lam=0.95, noise=True, field=None,
+            step_fn=None):
+  """Actor: full discounted return + terminal value. Critic: detached TD(lambda).
+
+  Caller freezes dynamics and target critic weights, preserving state gradients.
+  Reuse one terrain scan at each state, including the rollout's initial state.
+  """
+  state, previous, history = seed["state"], seed["previous"], seed["history"]
+  command, context, bias = seed["command"], seed["context"], seed["bias"]
+  pose = seed.get("world_pose")
+  timers = seed.get("pair_timer")
+  if timers is None:
+    timers = torch.zeros(state.shape[0], 4, device=state.device)
+  def value_input(s, prev, scan):
+    base = critic_input(s, command, context, prev)
+    return base if field is None else torch.cat((base, scan), -1)
+  if step_fn is None:
+    step_fn = make_imagined_step(model, actor, constants, weights, env_action_order,
+                                 dt, noise, field)
+  rewards, continuations, inputs, next_inputs = [], [], [], []
+  scan = None if field is None else field.scan(state, pose)
+  for _ in range(horizon):
+    inputs.append(value_input(state, previous, scan).detach())
+    state, previous, history, pose, scan, reward, cont, timers = step_fn(
+      state, previous, history, pose, scan, command, context, bias, timers)
+    rewards.append(reward)
+    continuations.append(cont)
+    next_inputs.append(value_input(state, previous, scan).detach())
+  final_value = target_critic(value_input(state, previous, scan)).squeeze(-1)
+  r, c = torch.stack(rewards), torch.stack(continuations)
+  objective = discounted_objective(r, c, final_value, gamma)
+  with torch.no_grad():
+    next_values = target_critic(torch.stack(next_inputs)).squeeze(-1)
+    targets = lambda_returns(r.detach(), c.detach(), next_values, gamma, lam)
+  return -objective.mean(), torch.stack(inputs), targets