Initial commit
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user