app.py

"""
🏓 Mr. Pong — an interactive demo for fromziro/MrPong.

Two ways to meet the agent:
  1. Play against it live in the browser. The 28K-parameter policy is exported
     from the checkpoint and evaluated client-side at 40 physics steps/second,
     so the paddle reacts with zero network latency.
  2. Watch / benchmark the real PyTorch checkpoint playing full matches against
     the baseline opponents from the model card, rendered to video on the server.
"""

import base64
import html as html_lib
import math
import os
import random
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoConfig, AutoModel

from pong_engine import (
    OPPONENTS,
    PhysicsConfig,
    StandalonePongEnv,
)

MODEL_ID = "fromziro/MrPong"
HERE = Path(__file__).parent

# ----------------------------------------------------------------------------
# Model (28,484 params — pure CPU, a forward pass costs microseconds)
# ----------------------------------------------------------------------------
config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()
OBS_DIM = int(getattr(config, "obs_dim", 12))
N_PARAMS = sum(p.numel() for p in model.parameters())
print(f"[OK] {MODEL_ID} loaded — {N_PARAMS:,} params, obs_dim={OBS_DIM}, "
      f"hidden_dims={config.hidden_dims}")


@torch.no_grad()
def policy_action(obs: np.ndarray, deterministic: bool = True) -> int:
    """Greedy action from the Mr. Pong policy (0 stay, 1 up, 2 down)."""
    x = obs[:OBS_DIM] if len(obs) >= OBS_DIM else np.pad(obs, (0, OBS_DIM - len(obs)))
    return int(model.act(x, deterministic=deterministic))


# ----------------------------------------------------------------------------
# Export the checkpoint for the in-browser copy of the policy
# ----------------------------------------------------------------------------
def _export_weights_b64() -> str:
    sd = model.state_dict()
    order = ["trunk.0.weight", "trunk.0.bias", "trunk.2.weight", "trunk.2.bias",
             "actor.weight", "actor.bias", "critic.weight", "critic.bias"]
    flat = np.concatenate([sd[k].detach().cpu().numpy().astype("<f4").ravel()
                           for k in order])
    assert flat.size == N_PARAMS, (flat.size, N_PARAMS)
    return base64.b64encode(flat.tobytes()).decode("ascii")


GAME_HTML = (HERE / "game_template.html").read_text().replace(
    "__WEIGHTS_B64__", _export_weights_b64()
)
GAME_IFRAME = (
    '<iframe title="Play Mr. Pong" srcdoc="{}" '
    'style="width:100%;height:700px;border:0;border-radius:12px;overflow:hidden" '
    'scrolling="no"></iframe>'
).format(html_lib.escape(GAME_HTML, quote=True))


# ----------------------------------------------------------------------------
# Server-side match rendering (real PyTorch checkpoint)
# ----------------------------------------------------------------------------
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 = 1400             # hard cap on rendered frames (35 s at 40 fps)
POINT_SUBSTEP_CAP = 4500      # == the author's 1500-step rally limit
MATCH_SUBSTEP_CAP = 12000     # enough for a first-to-2 to finish vs any beatable baseline
MAX_POINTS = 12               # safety net against endless draw sequences


def _font(size: int):
    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if os.path.exists(p):
            return ImageFont.truetype(p, size)
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


F_BIG, F_MED, F_SMALL = _font(34), _font(15), _font(13)
BLUE, ORANGE = (79, 140, 255), (249, 115, 22)


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)

    for y in range(0, VH, 26):
        d.line([(VW // 2, y), (VW // 2, min(VH, y + 12))], fill=(38, 48, 70), width=2)
    d.rectangle([0, 0, VW - 1, VH - 1], outline=(28, 36, 54), width=2)

    for i, (tx, ty) in enumerate(trail):
        a = (i + 1) / (len(trail) + 1)
        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, 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([bx * SC - r, by * SC - r, bx * SC + r, by * SC + r], fill=(255, 255, 255))

    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")
    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=ORANGE, anchor="mm")
    return np.asarray(img)


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 = 2,
                randomize_seed: bool = True, seed: int = 0):
    """Simulate and render a Mr. Pong match against a scripted baseline opponent.

    Args:
        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 replay of the match, a markdown scoreline, and the seed used.
    """
    import imageio.v2 as imageio

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)
    rng = random.Random(seed)
    random.seed(seed)
    np.random.seed(seed % (2**32))

    points_to_win = max(1, int(points_to_win))
    opponent = OPPONENTS[opponent_name](rng=rng)
    env = StandalonePongEnv(PhysicsConfig(), seed=seed)

    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

    # 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"

    decided = max(score_a, score_b) >= points_to_win
    if decided and score_a > score_b:
        banner, result = f"MR. PONG WINS {score_a}-{score_b}", "Mr. Pong takes the match"
    elif decided:
        banner = f"{opponent_name.upper()} WINS {score_b}-{score_a}"
        result = f"{opponent_name} takes the match"
    else:
        banner = f"{score_a}-{score_b} — TIME LIMIT"
        result = "stopped at the clip time limit (no side reached the target)"

    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):
        frames.append(_render_frame(last, opponent_name, [], speed_label, banner=banner))

    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"
        f"| | |\n|---|---|\n"
        f"| Result | {result} |\n"
        f"| Points played | {len(rallies)} |\n"
        f"| Longest rally | {max(rallies)} hits |\n"
        f"| Average rally | {np.mean(rallies):.1f} hits |\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


def benchmark(opponent_name: str, num_matches: int = 25, seed: int = 0):
    """Replay the model card's evaluation: N single-point matches, no rendering.

    Args:
        opponent_name: which scripted baseline to face.
        num_matches: how many matches to simulate.
        seed: RNG seed for reproducibility.

    Returns:
        A markdown summary of the win / draw / loss record.
    """
    seed = int(seed)
    rng = random.Random(seed)
    random.seed(seed)
    num_matches = int(num_matches)

    opponent = OPPONENTS[opponent_name](rng=rng)
    env = StandalonePongEnv(PhysicsConfig(), seed=seed)

    wins = draws = losses = 0
    rallies = []
    t0 = time.perf_counter()

    for i in range(1, num_matches + 1):
        obs = env.reset(serve_direction=1 if i % 2 == 0 else -1)
        done = False
        info = {}
        while not done:
            ego_act = policy_action(obs)
            opp_act = opponent.act(env)
            obs, done, info = env.step(ego_action=ego_act, opp_action=opp_act)
        rallies.append(env.rally_count)
        w = info.get("winner")
        if w == "ego":
            wins += 1
        elif w == "opponent":
            losses += 1
        else:
            draws += 1

    el = time.perf_counter() - t0
    n = max(1, num_matches)
    return (
        f"### Mr. Pong vs **{opponent_name}** — {num_matches} matches\n\n"
        f"| Metric | Value |\n|---|---|\n"
        f"| Win rate | **{wins / n * 100:.1f}%** |\n"
        f"| Draw rate | {draws / n * 100:.1f}% |\n"
        f"| Loss rate | {losses / n * 100:.1f}% |\n"
        f"| Record (W/D/L) | {wins}W / {draws}D / {losses}L |\n"
        f"| Average rally | {np.mean(rallies):.1f} hits (max {max(rallies)}) |\n"
        f"| Wall time | {el:.2f} s ({num_matches / max(el, 1e-6):.1f} matches/s) |\n\n"
        f"*A draw means the rally hit the {env.phys.max_rally_steps}-step limit "
        f"without either side scoring.*"
    )


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = f"""# 🏓 Mr. Pong

[**fromziro/MrPong**](https://huggingface.co/{MODEL_ID}) is a {N_PARAMS:,}-parameter
actor–critic MLP trained with PPO and self-play for 10M steps to play 2D table tennis.
It wins ~87% of its matches against the strongest scripted baseline it was trained on.

**Play it yourself below** — the policy runs in your browser, so the paddle answers
in real time. Or watch the real PyTorch checkpoint take on the model-card opponents.
"""

with gr.Blocks(title="Mr. Pong") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)

        with gr.Tabs():
            with gr.Tab("🎮 Play Mr. Pong"):
                gr.HTML(GAME_IFRAME)
                gr.Markdown(
                    "Click the court to start. Hold **W** / **↑** and **S** / **↓** "
                    "(or switch to mouse / touch control). First to 5 points.  \n"
                    "*V(s) is the critic's value estimate from Mr. Pong's point of "
                    "view — it climbs when the agent thinks it is about to win the point.*"
                )

            with gr.Tab("🎬 Watch a match"):
                with gr.Row():
                    opponent = gr.Dropdown(
                        list(OPPONENTS.keys()), value="Realistic Hard",
                        label="Opponent", scale=3)
                    watch_btn = gr.Button("Play match", variant="primary", scale=1)
                video_out = gr.Video(label="Match replay", autoplay=True,
                                     height=430, format="mp4")
                stats_out = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    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(
                    examples=[["Realistic Hard"], ["Impossible Hard"],
                              ["Medium Logic"], ["Random Agent"]],
                    inputs=[opponent],
                    outputs=[video_out, stats_out, seed_num],
                    fn=watch_match,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Opponents from the model card",
                )
                watch_btn.click(
                    watch_match,
                    inputs=[opponent, points, rand_seed, seed_num],
                    outputs=[video_out, stats_out, seed_num],
                    api_name="watch_match",
                )

            with gr.Tab("📊 Benchmark"):
                with gr.Row():
                    bench_opp = gr.Dropdown(
                        list(OPPONENTS.keys()), value="Realistic Hard",
                        label="Opponent", scale=3)
                    bench_btn = gr.Button("Run benchmark", variant="primary", scale=1)
                bench_out = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    n_matches = gr.Slider(5, 200, value=25, step=5, label="Matches")
                    bench_seed = gr.Number(value=0, precision=0, label="Seed")
                gr.Markdown(
                    "Reference numbers reported by the author over 1000 matches:  \n"
                    "Easy 99.7% · Medium 99.0% · Realistic Hard 86.6% · "
                    "Random 100% · Impossible Hard 0% W / 98.4% D."
                )
                bench_btn.click(
                    benchmark,
                    inputs=[bench_opp, n_matches, bench_seed],
                    outputs=[bench_out],
                    api_name="benchmark",
                )

        gr.Markdown(
            "Model, physics environment and baseline opponents by "
            "[FromZero](https://huggingface.co/fromziro) (Apache-2.0). "
            "The in-browser policy is the exact checkpoint re-evaluated in JavaScript; "
            "the *Watch* and *Benchmark* tabs run the PyTorch model itself."
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)