Blame for scripts/beam_play.py

"""Beam plays SUPERBEAM: the hosted Beam model picks every action from an ASCII frame.

The game runs locally in the reference CHIP-8 emulator (beamboot/superbeam_gym.py).
A small Modal function proxies chat requests to the HF router; the HF token stays
in the Modal secret `huggingface-secret` and is never printed.

  OMP_NUM_THREADS=1 .venv/bin/modal run scripts/beam_play.py               # real run
  OMP_NUM_THREADS=1 .venv/bin/modal run scripts/beam_play.py --dry-run     # scripted stand-in, no HF calls

Budget: <= --max-requests chat requests per attempt (default 80) and
<= --total-cap overall (default 200, i.e. up to 3 attempts).  Every HTTP attempt,
including retries, counts toward the cap.

Writes results/beam_play_run.json (every request/response, actions, rewards,
frames) and results/beam_key_schedule.json (best attempt's key schedule, for
replay on the in-weights backend).
"""
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request

import modal

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RUN_OUT = os.path.join(ROOT, "results", "beam_play_run.json")
SCHED_OUT = os.path.join(ROOT, "results", "beam_key_schedule.json")
MODEL = "reflection/Beam-1.0-501B-A23B:baseten"

app = modal.App("beamboot-beam-play")


@app.function(image=modal.Image.debian_slim(python_version="3.12"),
              secrets=[modal.Secret.from_name("huggingface-secret")], timeout=240)
def beam_chat(messages, tools, model=MODEL, max_tokens=6144):
    token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
             or os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN"))
    if not token:
        raise RuntimeError("Hugging Face token was not found under a supported environment "
                           "variable name in the Modal secret")
    payload = {"model": model, "messages": messages, "max_tokens": max_tokens}
    if tools:
        payload.update(tools=tools, tool_choice="auto")
    req = urllib.request.Request("https://router.huggingface.co/v1/chat/completions",
                                 data=json.dumps(payload).encode(),
                                 headers={"Authorization": "Bearer " + token,
                                          "Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=220) as response:
            out = json.load(response)
    except urllib.error.HTTPError as error:
        raise RuntimeError(f"Hugging Face router returned HTTP {error.code}: "
                           f"{error.read(1000).decode(errors='replace')}") from None
    out["_latency_s"] = round(time.time() - t0, 2)
    return out


# ----------------------------------------------------------------- prompt
ACTION_NAMES = ["noop", "left", "right", "jump", "right+jump", "left+jump"]
TOOL = {"type": "function", "function": {
    "name": "act",
    "description": "Press the controller for the next 4 frames.",
    "parameters": {"type": "object",
                   "properties": {"action": {"type": "string", "enum": ACTION_NAMES}},
                   "required": ["action"], "additionalProperties": False}}}

SYSTEM = """You are playing SUPERBEAM, an original side-scrolling platformer running on a CHIP-8 computer (64x32 pixel screen).
Each turn you get the screen as ASCII art plus a few facts, then you choose ONE action by calling the `act` tool.
Legend: P = you, E = walker enemy, # = ground, H = pipe (solid block), o = coin, F = flag pole, C = castle, . = empty.
Rules:
- Goal: reach the flag pole far to the right. The camera scrolls as you move right.
- Falling into a gap in the ground (columns where the ground rows are '.') kills you.
- Pipes are solid: you cannot walk through them, you can stand on top. A pipe is 8 px tall and can be
  jumped onto from the ground. A tall pipe is 16 px tall: climb it from a pipe next to it. You cannot
  jump while standing on a tall pipe; walk off it.
- Coins: touching one collects it (+points). Low coins: walk into them. Mid coins: jump from the ground
  or walk on a pipe. High coins: only from a pipe top (jump) or by standing on a tall pipe.
- Walkers walk and turn around at gaps and pipes. Touching one from the side kills you; landing on top stomps it.
- One action = 4 frames. Running moves you 8 px per action. A jump rises 10 px, lasts about 9 frames
  (~2 actions) and can only start while standing (on ground or a pipe).
- There is NO momentum: you only move sideways on turns where you choose right/left, in the air too.
  While airborne, choosing "noop" makes you fall straight down (often into the gap). To clear a gap:
  "right+jump" once, then keep choosing "right" every turn until you land (~18 px forward in total).
  To get onto a pipe: walk up to it, "right+jump", then keep choosing "right".
- Good timing: jump over a gap on the turn when it is 1 px ahead of you or less (it may already be
  "under you" while you still stand: you only fall once your back edge leaves the ground); jumping
  earlier lands in a 16 px gap. Jump when a walker is about 4-20 px ahead (landing on it stomps it).
  Otherwise keep running right.
Actions: noop, left, right, jump, right+jump, left+jump.
If you cannot call the tool, reply with exactly one action word and nothing else."""


def describe(env, obs):
    """Short factual legend derived from the game's registers/RAM (SuperBeamEnv.describe): your
    position and footing, gaps, solid blocks (pipes) and coins on screen, walkers, coins, progress."""
    return env.describe(obs)


def view(env):
    """Rows 0..25 of the labelled ASCII frame (rows 26-29 repeat the ground, 30-31 are the coin bar)."""
    rows = env.ascii(legend=True).splitlines()
    return "\n".join(rows[0:26])


def parse_action(message):
    for call in message.get("tool_calls") or []:
        try:
            if call["function"]["name"] == "act":
                a = json.loads(call["function"]["arguments"]).get("action")
                if a in ACTION_NAMES:
                    return a, "tool"
        except (KeyError, ValueError, TypeError, AttributeError):
            pass
    text = (message.get("content") or "").strip().lower()
    words = re.findall(r"(?:left|right)\+jump|noop|left|right|jump", text)
    if words:
        return words[-1], "text"
    return None, "invalid"


# ----------------------------------------------------------------- play loop
def play_attempt(env, chat, max_requests, history_turns=2, log=print):
    obs, info = env.reset()
    obs, _, _, _, info = env.step(0)            # one noop so the first frame is drawn
    history, history_turns_list, records, used = [], [], [], 0
    while used < max_requests:
        facts = describe(env, obs)
        user = {"role": "user", "content": "Screen (rows 0-25, 64 columns):\n%s\n%s\nChoose your action."
                % (view(env), facts)}
        messages = [{"role": "system", "content": SYSTEM}] + history + [user]
        response, error = None, None
        for _try in range(3):
            if used >= max_requests:
                break
            used += 1
            try:
                response = chat(messages, [TOOL])
                break
            except Exception as e:  # provider/network error: record, maybe retry
                error = str(e)[:2000]
                log("request %d failed: %s" % (used, error[:300]))
                if "HTTP 4" in error and "HTTP 429" not in error:
                    break
                time.sleep(3)
        if response is None:
            records.append({"request": used, "messages": messages, "error": error})
            if error and "HTTP 4" in error and "HTTP 429" not in error:
                return records, used, "provider_error: " + error[:500]
            continue
        msg = response["choices"][0]["message"]
        action, how = parse_action(msg)
        a = ACTION_NAMES.index(action) if action else 0
        frame_before = view(env)
        obs, reward, term, trunc, info = env.step(a)
        rec = {"request": used, "messages": messages, "response": response,
               "action": action or "noop", "parsed_from": how, "reward": reward,
               "frame": frame_before, "progress": info["best_progress"],
               "frames": info["frames"], "terminated": term, "truncated": trunc}
        records.append(rec)
        log("req %2d  %-10s (%s)  progress %5.1f%%  frames %3d  reward %+6.2f%s"
            % (used, action or "noop", how, 100 * info["best_progress"], info["frames"], reward,
               "  WIN" if info["won"] else ("  DIED" if info["dead"] else "")))
        # rolling history: assistant reply (with reasoning_content if returned) + result
        assistant = {k: msg[k] for k in ("role", "content", "reasoning_content", "tool_calls") if k in msg}
        assistant.setdefault("role", "assistant")
        result = "done; progress now %.0f%%" % (100 * info["best_progress"])
        # older turns keep only the facts line (no ASCII) to keep requests short
        turn = [{"role": "user", "content": "(earlier turn) " + facts}, assistant]
        calls = msg.get("tool_calls") or []
        for call in calls:
            turn.append({"role": "tool", "tool_call_id": call.get("id", ""), "content": result})
        if not calls:
            turn[1] = {"role": "assistant", "content": action or "noop"}
            if msg.get("reasoning_content"):   # provider rejects assistant turns without their reasoning
                turn[1]["reasoning_content"] = msg["reasoning_content"]
        history_turns_list = (history_turns_list + [turn])[-history_turns:] if history_turns else []
        history = [m for t in history_turns_list for m in t]
        if term or trunc:
            break
    status = "won" if env.won else ("died" if env.dead else
                                    ("truncated" if env.frames >= env.max_frames else "request_cap"))
    return records, used, status


def run(chat, max_requests=80, total_cap=200, frame_skip=4, max_attempts=3, model=MODEL, log=print,
        save=True, backend="emulator", level="classic"):
    global RUN_OUT, SCHED_OUT
    level = int(level) if str(level).isdigit() else level
    tag = "" if level == "classic" else "_level%s" % level
    if backend != "emulator":          # live run against the in-weights computer: separate files
        RUN_OUT = os.path.join(ROOT, "results", "beam_live_weights%s_run.json" % tag)
        SCHED_OUT = os.path.join(ROOT, "results", "beam_live_weights%s_schedule.json" % tag)
    elif tag:                          # held-out / generated level evaluation
        RUN_OUT = os.path.join(ROOT, "results", "levels", "beam%s_run.json" % tag)
        SCHED_OUT = os.path.join(ROOT, "results", "levels", "beam%s_schedule.json" % tag)
    sys.path.insert(0, ROOT)
    from beamboot.superbeam_gym import SuperBeamEnv, WIN_X, START_X, IPF
    est = -(-(WIN_X - START_X) // (2 * frame_skip))
    log("estimate: level needs >= %d frames = >= %d actions at frame_skip=%d (cap %d requests/attempt, %d total)"
        % (-(-(WIN_X - START_X) // 2), est, frame_skip, max_requests, total_cap))
    if est > max_requests:
        raise SystemExit("frame_skip too small for the request cap")
    attempts, total = [], 0
    for k in range(max_attempts):
        budget = min(max_requests, total_cap - total)
        if budget <= 0:
            break
        env = SuperBeamEnv(backend=backend, frame_skip=frame_skip, max_frames=600, level=level,
                           device="cuda" if backend != "emulator" else "cpu")
        log("=== attempt %d (budget %d requests, backend %s) ===" % (k + 1, budget, backend))
        t0 = time.time()
        records, used, status = play_attempt(env, chat, budget, log=log)
        total += used
        attempts.append({"attempt": k + 1, "status": status, "requests_used": used,
                         "won": env.won, "dead": env.dead, "progress": env._info()["best_progress"],
                         "frames": env.frames, "actions": [ACTION_NAMES[a] for a in env.actions],
                         "key_schedule": env.key_schedule(), "requests": records,
                         "backend": backend, "wall_seconds": round(time.time() - t0, 1)})
        if backend != "emulator":
            st = env.backend.session.stats
            attempts[-1]["weights"] = {
                "frame_rows": env.backend.session.frame_rows,
                "per_frame": [{"f": x.frame, "tokens": x.frame_tokens, "accepted_all": x.accepted_all,
                               "first_reject": x.first_reject, "s": round(x.seconds, 3)} for x in st],
                "frames_all_draft_accepted": sum(x.accepted_all for x in st),
                "compute_seconds": round(sum(x.seconds for x in st), 1)}
        log("attempt %d: %s, progress %.1f%%, %d requests" % (k + 1, status,
                                                               100 * attempts[-1]["progress"], used))
        if env.won or status.startswith("provider_error"):
            break
    best = max(attempts, key=lambda a: (a["won"], a["progress"]))
    # verify the best schedule replays to the same outcome
    venv = SuperBeamEnv(frame_skip=1, max_frames=600, level=level)
    vinfo = venv.replay_schedule(best["key_schedule"], n_frames=best["frames"])
    out = {"model": model, "level": level, "frame_skip": frame_skip, "ipf": IPF,
           "max_requests_per_attempt": max_requests,
           "total_cap": total_cap, "total_requests": total, "best_attempt": best["attempt"],
           "best": {k: best[k] for k in ("status", "won", "progress", "frames", "requests_used")},
           "replay_check": vinfo, "attempts": attempts}
    log("best: attempt %d %s progress %.1f%% | total requests %d | replay %s"
        % (best["attempt"], best["status"], 100 * best["progress"], total, vinfo))
    if not save:
        return out
    os.makedirs(os.path.dirname(RUN_OUT), exist_ok=True)
    with open(RUN_OUT, "w") as f:
        json.dump(out, f, indent=1, default=str)
    with open(SCHED_OUT, "w") as f:
        json.dump({"rom": "roms/superbeam.ch8", "ipf": IPF, "frame_skip": frame_skip,
                   "frames": best["frames"], "won": best["won"], "progress": best["progress"],
                   "format": "[[frame, key, 'down'|'up'], ...] applied before that frame's instructions",
                   "events": best["key_schedule"]}, f, indent=1)
    log("wrote %s and %s" % (os.path.relpath(RUN_OUT, ROOT), os.path.relpath(SCHED_OUT, ROOT)))
    return out


def _dry_chat(messages, tools):
    """Offline stand-in: reads the facts line like the prompt suggests (no HF call)."""
    text = messages[-1]["content"]
    m = re.search(r"standing on (ground|a pipe)", text)
    terrain, walkers = text.split("Walkers:")[0], text.split("Walkers:")[-1].split("Coins collected")[0]
    near = [int(v) for v in re.findall(r"\((-?\d+) px ahead\)", terrain)]
    foes = [int(v) for v in re.findall(r"\((-?\d+) px ahead\)", walkers)]
    act = "right"
    if m and (any(d <= 1 for d in near) or "under you" in terrain.split("Solid")[0]
              or any(4 <= d <= 20 for d in foes)):
        act = "right+jump"
    return {"choices": [{"message": {"role": "assistant", "content": None, "tool_calls": [
        {"id": "dry", "type": "function", "function": {"name": "act",
                                                       "arguments": json.dumps({"action": act})}}]}}]}


@app.local_entrypoint()
def main(dry_run: bool = False, max_requests: int = 80, total_cap: int = 200,
         frame_skip: int = 4, model: str = MODEL, backend: str = "emulator", max_attempts: int = 3,
         level: str = "classic"):
    if dry_run:
        run(_dry_chat, max_requests, total_cap, frame_skip, max_attempts, model="dry-run", save=False,
            backend=backend, level=level)
        return
    run(lambda msgs, tools: beam_chat.remote(msgs, tools, model), max_requests, total_cap,
        frame_skip, max_attempts, model=model, backend=backend, level=level)