Initial commit

scripts/export_onnx.py

@@ -0,0 +1,162 @@
+"""Export an S10 SuperMimic actor checkpoint to ONNX + deployment JSON.
+
+Example:
+  python scripts/export_onnx.py \
+    --checkpoint logs/s10_terrain_v4_20k_full/model_20000.pt \
+    --output logs/s10_terrain_v4_20k_full/model_20000.onnx
+
+The ONNX file contains only the deterministic actor (RunningNorm + MLP + tanh).
+Input ``obs`` is [batch, 265] for v4/v4-upward and for the new v7 (wheel speed
+restored); only legacy v7 checkpoints use [batch, 245]. Output is [batch, 16].
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from mjlab.utils.string import resolve_expr
+
+from s10_supermimic.core.policy import load_actor
+from s10_supermimic.robot.config import (
+  LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, get_s10_robot_cfg,
+)
+
+
+def task_matches_actor(task, actor_wheel_vel: bool, frame_dim: int) -> bool:
+  return task.actor.wheel_vel == actor_wheel_vel and task.actor.frame_dim == frame_dim
+
+
+def default_joint_pos(checkpoint, actor_wheel_vel: bool, frame_dim: int) -> list[float]:
+  """Use the checkpoint task's default pose when its Actor contract matches.
+
+  Legacy 49-value v7 checkpoints predate the ±5° safe stance, so they fall
+  back to the legacy v4 ±20° pose instead of silently exporting a new target.
+  """
+  robot_cfg = get_s10_robot_cfg()
+  task_name = checkpoint.get("config", {}).get("task")
+  if task_name:
+    from s10_supermimic.tasks import get_task
+    try:
+      task = get_task(task_name)
+      if task_matches_actor(task, actor_wheel_vel, frame_dim):
+        cfg = task.build_env_cfg(
+          play=True, history=int(checkpoint["history"]),
+          seed=int(checkpoint.get("config", {}).get("seed", 0)))
+        robot_cfg = cfg.scene.entities["robot"]
+    except KeyError:
+      pass
+  names = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
+  return [float(v) for v in resolve_expr(robot_cfg.init_state.joint_pos, names, 0.0)]
+
+
+def main() -> None:
+  p = argparse.ArgumentParser(description=__doc__)
+  p.add_argument("--checkpoint", type=Path, required=True)
+  p.add_argument("--output", type=Path, default=None,
+                 help="ONNX path; defaults to <checkpoint>.onnx")
+  p.add_argument("--opset", type=int, default=17)
+  p.add_argument("--device", default="cpu")
+  args = p.parse_args()
+
+  ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
+  actor = load_actor(ckpt, args.device).eval()
+  history = int(ckpt["history"])
+  if ckpt.get("format") not in ("s10-flat-policy-v1", "s10-terrain-policy-v1"):
+    raise ValueError(f"Not an S10 policy checkpoint: {ckpt.get('format')}")
+  if history < 1:
+    raise ValueError("checkpoint history must be positive")
+
+  output = args.output or args.checkpoint.with_suffix(".onnx")
+  output.parent.mkdir(parents=True, exist_ok=True)
+  dummy = torch.zeros(1, actor.frame_dim * history, device=args.device)
+
+  with torch.no_grad():
+    reference = actor(dummy).cpu().numpy()
+
+  torch.onnx.export(
+    actor,
+    (dummy,),
+    str(output),
+    input_names=["obs"],
+    output_names=["actions"],
+    opset_version=args.opset,
+    dynamo=False,
+    do_constant_folding=True,
+    dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
+  )
+
+  import onnx
+  import onnxruntime as ort
+  onnx.checker.check_model(str(output))
+  session = ort.InferenceSession(str(output), providers=["CPUExecutionProvider"])
+  onnx_out = session.run(["actions"], {"obs": dummy.cpu().numpy()})[0]
+  max_error = float(np.max(np.abs(onnx_out - reference)))
+  if max_error > 1e-5:
+    raise RuntimeError(f"ONNX/PyTorch mismatch: max error {max_error}")
+
+  revision = int(ckpt.get("config", {}).get("terrain_revision", 1))
+  leg_scales = ([0.25, 0.6, 0.8] * 4) if revision >= 2 else ([0.25] * 12)
+  wheel_scale = 35.0 if revision >= 3 else 15.0
+  task_name = ckpt.get("config", {}).get("task")
+  task = None
+  if task_name:
+    from s10_supermimic.tasks import get_task
+    try:
+      task = get_task(task_name)
+    except KeyError:
+      task = None
+  if task is not None and task_matches_actor(task, actor.actor_wheel_vel, actor.frame_dim):
+    actor_schema = task.actor.name
+  else:
+    actor_schema = "s10-v4-53" if actor.actor_wheel_vel else "s10-v7-49"
+  sidecar = {
+    "format": ckpt["format"],
+    "step": ckpt.get("step"),
+    "history": history,
+    "actor_schema": actor_schema,
+    "actor_wheel_vel": actor.actor_wheel_vel,
+    "frame_dim": actor.frame_dim,
+    "observation_dim": actor.frame_dim * history,
+    "action_dim": 16,
+    "dt": ckpt.get("dt", 0.02),
+    "actor_hidden": list(ckpt.get("actor_hidden", ())),
+    "actor_params": sum(x.numel() for x in actor.parameters()),
+    "frame_terms": [
+      {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
+      {"name": "projected_gravity", "dim": 3, "scale": 1.0},
+      {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
+      {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
+      *([{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05}]
+        if actor.actor_wheel_vel else []),
+      {"name": "last_action", "dim": 16, "scale": 1.0},
+      {"name": "command", "dim": 3, "scale": 1.0},
+    ],
+    "history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
+    "env_action_order": list(ckpt["env_action_order"]),
+    "leg_scale": leg_scales,
+    "wheel_scale": wheel_scale,
+    "wheel_clip": 35.0,
+    "default_joint_pos": default_joint_pos(
+      ckpt, actor.actor_wheel_vel, actor.frame_dim),
+    "leg_joint_names": list(LEG_JOINT_NAMES),
+    "wheel_joint_names": list(WHEEL_JOINT_NAMES),
+    "onnx": {
+      "input": "obs",
+      "output": "actions",
+      "opset": args.opset,
+      "max_abs_error_vs_torch": max_error,
+    },
+  }
+  sidecar_path = output.with_suffix(".json")
+  sidecar_path.write_text(json.dumps(sidecar, indent=2) + "\n")
+  print(f"Exported ONNX: {output}")
+  print(f"Deployment metadata: {sidecar_path}")
+  print(f"ONNX/PyTorch max error: {max_error:.3e}")
+
+
+if __name__ == "__main__":
+  main()

scripts/play_gui.py

@@ -0,0 +1,211 @@
+"""S10 browser sliders with one CPU environment and bounded JPEG storage."""
+import argparse
+import gc
+from io import BytesIO
+import os
+from pathlib import Path
+import threading
+import time
+
+from play_keyboard import DEFAULT_CHECKPOINT, keyboard_command_cfg
+from s10_supermimic.core.gui import ControlState, make_server
+
+
+def main():
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
+  parser.add_argument("--port", type=int, default=8765)
+  parser.add_argument("--device", default="cpu")
+  parser.add_argument("--terrain-kind", choices=("flat", "up", "down", "platform", "slope", "rubble"), default="rubble")
+  parser.add_argument("--terrain-level", type=int, choices=range(6), default=0)
+  parser.add_argument("--fps", type=int, default=50)
+  parser.add_argument("--shadows", action="store_true", help="Enable costly shadows/reflections")
+  parser.add_argument("--ghost-horizon", type=int, default=32,
+                      help="World-model future steps drawn as ghost poses")
+  parser.add_argument("--ghost-stride", type=int, default=4,
+                      help="Draw every Nth ghost (1 = all, 4 = 8 ghosts)")
+  parser.add_argument("--no-ghost", action="store_true",
+                      help="Disable the world-model ghost/error overlay")
+  parser.add_argument("--steps", type=int, default=0, help="Exit after N simulation steps; 0 means until Ctrl+C")
+  args = parser.parse_args()
+  if not args.checkpoint.is_file():
+    parser.error(f"Checkpoint not found: {args.checkpoint}")
+  if not 1 <= args.fps <= 50 or not 1 <= args.port <= 65535 or args.steps < 0:
+    parser.error("Require fps 1..50, port 1..65535 and nonnegative steps")
+  if args.ghost_horizon < 1 or args.ghost_stride < 1:
+    parser.error("ghost-horizon and ghost-stride must be positive")
+  os.environ.setdefault("MUJOCO_GL", "egl")
+  os.environ.setdefault("OMP_NUM_THREADS", "1")
+  import torch
+  from PIL import Image
+  from mjlab.envs import ManagerBasedRlEnv
+  from s10_supermimic.core.bridge import S10Bridge
+  from s10_supermimic.core.policy import load_actor, evaluation_env_cfg
+  from s10_supermimic.core.world_preview import WorldModelPreview
+
+  torch.set_num_threads(1)
+  state = ControlState()
+  server = make_server(state, args.port)
+  env = None
+  worker = None
+  try:
+    checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
+    actor = load_actor(checkpoint, args.device)
+    # Keep the world model so the play view can draw the predicted future.
+    checkpoint = {k: checkpoint[k] for k in
+                  ("format", "history", "config", "env_action_order", "dt", "step",
+                   "world_model")}
+    checkpoint["actor_wheel_vel"] = actor.actor_wheel_vel
+    def build_scene(kind, level):
+      cfg = evaluation_env_cfg(checkpoint, level, kind)
+      cfg.scene.num_envs = 1
+      # Playback only needs the actor and the world-model preview; training
+      # rewards/scans are not consumed here.
+      cfg.rewards = {}
+      cfg.metrics = {}
+      cfg.observations.pop("critic", None)
+      cfg.episode_length_s = 1e9
+      cfg.commands["twist"] = keyboard_command_cfg(state)
+      cfg.viewer.width, cfg.viewer.height = 640, 480
+      cfg.viewer.enable_shadows = args.shadows
+      cfg.viewer.enable_reflections = args.shadows
+      cfg.viewer.distance, cfg.viewer.elevation, cfg.viewer.azimuth = 2.5, -20., 140.
+      new_env = ManagerBasedRlEnv(cfg=cfg, device=args.device, render_mode="rgb_array")
+      if checkpoint["format"] == "s10-terrain-policy-v1":
+        from s10_supermimic.core.terrain import TerrainBridge, TerrainField
+        field = TerrainField.from_env(new_env)
+        bridge = TerrainBridge(new_env, field)
+      else:
+        field = None
+        bridge = S10Bridge(new_env)
+      if checkpoint["env_action_order"] != bridge.to_env or abs(checkpoint["dt"] - new_env.step_dt) > 1e-8:
+        new_env.close()
+        raise ValueError("Checkpoint action order/control timestep mismatch")
+      obs, _ = new_env.reset()
+      preview = None
+      if not args.no_ghost:
+        try:
+          preview = WorldModelPreview(
+            checkpoint, new_env, bridge, field, actor,
+            horizon=args.ghost_horizon, device=args.device,
+            env_action_order=checkpoint["env_action_order"])
+        except Exception as exc:
+          print(f"World-model preview disabled: {exc}", flush=True)
+      return new_env, bridge, obs, preview
+
+    kind, level = args.terrain_kind, args.terrain_level
+    env, bridge, obs, preview = build_scene(kind, level)
+    metadata = dict(checkpoint=args.checkpoint.name, iteration=checkpoint["step"],
+                    terrain=checkpoint["format"] == "s10-terrain-policy-v1", device=args.device,
+                    terrain_revision=checkpoint["config"].get("terrain_revision", 1))
+    metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
+    state.publish([0., 0., 0.], 0, **metadata)
+    worker = threading.Thread(target=server.serve_forever, daemon=True)
+    worker.start()
+    print(f"S10 GUI ready: http://127.0.0.1:{server.server_port}/", flush=True)
+    print(f"One environment on {args.device}. Ctrl+C exits. Commands expire after 2 seconds without a browser heartbeat.", flush=True)
+    step, next_render = 0, 0.
+    stats_since = time.monotonic()
+    stats_step, frames = 0, 0
+    render_fps = realtime_factor = 0.
+    previous = torch.zeros(1, 16, device=args.device)
+    ghosts = []
+    with torch.inference_mode():
+      while not args.steps or step < args.steps:
+        began = time.monotonic()
+        selection = state.take_scene()
+        if selection is not None:
+          # Release the old simulator before constructing the new one.
+          state.clear()
+          env.close()
+          del bridge, obs
+          d = None
+          env = None
+          gc.collect()
+          if args.device.startswith("cuda"):
+            torch.cuda.empty_cache()
+          kind, level = selection
+          env, bridge, obs, preview = build_scene(kind, level)
+          previous.zero_()
+          ghosts = []
+          metadata.update(scene=f"{kind} · 等级 {level}", terrain_kind=kind, terrain_level=level)
+          state.finish_scene()
+          stats_since = next_render = time.monotonic()
+          stats_step, frames = step, 0
+        if state.take_reset():
+          obs, _ = env.reset()
+          previous.zero_()
+          ghosts = []
+          if preview is not None:
+            preview.clear()
+        with state.lock:
+          paused = state.paused
+        if not paused:
+          state_before, pose_before = bridge.state(), bridge.pose()
+          if preview is not None:
+            preview.update_errors(step, state_before, pose_before)
+            ghosts = preview.predict(
+              obs["actor"], previous, env.command_manager.get_command("twist"),
+              bridge.context(), bridge.robot.data.encoder_bias[:, bridge.legs],
+              state_before, pose_before)
+            preview.push_errors(step, ghosts)
+          action = actor(obs["actor"])
+          if not torch.isfinite(action).all():
+            raise RuntimeError("Non-finite policy action")
+          obs, _, terminated, truncated, _ = env.step(bridge.env_action(action))
+          previous = action
+          if not torch.isfinite(obs["actor"]).all():
+            raise RuntimeError("Non-finite observation")
+          if (terminated | truncated).any():
+            state.clear()
+            previous.zero_()
+            ghosts = []
+            if preview is not None:
+              preview.clear()
+          step += 1
+        d = bridge.robot.data
+        actual = [d.root_link_lin_vel_b[0, 0].item(), d.root_link_lin_vel_b[0, 1].item(),
+                  d.root_link_ang_vel_b[0, 2].item()]
+        state.publish(actual, step, render_fps=render_fps, realtime_factor=realtime_factor,
+                      wm_error=(preview.summary() if preview is not None else ""),
+                      wm_horizon=(args.ghost_horizon if preview is not None else 0),
+                      **metadata)
+        if began >= next_render:
+          # Only the simulation thread touches the native renderer camera.
+          with state.lock:
+            azimuth, elevation, distance = state.camera
+          camera = env._offline_renderer._cam
+          camera.azimuth, camera.elevation, camera.distance = azimuth, elevation, distance
+          rgb = (preview.render(ghosts, stride=args.ghost_stride)
+                 if preview is not None and ghosts else env.render())
+          buffer = BytesIO()
+          Image.fromarray(rgb).save(buffer, format="JPEG", quality=80)
+          with state.lock:
+            state.jpeg = buffer.getvalue()
+          frames += 1
+          # Schedule from frame start, not completion: rendering time is part
+          # of the frame budget rather than an extra delay on every frame.
+          next_render += 1 / args.fps
+          if next_render < began:
+            next_render = began + 1 / args.fps
+        now = time.monotonic()
+        if now - stats_since >= 1.:
+          render_fps = frames / (now - stats_since)
+          realtime_factor = (step - stats_step) * env.step_dt / (now - stats_since)
+          frames, stats_step, stats_since = 0, step, now
+        time.sleep(max(0., env.step_dt - (time.monotonic() - began)))
+    print(f"Render: {render_fps:.1f} FPS | realtime: {realtime_factor:.2f}x", flush=True)
+  except KeyboardInterrupt:
+    print("Closing S10 GUI.", flush=True)
+  finally:
+    state.clear()
+    if worker is not None:
+      server.shutdown()
+      worker.join(timeout=2.)
+    server.server_close()
+    if env is not None:
+      env.close()
+
+
+if __name__ == "__main__":
+  main()

scripts/play_gui.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+set -euo pipefail
+project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$project_dir"
+if [[ -n "${S10_PYTHON:-}" && -x "${S10_PYTHON}" ]]; then
+  gui_python="$S10_PYTHON"
+elif [[ -x "$project_dir/.venv/bin/python" ]]; then
+  gui_python="$project_dir/.venv/bin/python"
+else
+  gui_python=python
+fi
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+exec "$gui_python" scripts/play_gui.py "$@"

scripts/play_keyboard.py

@@ -0,0 +1,260 @@
+"""Keyboard velocity control of the trained S10 policy (one environment)."""
+from __future__ import annotations
+
+import argparse
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+import math
+import os
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CHECKPOINT = ROOT / "logs/s10_terrain_v4_20k_full/model_20000.pt"
+
+
+@dataclass
+class KeyboardInput:
+  speed: float = 0.6
+  lateral_speed: float = 0.2
+  turn_speed: float = 0.4
+  keys: set[str] = field(default_factory=set)
+
+  def clear(self):
+    self.keys.clear()
+
+  def command(self):
+    if "space" in self.keys:
+      return (0., 0., 0.)
+    x = float("up" in self.keys) - float("down" in self.keys)
+    y = float("left" in self.keys) - float("right" in self.keys)
+    scale = max(1., math.hypot(x, y))
+    yaw = float("q" in self.keys) - float("e" in self.keys)
+    return (self.speed * x / scale, self.lateral_speed * y / scale, self.turn_speed * yaw)
+
+
+def configure_rendering(headless, backend="auto"):
+  backend = ("egl" if headless else "glfw") if backend == "auto" else backend
+  os.environ["MUJOCO_GL"] = backend
+  os.environ["SDL_FRAMEBUFFER_ACCELERATION"] = "0"
+  os.environ["SDL_RENDER_DRIVER"] = "software"
+  os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")
+  return backend
+
+
+def keyboard_command_cfg(controls):
+  from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
+
+  class KeyboardVelocity(UniformVelocityCommand):
+    def _resample_command(self, env_ids):
+      self._update_command(env_ids)
+
+    def _update_command(self, env_ids=None):
+      # Commands are in the robot body frame, exactly as during training.
+      for axis, value in enumerate(controls.command()):
+        self.vel_command_b[:, axis] = value
+
+    def reset(self, env_ids=None):
+      controls.clear()
+      return super().reset(env_ids)
+
+  @dataclass(kw_only=True)
+  class KeyboardVelocityCfg(UniformVelocityCommandCfg):
+    def build(self, env):
+      return KeyboardVelocity(self, env)
+
+  return KeyboardVelocityCfg(
+    entity_name="robot", resampling_time_range=(1e9, 1e9),
+    heading_command=False, rel_heading_envs=0., rel_standing_envs=0.,
+    rel_forward_envs=0., debug_vis=True,
+    ranges=UniformVelocityCommandCfg.Ranges(
+      lin_vel_x=(-controls.speed, controls.speed),
+      lin_vel_y=(-controls.lateral_speed, controls.lateral_speed),
+      ang_vel_z=(-controls.turn_speed, controls.turn_speed)),
+  )
+
+
+def main():
+  p = argparse.ArgumentParser(description=__doc__)
+  p.add_argument("--terrain-level", type=int, default=0, choices=range(6))
+  p.add_argument("--terrain-kind", choices=("flat","up","down","platform","slope","rubble"), default="rubble")
+  p.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
+  p.add_argument("--device", default="cpu", help="One CPU environment by default to limit GPU memory.")
+  p.add_argument("--speed", type=float, default=0.6)
+  p.add_argument("--lateral-speed", type=float, default=0.2)
+  p.add_argument("--turn-speed", type=float, default=0.4)
+  p.add_argument("--width", type=int, default=800)
+  p.add_argument("--height", type=int, default=600)
+  p.add_argument("--headless", action="store_true")
+  p.add_argument("--steps", type=int, default=0, help="0: interactive until Esc; otherwise exit after N steps.")
+  p.add_argument("--demo", action="store_true", help="Automatically exercise forward/stop/left/turn/stop.")
+  p.add_argument("--video", type=Path, help="Stream rendered frames to MP4.")
+  p.add_argument("--ghost-horizon", type=int, default=32,
+                 help="World-model future steps drawn as ghost poses")
+  p.add_argument("--ghost-stride", type=int, default=4,
+                 help="Draw every Nth ghost (1 = all, 4 = 8 ghosts)")
+  p.add_argument("--no-ghost", action="store_true",
+                 help="Disable the world-model ghost/error overlay")
+  p.add_argument("--gl-backend", choices=("auto", "glfw", "egl"), default="auto")
+  args = p.parse_args()
+  if not args.checkpoint.is_file():
+    p.error(f"Checkpoint not found: {args.checkpoint}")
+  if args.steps < 0 or (args.headless and args.steps == 0):
+    p.error("Headless mode requires --steps > 0; steps cannot be negative")
+  if args.ghost_horizon < 1 or args.ghost_stride < 1:
+    p.error("ghost-horizon and ghost-stride must be positive")
+  if min(args.width, args.height) < 320 or args.width % 2 or args.height % 2:
+    p.error("width and height must be even and at least 320")
+  for name, value, maximum in (("speed", args.speed, 1.), ("lateral-speed", args.lateral_speed, 0.2),
+                               ("turn-speed", args.turn_speed, 0.5)):
+    if not math.isfinite(value) or not 0 < value <= maximum:
+      p.error(f"{name} must be in (0, {maximum}] to stay within the training command range")
+  backend = configure_rendering(args.headless, args.gl_backend)
+
+  import torch
+  from mjlab.envs import ManagerBasedRlEnv
+  from s10_supermimic.core.bridge import S10Bridge
+  from s10_supermimic.core.policy import load_actor
+  from s10_supermimic.core.world_preview import WorldModelPreview, annotate_rgb
+
+  pg = None
+  if not args.headless:
+    try:
+      import pygame as pg
+    except ImportError:
+      p.error("Keyboard window requires pygame: .venv/bin/python -m pip install 'pygame>=2.5'")
+  checkpoint = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
+  actor = load_actor(checkpoint, args.device)
+  controls = KeyboardInput(args.speed, args.lateral_speed, args.turn_speed)
+  from s10_supermimic.core.policy import evaluation_env_cfg
+  cfg = evaluation_env_cfg(checkpoint, args.terrain_level, args.terrain_kind)
+  cfg.scene.num_envs = 1
+  cfg.seed = 0
+  cfg.episode_length_s = 1e9  # No automatic timed teleport during keyboard play.
+  cfg.commands["twist"] = keyboard_command_cfg(controls)
+  cfg.viewer.width, cfg.viewer.height = args.width, args.height
+  cfg.viewer.distance, cfg.viewer.elevation, cfg.viewer.azimuth = 2.5, -20., 140.
+
+  with ExitStack() as stack:
+    env = ManagerBasedRlEnv(cfg=cfg, device=args.device,
+                            render_mode="rgb_array" if pg or args.video else None)
+    stack.callback(env.close)
+    obs, _ = env.reset()
+    if checkpoint["format"] == "s10-terrain-policy-v1":
+      from s10_supermimic.core.terrain import TerrainBridge, TerrainField
+      field = TerrainField.from_env(env)
+      bridge = TerrainBridge(env, field)
+    else:
+      field = None
+      bridge = S10Bridge(env)
+    if checkpoint["env_action_order"] != bridge.to_env or abs(checkpoint["dt"] - env.step_dt) > 1e-8:
+      raise ValueError("Checkpoint action order/control timestep differs from this environment")
+    preview = None
+    if not args.no_ghost:
+      try:
+        preview = WorldModelPreview(
+          checkpoint, env, bridge, field, actor,
+          horizon=args.ghost_horizon, device=args.device,
+          env_action_order=checkpoint["env_action_order"])
+      except Exception as exc:
+        print(f"World-model preview disabled: {exc}", flush=True)
+    if pg:
+      stack.callback(pg.quit)
+      # Open SDL only after MuJoCo initialization; SDL displays RGB in software.
+      pg.display.init()
+      pg.font.init()
+      screen = pg.display.set_mode((args.width, args.height))
+      pg.display.set_caption("S10 | Arrows move | Q/E turn | Space stop | R reset | Esc quit")
+      font = pg.font.Font(None, 22)
+      clock = pg.time.Clock()
+      key_names = {pg.K_UP: "up", pg.K_DOWN: "down", pg.K_LEFT: "left", pg.K_RIGHT: "right",
+                   pg.K_q: "q", pg.K_e: "e", pg.K_SPACE: "space"}
+    video = None
+    if args.video:
+      import mediapy
+      args.video.parent.mkdir(parents=True, exist_ok=True)
+      video = stack.enter_context(mediapy.VideoWriter(str(args.video), shape=(args.height, args.width),
+                                                       fps=round(1 / env.step_dt)))
+    iteration = checkpoint["step"]
+    del checkpoint  # Drop optimizer/world-model tensors; inference uses actor only.
+    print(f"Loaded S10 iteration {iteration}: {args.checkpoint}", flush=True)
+    print(f"One environment | {args.device} | {backend} | {env.step_dt:.3f}s control period", flush=True)
+    print("Arrows: forward/back/strafe | Q/E: turn | release/Space: zero command | R: reset | Esc: quit", flush=True)
+    print("Keyboard changes enter the next observation (one control-step latency).", flush=True)
+    previous = torch.zeros(1, 16, device=args.device)
+    ghosts = []
+    step, running = 0, True
+    with torch.inference_mode():
+      while running and (args.steps == 0 or step < args.steps):
+        if pg:
+          for event in pg.event.get():
+            if event.type == pg.QUIT or (event.type == pg.KEYDOWN and event.key == pg.K_ESCAPE):
+              running = False
+            elif event.type == pg.WINDOWFOCUSLOST:
+              controls.clear()
+            elif event.type == pg.KEYDOWN:
+              if event.key == pg.K_r:
+                controls.clear()
+                obs, _ = env.reset()
+                previous.zero_()
+                ghosts = []
+                if preview is not None:
+                  preview.clear()
+              elif event.key in key_names:
+                controls.keys.add(key_names[event.key])
+            elif event.type == pg.KEYUP and event.key in key_names:
+              controls.keys.discard(key_names[event.key])
+        if not running:
+          break
+        if args.demo:
+          phases = ({"up"}, set(), {"left"}, {"q"}, set())
+          controls.keys = set(phases[(step // 100) % len(phases)])
+        state_before, pose_before = bridge.state(), bridge.pose()
+        if preview is not None:
+          preview.update_errors(step, state_before, pose_before)
+          ghosts = preview.predict(
+            obs["actor"], previous, env.command_manager.get_command("twist"),
+            bridge.context(), bridge.robot.data.encoder_bias[:, bridge.legs],
+            state_before, pose_before)
+          preview.push_errors(step, ghosts)
+        action = actor(obs["actor"])
+        if not torch.isfinite(action).all():
+          raise RuntimeError("Non-finite policy action")
+        obs, _, terminated, _, _ = env.step(bridge.env_action(action))
+        previous = action
+        if not torch.isfinite(obs["actor"]).all():
+          raise RuntimeError("Non-finite observation")
+        if terminated.any():
+          controls.clear()
+          previous.zero_()
+          ghosts = []
+          if preview is not None:
+            preview.clear()
+          print("Fall reset; keys cleared.", flush=True)
+        if preview is not None and step % 20 == 0 and step > 0:
+          print("WM " + preview.summary(), flush=True)
+        if pg or video:
+          rgb = (preview.render(ghosts, stride=args.ghost_stride)
+                 if preview is not None and ghosts else env.render())
+          if video:
+            lines = ([f"WM {preview.summary()}"] if preview is not None else [])
+            video.add_image(annotate_rgb(rgb, lines))
+          if pg:
+            pg.surfarray.blit_array(screen, rgb.swapaxes(0, 1))
+            vx, vy, yaw = controls.command()
+            actual = bridge.robot.data
+            lines = [f"S10 model {iteration} | Arrows move | Q/E turn | Space stop | R reset | Esc exit",
+                     f"Command: vx {vx:+.2f}  vy {vy:+.2f} m/s  yaw {yaw:+.2f} rad/s",
+                     f"Actual: vx {actual.root_link_lin_vel_b[0, 0]:+.2f}  vy {actual.root_link_lin_vel_b[0, 1]:+.2f}"
+                     f"  yaw {actual.root_link_ang_vel_b[0, 2]:+.2f}"]
+            if preview is not None:
+              lines.append("WM " + preview.summary())
+            for i, line in enumerate(lines):
+              screen.blit(font.render(line, True, (240, 240, 240), (20, 24, 30)), (8, 8 + i * 22))
+            pg.display.flip()
+            clock.tick(round(1 / env.step_dt))
+        step += 1
+    print(f"Completed {step} control steps.", flush=True)
+
+
+if __name__ == "__main__":
+  main()

scripts/play_keyboard.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+set -euo pipefail
+project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$project_dir"
+if [[ -n "${S10_PYTHON:-}" && -x "${S10_PYTHON}" ]]; then
+  keyboard_python="$S10_PYTHON"
+elif [[ -x "$project_dir/.venv/bin/python" ]]; then
+  keyboard_python="$project_dir/.venv/bin/python"
+else
+  keyboard_python=python
+fi
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+exec "$keyboard_python" scripts/play_keyboard.py "$@"

scripts/run_train_v4_50k.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# S10 terrain-v4 50k training from scratch (16 GB recipe).
+# Override S10_OUTPUT for a fixed output directory.
+set -euo pipefail
+
+PROJECT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$PROJECT_DIR"
+
+export PYTHONPATH="$PROJECT_DIR"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export MUJOCO_GL="${MUJOCO_GL:-egl}"
+
+PY="${S10_PYTHON:-python}"
+OUT="${S10_OUTPUT:-logs/s10_terrain_v4_50k_$(date +%Y%m%d_%H%M%S)}"
+mkdir -p "$OUT"
+
+exec "$PY" scripts/train.py \
+  --task s10-v4 \
+  --terrain --terrain-revision 4 --seed 0 \
+  --num-envs 4096 --num-steps-per-env 1 \
+  --buffer-len 512 --wm-window 8 \
+  --wm-batch 2048 --wm-updates 4 \
+  --imag-batch 2048 --horizon 32 \
+  --warmup 50 --critic-epochs 5 --critic-minibatches 4 \
+  --iterations 50000 --save-interval 500 --log-interval 10 \
+  --command-curriculum --command-curriculum-start 0 \
+  --prioritized-replay \
+  --video --video-length 500 --video-interval 2700 \
+  --output "$OUT"

scripts/run_train_v7_50k.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# S10 terrain-v7 safe all-terrain 50k training from the task preset.
+# Set S10_RESUME to continue from a checkpoint; S10_VIDEO=0 disables video.
+set -euo pipefail
+
+PROJECT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$PROJECT_DIR"
+
+export PYTHONPATH="$PROJECT_DIR"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export MUJOCO_GL="${MUJOCO_GL:-egl}"
+
+PY="${S10_PYTHON:-python}"
+OUT="${S10_OUTPUT:-logs/s10_terrain_v7_50k_$(date +%Y%m%d_%H%M%S)}"
+mkdir -p "$OUT"
+
+args=(--task s10-v7 --iterations 50000 --output "$OUT")
+if [[ -n "${S10_RESUME:-}" ]]; then
+  args+=(--resume "$S10_RESUME")
+fi
+if [[ "${S10_VIDEO:-1}" != "0" ]]; then
+  args+=(--video --video-length 500 --video-interval 2700)
+fi
+
+exec "$PY" scripts/train.py "${args[@]}" "$@"

scripts/train.py

@@ -0,0 +1,5 @@
+"""Independent online S10 model-based training entry point."""
+from s10_supermimic.core.train import main
+
+if __name__ == "__main__":
+  main()