"""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()