Initial commit

s10_supermimic/core/world_preview.py

@@ -0,0 +1,296 @@
+"""World-model future preview for the play scripts.
+
+The actor checkpoint already contains the world model (``world_model``) and the
+training config. We use it directly; no separate ONNX export is required.
+
+The preview has two independent outputs:
+
+* a kinematic ghost trail: every predicted state is mapped back to MJCF qpos
+  and transparently added to the offscreen renderer;
+* a prediction-error ledger: predictions are stored with their target step and
+  compared with the real state when that step arrives (xy/z/yaw/gravity/leg q).
+
+The ghost is only a visualization. It never writes into ``env.sim.data`` or the
+physics model.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+
+import mujoco
+import numpy as np
+import torch
+
+from .imagine import frame, push_history
+from .policy import Actor
+from .terrain import advance_pose
+from .terrain_world_model import TerrainWorldModel
+from .world_model import GRAVITY, LEG_POS, S10WorldModel
+
+DEFAULT_HORIZON = 32
+DEFAULT_STRIDE = 4
+GHOST_NEAR = (1.0, 0.86, 0.08, 0.20)
+GHOST_FAR = (1.0, 0.24, 0.02, 0.055)
+
+
+@dataclass
+class GhostPrediction:
+  horizon: int
+  state: np.ndarray       # 38
+  pose: np.ndarray        # x, y, yaw
+  qpos: np.ndarray        # MJCF qpos
+
+
+def load_world_model(checkpoint, device="cpu"):
+  payload = checkpoint.get("world_model")
+  if not isinstance(payload, dict):
+    raise ValueError("Checkpoint has no world_model payload")
+  fmt = payload.get("format")
+  if fmt == "s10-terrain-dynamics-v1":
+    model = TerrainWorldModel.from_checkpoint(payload, device)
+  elif fmt == "s10-flat-dynamics-v1":
+    model = S10WorldModel.from_checkpoint(payload, device)
+  else:
+    raise ValueError(f"Unsupported world-model format: {fmt}")
+  return model.to(device).eval()
+
+
+def _quat_mul(a, b):
+  w1, x1, y1, z1 = a
+  w2, x2, y2, z2 = b
+  return np.array([
+    w1*w2 - x1*x2 - y1*y2 - z1*z2,
+    w1*x2 + x1*w2 + y1*z2 - z1*y2,
+    w1*y2 - x1*z2 + y1*w2 + z1*x2,
+    w1*z2 + x1*y2 - y1*x2 + z1*w2,
+  ], dtype=np.float64)
+
+
+def _quat_from_up(up_b, yaw):
+  """Rotation mapping body-frame up to world +z, followed by world yaw."""
+  up = np.asarray(up_b, dtype=np.float64)
+  up /= max(np.linalg.norm(up), 1e-9)
+  target = np.array([0., 0., 1.])
+  dot = float(np.clip(up @ target, -1., 1.))
+  if dot < -0.999999:
+    align = np.array([0., 1., 0., 0.])
+  else:
+    axis = np.cross(up, target)
+    align = np.concatenate(([1. + dot], axis))
+    align /= max(np.linalg.norm(align), 1e-9)
+  yaw_q = np.array([math.cos(yaw / 2), 0., 0., math.sin(yaw / 2)])
+  return _quat_mul(yaw_q, align)
+
+
+class WorldModelPreview:
+  def __init__(self, checkpoint, env, bridge, field, actor, horizon=DEFAULT_HORIZON,
+               device="cpu", env_action_order=None):
+    self.env = env
+    self.bridge = bridge
+    self.field = field
+    self.actor = actor
+    self.model = load_world_model(checkpoint, device)
+    self.horizon = int(horizon)
+    self.device = device
+    self.env_action_order = list(env_action_order or checkpoint["env_action_order"])
+    self.default_q = bridge.robot.data.default_joint_pos[0, bridge.legs].detach().cpu().numpy()
+    model = env.sim.mj_model
+    free = np.flatnonzero(model.jnt_type == mujoco.mjtJoint.mjJNT_FREE)
+    if len(free) != 1:
+      raise ValueError(f"Expected one free root joint, found {len(free)}")
+    self.root_qadr = int(model.jnt_qposadr[int(free[0])])
+    # ``bridge.legs``/``bridge.wheels`` are entity joint indices, not MuJoCo
+    # joint ids (the free joint is not in the entity list). Use the entity's
+    # own qpos-address table so the ghost writes the intended joints.
+    qadr = bridge.robot.indexing.joint_q_adr.cpu().numpy()
+    self.leg_qadr = [int(qadr[int(j)]) for j in bridge.legs]
+    self.wheel_qadr = [int(qadr[int(j)]) for j in bridge.wheels]
+    self.qpos_template = env.sim.data.qpos[0].detach().cpu().numpy().copy()
+    # root_link_pos_w may be env-local while qpos is world. Store the constant
+    # offset so predicted future positions land in the renderer's world frame.
+    self.world_offset = (self.qpos_template[self.root_qadr:self.root_qadr + 2]
+                         - bridge.pose()[0, :2].detach().cpu().numpy())
+    self.tracker = PredictionTracker(self.horizon)
+
+  def _ghost(self, state, pose, horizon):
+    s = state[0].detach().cpu().numpy().astype(np.float64)
+    p = pose[0].detach().cpu().numpy().astype(np.float64)
+    qpos = self.qpos_template.copy()
+    z = float(s[0])
+    if self.field is not None:
+      z += float(self.field.sample(pose[:, :2])[0].detach().cpu())
+    quat = _quat_from_up(-s[1:4], float(p[2]))
+    qpos[self.root_qadr:self.root_qadr + 3] = (
+      p[0] + self.world_offset[0], p[1] + self.world_offset[1], z)
+    qpos[self.root_qadr + 3:self.root_qadr + 7] = quat
+    qpos[self.leg_qadr] = s[LEG_POS] + self.default_q
+    qpos[self.wheel_qadr] = self.qpos_template[self.wheel_qadr]
+    return GhostPrediction(horizon, s, p, qpos)
+
+  @torch.no_grad()
+  def predict(self, history, previous, command, context, bias, state, pose):
+    """Roll the actor + world model forward from the current real state."""
+    ghosts = []
+    history = history.clone()
+    previous = previous.clone()
+    state = state.clone()
+    pose = pose.clone()
+    scan = None if self.field is None else self.field.scan(state, pose)
+    for i in range(self.horizon):
+      action = self.actor(history)
+      model_context = context if self.field is None else torch.cat((context, scan), -1)
+      if self.field is not None:
+        nxt, delta, _, _ = self.model.forward_all(state, action, previous, model_context)
+      else:
+        nxt, delta = self.model(state, action, previous, model_context)
+      pose = advance_pose(pose, delta)
+      if self.field is not None:
+        scan = self.field.scan(nxt, pose)
+      latest = frame(nxt, action[:, self.env_action_order], command, bias,
+                     noise=False, wheel_vel=self.actor.actor_wheel_vel)
+      history = push_history(history, latest, self.actor.history, self.actor.terms)
+      previous = action
+      state = nxt
+      ghosts.append(self._ghost(state, pose, i + 1))
+    return ghosts
+
+  def update_errors(self, step, state, pose):
+    self.tracker.update(step, state[0].detach().cpu().numpy(),
+                        pose[0].detach().cpu().numpy())
+
+  def push_errors(self, step, ghosts):
+    self.tracker.push(step, ghosts)
+
+  def summary(self):
+    return self.tracker.summary()
+
+  def clear(self):
+    self.tracker.clear()
+
+  def render(self, ghosts, stride=DEFAULT_STRIDE, rgba=None):
+    """Actual frame plus transparent ghost geoms, using the env's renderer."""
+    return render_with_ghosts(self.env, ghosts, stride=stride, rgba=rgba)
+
+
+class PredictionTracker:
+  """Compare predictions with the state that actually arrives later."""
+  def __init__(self, horizon, keep=200):
+    self.horizon = int(horizon)
+    self.keep = int(keep)
+    self.pending = []
+    self.errors = {h: [] for h in (1, 4, 8, 16, 32) if h <= self.horizon}
+
+  def clear(self):
+    self.pending.clear()
+    for values in self.errors.values():
+      values.clear()
+
+  @staticmethod
+  def _error(pred_state, pred_pose, state, pose):
+    dxy = float(np.linalg.norm(pred_pose[:2] - pose[:2]))
+    dz = float(abs(pred_state[0] - state[0]))
+    dyaw = float(abs(math.atan2(math.sin(pred_pose[2] - pose[2]),
+                                math.cos(pred_pose[2] - pose[2]))))
+    pg = pred_state[GRAVITY]
+    ag = state[GRAVITY]
+    pg = pg / max(np.linalg.norm(pg), 1e-9)
+    ag = ag / max(np.linalg.norm(ag), 1e-9)
+    dg = float(math.degrees(math.acos(float(np.clip(pg @ ag, -1., 1.)))))
+    dq = float(np.sqrt(np.mean((pred_state[LEG_POS] - state[LEG_POS]) ** 2)))
+    return dxy, dz, dyaw, dg, dq
+
+  def push(self, step, ghosts):
+    for ghost in ghosts:
+      if ghost.horizon in self.errors:
+        self.pending.append((step + ghost.horizon, ghost.horizon,
+                             ghost.state, ghost.pose))
+
+  def update(self, step, state, pose):
+    keep = []
+    for target, horizon, pred_state, pred_pose in self.pending:
+      if target == step:
+        values = self.errors[horizon]
+        values.append(self._error(pred_state, pred_pose, state, pose))
+        if len(values) > self.keep:
+          del values[:-self.keep]
+      elif target > step:
+        keep.append((target, horizon, pred_state, pred_pose))
+    self.pending = keep
+
+  def summary(self):
+    parts = []
+    for horizon in sorted(self.errors):
+      values = self.errors[horizon]
+      if not values:
+        continue
+      mean = np.mean(values, axis=0)
+      parts.append(f"h{horizon} xy {mean[0]:.3f}m yaw {math.degrees(mean[2]):.1f}deg "
+                   f"leg {mean[4]:.3f}")
+    return " | ".join(parts) if parts else "collecting..."
+
+
+def annotate_rgb(rgb, lines, color=(255, 185, 90)):
+  """Small bottom-left text overlay for MP4 frames (Pillow optional)."""
+  if not lines:
+    return rgb
+  try:
+    from PIL import Image, ImageDraw, ImageFont
+  except ImportError:
+    return rgb
+  image = Image.fromarray(rgb)
+  draw = ImageDraw.Draw(image)
+  font = ImageFont.load_default()
+  y = image.height - 8 - 14 * len(lines)
+  for line in lines:
+    draw.rectangle((6, y - 2, 6 + 8 * len(line), y + 13), fill=(8, 10, 14))
+    draw.text((10, y), line, fill=color, font=font)
+    y += 14
+  return np.asarray(image)
+
+
+def render_with_ghosts(env, ghosts, stride=DEFAULT_STRIDE, rgba=None):
+  """Add ghost robot geoms to the existing offscreen scene and render once."""
+  rgb = env.render()
+  renderer = env._offline_renderer
+  if renderer is None or not ghosts:
+    return rgb
+  import mujoco
+
+  model, data = renderer._model, renderer._data
+  scene = renderer.renderer.scene
+  old_qpos, old_qvel = data.qpos.copy(), data.qvel.copy()
+  old_group = renderer._opt.geomgroup.copy()
+  # S10 visual meshes are group 2; group 0 is terrain and group 1 collision
+  # geometry. Adding only group 2 gives robot ghosts without duplicating terrain.
+  visual = np.flatnonzero(model.geom_group == 2)
+  old_geom_rgba = model.geom_rgba[visual].copy()
+  mat_ids = np.unique(model.geom_matid[visual])
+  mat_ids = mat_ids[mat_ids >= 0]
+  old_mat_rgba = {int(i): model.mat_rgba[int(i)].copy() for i in mat_ids}
+  near = np.asarray(rgba if rgba is not None else GHOST_NEAR, dtype=np.float32)
+  far = np.asarray(GHOST_FAR, dtype=np.float32)
+  try:
+    renderer._opt.geomgroup[:] = 0
+    renderer._opt.geomgroup[2] = 1
+    for ghost in ghosts[max(0, int(stride) - 1)::max(1, int(stride))]:
+      # Near future is yellow and stronger; the far tail fades to red.
+      t = ghost.horizon / max(1, ghosts[-1].horizon)
+      tint = near * (1. - t) + far * t
+      for mat_id in mat_ids:
+        model.mat_rgba[int(mat_id)] = tint
+      model.geom_rgba[visual] = tint
+      data.qpos[:] = ghost.qpos
+      data.qvel[:] = 0.
+      mujoco.mj_forward(model, data)
+      mujoco.mjv_addGeoms(model, data, renderer._opt, renderer._pert,
+                          renderer._catmask.value, scene)
+  finally:
+    renderer._opt.geomgroup[:] = old_group
+    model.geom_rgba[visual] = old_geom_rgba
+    for mat_id, value in old_mat_rgba.items():
+      model.mat_rgba[mat_id] = value
+    data.qpos[:] = old_qpos
+    data.qvel[:] = old_qvel
+    mujoco.mj_forward(model, data)
+  return renderer.render()