Decouple physics rate from frame rate; Citrus theme via launch()

app.py

@@ -81,8 +81,9 @@ GAME_IFRAME = (
 SC = 0.8                      # render scale: 800x500 table -> 640x400 video
 VW, VH = int(800 * SC), int(500 * SC)
 FPS = 40                      # one frame per physics sub-step == real time
-MAX_FRAMES = 1100             # hard cap on video length (~27 s)
-POINT_SUBSTEP_CAP = 420       # a single point is called a draw after this
+MAX_FRAMES = 1600             # hard cap on rendered frames (40 s at 40 fps)
+POINT_SUBSTEP_CAP = 4500      # == the author's 1500-step rally limit
+MATCH_SUBSTEP_CAP = 4800      # keeps playback speed at 3x or below
 MAX_POINTS = 12               # safety net against endless draw sequences
 
 
@@ -98,9 +99,11 @@ def _font(size: int):
 
 
 F_BIG, F_MED, F_SMALL = _font(34), _font(15), _font(13)
+BLUE, ORANGE = (79, 140, 255), (249, 115, 22)
 
 
-def _render_frame(env, score_a, score_b, opp_name, trail, banner=""):
+def _render_frame(state, opp_name, trail, speed_label, banner=""):
+    bx, by, ego_y, opp_y, ego_h, opp_h, rally, sa, sb, bspeed = state
     img = Image.new("RGB", (VW, VH), (15, 22, 38))
     d = ImageDraw.Draw(img)
 
@@ -110,73 +113,93 @@ def _render_frame(env, score_a, score_b, opp_name, trail, banner=""):
 
     for i, (tx, ty) in enumerate(trail):
         a = (i + 1) / (len(trail) + 1)
-        c = int(60 + 150 * a)
-        r = 8 * SC * (0.35 + 0.6 * a)
+        c = int(45 + 150 * a)
+        r = 8 * SC * (0.3 + 0.65 * a)
         d.ellipse([tx * SC - r, ty * SC - r, tx * SC + r, ty * SC + r], fill=(c, c, c))
 
     pw = 14 * SC
-    eh = env.ego_paddle_h * SC / 2
-    oh = env.opp_paddle_h * SC / 2
-    d.rounded_rectangle([2, env.ego_y * SC - eh, 2 + pw, env.ego_y * SC + eh],
-                        radius=5, fill=(79, 140, 255))          # Mr. Pong (blue)
-    d.rounded_rectangle([VW - pw - 2, env.opp_y * SC - oh, VW - 2, env.opp_y * SC + oh],
-                        radius=5, fill=(249, 115, 22))          # baseline (orange)
+    eh, oh = ego_h * SC / 2, opp_h * SC / 2
+    d.rounded_rectangle([2, ego_y * SC - eh, 2 + pw, ego_y * SC + eh],
+                        radius=5, fill=BLUE)        # Mr. Pong
+    d.rounded_rectangle([VW - pw - 2, opp_y * SC - oh, VW - 2, opp_y * SC + oh],
+                        radius=5, fill=ORANGE)      # scripted baseline
 
     r = 8 * SC
-    d.ellipse([env.ball_x * SC - r, env.ball_y * SC - r,
-               env.ball_x * SC + r, env.ball_y * SC + r], fill=(255, 255, 255))
+    d.ellipse([bx * SC - r, by * SC - r, bx * SC + r, by * SC + r], fill=(255, 255, 255))
 
-    d.text((VW // 2 - 46, 16), str(score_a), font=F_BIG, fill=(79, 140, 255), anchor="ma")
-    d.text((VW // 2 + 46, 16), str(score_b), font=F_BIG, fill=(249, 115, 22), anchor="ma")
+    d.text((VW // 2 - 46, 14), str(sa), font=F_BIG, fill=BLUE, anchor="ma")
+    d.text((VW // 2 + 46, 14), str(sb), font=F_BIG, fill=ORANGE, anchor="ma")
     d.text((20, 20), "MR. PONG", font=F_MED, fill=(120, 150, 200))
     d.text((VW - 20, 20), opp_name.upper(), font=F_MED, fill=(190, 130, 80), anchor="ra")
-
-    speed = math.hypot(env.ball_vx, env.ball_vy)
-    d.text((20, VH - 26),
-           f"rally {env.rally_count} hits   ball {speed:4.1f} px/f", font=F_SMALL,
-           fill=(120, 135, 160))
+    d.text((20, VH - 26), f"rally {rally} hits   ball {bspeed:4.1f} px/f",
+           font=F_SMALL, fill=(120, 135, 160))
+    d.text((VW - 20, VH - 26), speed_label, font=F_SMALL, fill=(120, 135, 160),
+           anchor="ra")
 
     if banner:
         d.rectangle([0, VH // 2 - 34, VW, VH // 2 + 34], fill=(11, 15, 25))
-        d.text((VW // 2, VH // 2), banner, font=F_BIG, fill=(249, 115, 22), anchor="mm")
+        d.text((VW // 2, VH // 2), banner, font=F_BIG, fill=ORANGE, anchor="mm")
     return np.asarray(img)
 
 
-def _play_point(env, opponent, frames, opp_name, score_a, score_b, serve_dir, record):
-    """Run one point sub-step by sub-step; returns the winner string."""
-    env.reset(serve_direction=serve_dir)
-    trail = []
-    ai_act = 0
-    opp_act = 0
-    for i in range(POINT_SUBSTEP_CAP):
-        # Both policies are re-queried every `frame_skip` sub-steps, exactly as
-        # StandalonePongEnv.step() does in the author's simulate() loop.
-        if i % env.phys.frame_skip == 0:
-            ai_act = policy_action(env.get_ego_observation())
-            opp_act = opponent.act(env)
-        trail.append((env.ball_x, env.ball_y))
-        if len(trail) > 8:
-            trail.pop(0)
-        done, info = env.physics_substep(ego_action=ai_act, opp_action=opp_act)
-        if record and len(frames) < MAX_FRAMES:
-            frames.append(_render_frame(env, score_a, score_b, opp_name, trail))
-        if done:
-            return info.get("winner")
-    return "draw"
+def _snapshot(env, sa, sb):
+    return (env.ball_x, env.ball_y, env.ego_y, env.opp_y,
+            env.ego_paddle_h, env.opp_paddle_h, env.rally_count, sa, sb,
+            math.hypot(env.ball_vx, env.ball_vy))
+
+
+def _simulate_match(env, opponent, points_to_win, budget):
+    """Play a match sub-step by sub-step, recording every state for later replay."""
+    states, rallies = [], []
+    score_a = score_b = 0
+    serve_dir = 1
+
+    while (score_a < points_to_win and score_b < points_to_win
+           and len(rallies) < MAX_POINTS and len(states) < budget):
+        env.reset(serve_direction=serve_dir)
+        ai_act = opp_act = 0
+        winner = "draw"
+        for i in range(POINT_SUBSTEP_CAP):
+            # Both sides re-decide every `frame_skip` sub-steps, exactly as
+            # StandalonePongEnv.step() does in the author's simulate() loop.
+            if i % env.phys.frame_skip == 0:
+                ai_act = policy_action(env.get_ego_observation())
+                opp_act = opponent.act(env)
+            done, info = env.physics_substep(ego_action=ai_act, opp_action=opp_act)
+            states.append(_snapshot(env, score_a, score_b))
+            if done:
+                winner = info.get("winner")
+                break
+            if len(states) >= budget:
+                break
+
+        rallies.append(env.rally_count)
+        if winner == "ego":
+            score_a += 1
+            serve_dir = 1
+        elif winner == "opponent":
+            score_b += 1
+            serve_dir = -1
+        else:
+            serve_dir = -serve_dir
+        # keep the final state of the point on screen for a beat
+        states.extend([_snapshot(env, score_a, score_b)] * env.phys.frame_skip * 8)
+
+    return states, rallies, score_a, score_b
 
 
-def watch_match(opponent_name: str, points_to_win: int = 3,
+def watch_match(opponent_name: str, points_to_win: int = 2,
                 randomize_seed: bool = True, seed: int = 0):
-    """Simulate and render a Mr. Pong match against a baseline opponent.
+    """Simulate and render a Mr. Pong match against a scripted baseline opponent.
 
     Args:
-        opponent_name: which scripted baseline to face (from the model card).
+        opponent_name: which baseline from the model card to face.
         points_to_win: number of points needed to take the match.
         randomize_seed: draw a fresh random seed for this match.
         seed: RNG seed used when randomize_seed is off.
 
     Returns:
-        An MP4 of the match, a markdown scoreline, and the seed actually used.
+        An MP4 replay of the match, a markdown scoreline, and the seed used.
     """
     import imageio.v2 as imageio
 
@@ -187,48 +210,44 @@ def watch_match(opponent_name: str, points_to_win: int = 3,
     random.seed(seed)
     np.random.seed(seed % (2**32))
 
-    points_to_win = int(points_to_win)
-    opp_cls = OPPONENTS[opponent_name]
-    opponent = opp_cls(rng=rng)
-
+    points_to_win = max(1, int(points_to_win))
+    opponent = OPPONENTS[opponent_name](rng=rng)
     env = StandalonePongEnv(PhysicsConfig(), seed=seed)
-    frames = []
-    score_a = score_b = 0
-    serve_dir = 1
-    rallies = []
-    t0 = time.perf_counter()
 
-    while (score_a < points_to_win and score_b < points_to_win
-           and len(rallies) < MAX_POINTS):
-        winner = _play_point(env, opponent, frames, opponent_name,
-                             score_a, score_b, serve_dir,
-                             record=len(frames) < MAX_FRAMES)
-        rallies.append(env.rally_count)
-        if winner == "ego":
-            score_a += 1
-            serve_dir = 1
-        elif winner == "opponent":
-            score_b += 1
-            serve_dir = -1
-        else:
-            serve_dir = -serve_dir
-        if len(frames) >= MAX_FRAMES:
-            break
+    t0 = time.perf_counter()
+    states, rallies, score_a, score_b = _simulate_match(
+        env, opponent, points_to_win, MATCH_SUBSTEP_CAP)
+    sim_s = time.perf_counter() - t0
 
-    banner = ("MR. PONG WINS %d-%d" % (score_a, score_b) if score_a > score_b
-              else ("%s WINS %d-%d" % (opponent_name.upper(), score_b, score_a)
-                    if score_b > score_a else "TIME LIMIT"))
+    # Physics runs at 40 sub-steps/s; drop every n-th state so the clip fits the
+    # frame budget, and say so on screen.
+    stride = max(1, math.ceil(len(states) / MAX_FRAMES))
+    speed_label = "real time" if stride == 1 else f"{stride}x speed"
+
+    if score_a > score_b:
+        banner = f"MR. PONG WINS {score_a}-{score_b}"
+    elif score_b > score_a:
+        banner = f"{opponent_name.upper()} WINS {score_b}-{score_a}"
+    else:
+        banner = f"{score_a}-{score_b} — TIME LIMIT"
+
+    t1 = time.perf_counter()
+    frames, trail = [], []
+    for s in states[::stride]:
+        trail.append((s[0], s[1]))
+        if len(trail) > 8:
+            trail.pop(0)
+        frames.append(_render_frame(s, opponent_name, trail, speed_label))
+    last = states[-1]
     for _ in range(FPS * 2):
-        if len(frames) < MAX_FRAMES + FPS * 2:
-            frames.append(_render_frame(env, score_a, score_b, opponent_name,
-                                        [], banner=banner))
+        frames.append(_render_frame(last, opponent_name, [], speed_label, banner=banner))
 
-    sim_s = time.perf_counter() - t0
     out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     with imageio.get_writer(out, fps=FPS, codec="libx264", quality=7,
                             macro_block_size=1, pixelformat="yuv420p") as w:
         for f in frames:
             w.append_data(f)
+    render_s = time.perf_counter() - t1
 
     stats = (
         f"### Mr. Pong **{score_a} – {score_b}** {opponent_name}\n"
@@ -236,8 +255,9 @@ def watch_match(opponent_name: str, points_to_win: int = 3,
         f"| Points played | {len(rallies)} |\n"
         f"| Longest rally | {max(rallies)} hits |\n"
         f"| Average rally | {np.mean(rallies):.1f} hits |\n"
-        f"| Video length | {len(frames) / FPS:.1f} s ({len(frames)} frames) |\n"
-        f"| Simulation time | {sim_s:.2f} s |\n"
+        f"| Game time simulated | {len(states) / FPS:.1f} s |\n"
+        f"| Playback | {speed_label} ({len(frames) / FPS:.1f} s clip) |\n"
+        f"| Simulate / render | {sim_s:.2f} s / {render_s:.2f} s |\n"
         f"| Seed | `{seed}` |\n"
     )
     return out, stats, seed
@@ -317,7 +337,7 @@ It wins ~87% of its matches against the strongest scripted baseline it was train
 in real time. Or watch the real PyTorch checkpoint take on the model-card opponents.
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Mr. Pong") as demo:
+with gr.Blocks(title="Mr. Pong") as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(INTRO)
 
@@ -341,7 +361,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Mr. Pong") as demo:
                                      height=430, format="mp4")
                 stats_out = gr.Markdown()
                 with gr.Accordion("Advanced settings", open=False):
-                    points = gr.Slider(1, 5, value=3, step=1, label="Points to win")
+                    points = gr.Slider(1, 5, value=2, step=1, label="Points to win")
                     rand_seed = gr.Checkbox(value=True, label="Randomize seed")
                     seed_num = gr.Number(value=0, precision=0, label="Seed")
                 gr.Examples(
@@ -391,4 +411,4 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Mr. Pong") as demo:
         )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)