297 lines
11 KiB
Python
297 lines
11 KiB
Python
"""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()
|