"""Watch (or play) a sumo match. Examples: python watch.py --p0 L3 --p1 L1 --save match.gif # render a GIF (no window needed) python watch.py --p0 L4 --p1 L2 --perturb icy --save icy.mp4 # MP4 needs ffmpeg python watch.py --p0 L3 --p1 L4 --live # open a window python watch.py --human L2 # you (blue) vs L2: arrows / WASD, R restarts Robot 0 (blue) is --p0 (or you), robot 1 (orange) is --p1. Perturbations follow evaluate.py: they handicap robot 0 or change the floor (see sumo_env.apply_perturbation). """ from __future__ import annotations import argparse import sys import matplotlib import numpy as np from policies import HumanPolicy, load_level, run_matches from sumo_env import ( CONTROL_DT, MAX_STEPS, PERTURBATIONS, RING_RADIUS, ROBOT_RADIUS, PhysParams, SumoEnv, apply_perturbation, ) COLORS = ["#2a78d6", "#eb6834"] # robot 0 blue, robot 1 orange RING_FILL, RING_EDGE, BG, INK = "#2b2b29", "#fcfcfb", "#f0efec", "#0b0b0b" TRAIL = 25 # frames of trail to draw class Arena: """Matplotlib artists for one match; ``draw(snapshot)`` updates them.""" def __init__(self, ax, names): import matplotlib.patches as mp self.ax = ax ax.set_aspect("equal") lim = RING_RADIUS * 1.25 ax.set_xlim(-lim, lim) ax.set_ylim(-lim, lim) ax.axis("off") ax.figure.set_facecolor(BG) ax.add_patch(mp.Circle((0, 0), RING_RADIUS, facecolor=RING_FILL, edgecolor=RING_EDGE, lw=4)) # the two start lines ("shikiri-sen") of a real sumo ring, for flavour for s in (-1, 1): ax.plot([s * 0.12, s * 0.12], [-0.15, 0.15], color="#8c8b86", lw=3) self.bodies, self.noses, self.trails = [], [], [] for k in range(2): self.trails.append(ax.plot([], [], color=COLORS[k], lw=1.5, alpha=0.5)[0]) body = mp.Circle((0, 0), ROBOT_RADIUS, facecolor=COLORS[k], edgecolor="white", lw=1.5) ax.add_patch(body) self.bodies.append(body) self.noses.append(ax.plot([], [], color="white", lw=3, solid_capstyle="round")[0]) self.title = ax.set_title("", color=INK, fontsize=11) self.banner = ax.text(0, -RING_RADIUS * 1.15, "", ha="center", va="center", fontsize=12, color=INK, fontweight="bold") self.names = names self.history = [] def draw(self, snap, extra=""): self.history.append(snap["pos"][0].copy()) hist = np.array(self.history[-TRAIL:]) for k in range(2): p, h = snap["pos"][0, k], snap["heading"][0, k] self.bodies[k].center = p nose = p + ROBOT_RADIUS * 0.95 * np.array([np.cos(h), np.sin(h)]) self.noses[k].set_data([p[0], nose[0]], [p[1], nose[1]]) self.trails[k].set_data(hist[:, k, 0], hist[:, k, 1]) t = snap["t"] * CONTROL_DT self.title.set_text(f"{self.names[0]} (blue) vs {self.names[1]} (orange) t = {t:4.1f} s{extra}") if snap["done"][0]: o = snap["outcome"][0] tie = "Tie (time up)" if snap["t"] >= MAX_STEPS else "Tie (both out together)" self.banner.set_text({1: f"{self.names[0]} wins!", -1: f"{self.names[1]} wins!", 0: tie}[o]) else: self.banner.set_text("") return [*self.bodies, *self.noses, *self.trails, self.title, self.banner] def make_env(perturb, seed): return SumoEnv(1, apply_perturbation(PhysParams.nominal(1), perturb), seed=seed) def replay(args): """Simulate a whole match first, then animate it (saved and/or shown).""" if not args.live: matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.animation import FFMpegWriter, FuncAnimation, PillowWriter pol0, pol1 = load_level(args.p0, args.weights_dir), load_level(args.p1, args.weights_dir) env = make_env(args.perturb, args.seed) res = run_matches(env, pol0, pol1, np.random.default_rng(args.seed), record=True) frames = res["frames"] + [res["frames"][-1]] * 15 # hold the final frame briefly o = res["outcome"][0] print(f"{args.p0} vs {args.p1} [{args.perturb}]: " f"{ {1: args.p0 + ' wins', -1: args.p1 + ' wins', 0: 'tie'}[o] } after {res['end_step'][0] * CONTROL_DT:.2f}s") fig, ax = plt.subplots(figsize=(5.5, 5.8)) arena = Arena(ax, (args.p0, args.p1)) extra = "" if args.perturb == "normal" else f" [{args.perturb}]" anim = FuncAnimation(fig, lambda i: arena.draw(frames[i], extra), frames=len(frames), interval=1000 * CONTROL_DT / args.speed, blit=False, repeat=args.live) if args.save: fps = int(round(args.speed / CONTROL_DT)) if args.save.endswith(".gif"): anim.save(args.save, writer=PillowWriter(fps=fps), dpi=args.dpi) else: anim.save(args.save, writer=FFMpegWriter(fps=fps), dpi=args.dpi) print(f"saved {args.save}") if args.live: plt.show() plt.close(fig) def human(args): """Interactive: you drive robot 0 with the keyboard against a chosen level.""" import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation # free the keys we use from matplotlib's default shortcuts (s = save, left = back, ...) for k in list(plt.rcParams): if k.startswith("keymap.") and k != "keymap.quit": plt.rcParams[k] = [] me, opp = HumanPolicy(), load_level(args.human, args.weights_dir) env = make_env(args.perturb, args.seed) rng = np.random.default_rng(args.seed) state = {"obs": None, "pressed": set()} def restart(): o0, o1 = env.reset() me.reset(1) opp.reset(1, rng) state["obs"] = (o0, o1) arena.history.clear() keymap = {"up": (0, 1), "w": (0, 1), "down": (0, -1), "s": (0, -1), "left": (1, 1), "a": (1, 1), "right": (1, -1), "d": (1, -1)} def update_action(): a = np.zeros(2) for key in state["pressed"]: if key in keymap: idx, val = keymap[key] a[idx] += val me.action[:] = np.clip(a, -1, 1) def on_press(ev): if ev.key == "r": restart() state["pressed"].add(ev.key) update_action() def on_release(ev): state["pressed"].discard(ev.key) update_action() fig, ax = plt.subplots(figsize=(6, 6.3)) arena = Arena(ax, ("You", args.human)) fig.text(0.5, 0.02, "arrows / WASD to drive · R = restart · Q = quit", ha="center", color=INK) fig.canvas.mpl_connect("key_press_event", on_press) fig.canvas.mpl_connect("key_release_event", on_release) restart() def tick(_): if not env.done[0]: o0, o1 = state["obs"] o0, o1, _, _ = env.step(me.act(o0), opp.act(o1)) state["obs"] = (o0, o1) return arena.draw(env.snapshot(), " (R = restart)" if env.done[0] else "") fig._anim = FuncAnimation(fig, tick, interval=1000 * CONTROL_DT, blit=False, cache_frame_data=False) plt.show() def main(argv=None): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--p0", default="L3", help="level for robot 0 (L0, L1, Dodger, L2, L3, L4, L3_gen0100, x.npz)") ap.add_argument("--p1", default="L1", help="level for robot 1") ap.add_argument("--perturb", default="normal", choices=PERTURBATIONS) ap.add_argument("--seed", type=int, default=0, help="start-position seed") ap.add_argument("--save", default=None, help="output .gif (Pillow) or .mp4 (ffmpeg)") ap.add_argument("--live", action="store_true", help="show in a window") ap.add_argument("--human", default=None, metavar="LEVEL", help="play yourself against LEVEL") ap.add_argument("--speed", type=float, default=1.0, help="playback speed multiplier") ap.add_argument("--dpi", type=int, default=72) ap.add_argument("--weights-dir", default=None) args = ap.parse_args(argv) if args.weights_dir is None: from policies import WEIGHTS_DIR args.weights_dir = WEIGHTS_DIR if args.human: human(args) elif args.save or args.live: replay(args) else: print("nothing to do: pass --save FILE, --live, or --human LEVEL", file=sys.stderr) sys.exit(2) if __name__ == "__main__": main()