scripts/play.py

#!/usr/bin/env python3
"""Closed-loop gameplay evaluation.

Controllers are kept deliberately separable so it is clear what is doing the
work.  ``model`` asks the network for the action distribution and does nothing
else -- no search, no visited set, no memory.  ``model+memory`` adds the thin
bookkeeping a deployed agent would have, which is roughly what NanoJev's
showcase controller provides, and is reported separately rather than folded
into the headline number.

``random+memory`` is that second row's control, and the reason it exists: the
memory is doing enough of the work that ``model+memory`` says very little on
its own.  Running the identical bookkeeping with no network behind it says how
much.
"""
from __future__ import annotations

import argparse
import json
import math
import random
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch

from envs import maze as mz
from envs import snake as sk
from jevon.inference import JevonRunner, maze_state_sample, snake_state_sample


# ---------------------------------------------------------------- policies

def model_action(runner, sample, legal, rng, temperature):
    answers = {a["question"]: a for a in runner.answer([sample])}
    row = answers.get("action")
    if row is None:
        # question_profile omits the action question when fewer than two moves
        # are legal, because there is nothing to choose between. The answer is
        # then forced, and taking it is not a guess -- but taking legal[0] is.
        # That fallback used to fire at every one-exit cell and return `north`,
        # a wall, so the agent never moved and was handed the identical state
        # forever. Every maze figure for this controller measured that, not the
        # network.
        return (legal[0] if len(legal) == 1 else rng.choice(legal)), {}
    distribution = {k: v for k, v in row["probabilities"].items() if k in legal}
    if not distribution:
        return rng.choice(legal), row["probabilities"]
    if temperature <= 0:
        return max(distribution, key=distribution.get), distribution
    keys = list(distribution)
    weights = [max(distribution[k], 1e-9) ** (1.0 / temperature) for k in keys]
    return rng.choices(keys, weights=weights, k=1)[0], distribution


def field_action(runner, sample, state, legal, rng):
    """Move to the neighbouring cell the planner thinks is closest to the goal.

    Pure greedy descent on the model's own learned field -- no search, no
    memory, no visited set. It reads the planner's training read-out instead of
    the text-conditioned action head, so it measures whether the recurrence
    learned a usable value function rather than whether the language side
    learned to describe one.
    """
    field = runner.distance_field(sample)
    scores = {}
    for action in legal:
        dr, dc = mz.DIRECTIONS[action]
        cell = (state.position[0] + dr, state.position[1] + dc)
        if not state.free(cell):
            continue
        scores[action] = float(field[cell])
    if not scores:
        return rng.choice(legal), {}
    best = min(scores.values())
    # Report it as a distribution so traces render the same way as the others.
    weights = {a: math.exp(-(v - best) * 8.0) for a, v in scores.items()}
    total = sum(weights.values())
    return min(scores, key=scores.get), {a: w / total for a, w in weights.items()}


def episode_field(runner, state):
    """The model's predicted distance-to-goal for every cell, once per episode.

    The planner never sees the agent, so this is constant for the whole
    episode -- the same property that lets the controller cache it.  Recorded
    so the viewer can show what the model believes about the board, not just
    which move it picked.
    """
    if runner is None:
        return None
    try:
        field = runner.distance_field(maze_state_sample(state))
    except RuntimeError:                       # checkpoint trained without a planner
        return None
    return [[round(float(v), 4) for v in row] for row in field]


class Recorder(list):
    """A step log with a ceiling, so a trace stays viewable without lying.

    A random walk on a 31x31 maze runs to its 3844-step cap, and thirty-six of
    those come to roughly six megabytes of JSON that nobody will ever scrub
    through. Truncating the *recording* keeps the episode itself untouched, so
    the summary -- solve rate, step count, collisions -- is still the real one;
    only the replay is shortened, and ``truncated`` says by how much.
    """

    def __init__(self, cap: int = 0):
        super().__init__()
        self.cap = cap
        self.seen = 0

    def append(self, item) -> None:
        self.seen += 1
        if not self.cap or len(self) < self.cap:
            super().append(item)

    @property
    def truncated(self) -> int:
        return self.seen - len(self)


def maze_steps(runner, state, controller, max_steps, rng, temperature=0.0):
    """One decision per ``next()``, yielding the step it just recorded.

    Written as a generator so the live server in the arcade repo can take a
    single forward pass per request while ``play_maze`` below drains it in a
    loop. The alternative was a second copy of this decision logic behind an
    HTTP handler, which is the failure mode these two repos have hit more than
    any other -- and this is the loop every shipped number is measured from,
    so a divergent copy would be a viewer demonstrating a model nobody
    benchmarked. Returns the episode result, so callers need ``StopIteration``
    or ``yield from``.
    """
    # The set of cells stood on, read twice. It is the ``model+memory``
    # controller's whole memory -- prefer a move that leaves ground already
    # covered -- and it is also the wandering diagnostic, because an agent that
    # visited one cell and an agent that toured forty both score 0.00.
    #
    # Memory used to mean remembering which (cell, action) pairs had bounced
    # off a wall. Restricting the candidate set to legal moves made that set
    # permanently empty, which would have left ``model+memory`` a silent
    # duplicate of ``model`` and its row in the results table a lie. Cycles are
    # what end these episodes now: the planner is stateless, so a greedy policy
    # that steps A->B->A repeats it for the whole budget.
    seen: set[tuple[int, int]] = {state.position}
    stuck = longest_stuck = 0
    while state.steps < max_steps and not state.solved():
        # The action question offers only legal moves as criteria, so ranking
        # all four here asks the model to score candidates it was never trained
        # on and lets a wall win.
        legal = state.legal_actions() or list(mz.ACTIONS)
        # The memory, applied before any policy sees the list, so `model+memory`
        # and `random+memory` differ in exactly one thing: who picks among the
        # survivors. Anything `random+memory` already achieves is the memory's
        # doing, not the network's.
        options = legal
        if controller.endswith("+memory"):
            r, c = state.position
            fresh = []
            for a in legal:
                dr, dc = mz.DIRECTIONS[a]
                if (r + dr, c + dc) not in seen:
                    fresh.append(a)
            options = fresh or legal

        if controller == "reference":
            probs = mz.optimal_action_distribution(state)
            action = max(probs, key=probs.get) if probs else rng.choice(legal)
            distribution = probs
        elif controller.startswith("random"):
            action, distribution = rng.choice(options), {}
        elif controller == "model-field":
            sample = maze_state_sample(state)
            action, distribution = field_action(runner, sample, state, legal, rng)
        else:
            sample = maze_state_sample(state)
            action, distribution = model_action(runner, sample, options, rng, temperature)
        before = state.position
        moved = state.step(action)
        if not moved:
            stuck += 1
            longest_stuck = max(longest_stuck, stuck)
        else:
            stuck = 0
            seen.add(state.position)
        yield {"position": list(before), "action": action, "moved": moved,
               "probabilities": {k: round(float(v), 4)
                                 for k, v in (distribution or {}).items()}}
    return {"solved": state.solved(), "steps": state.steps,
            "collisions": state.collisions, "distinct_cells": len(seen),
            "longest_stuck": longest_stuck,
            # Kept as a guard rather than a finding. It fired on every episode
            # while blocked moves were being offered as candidates; with only
            # legal moves on the list it should stay at zero, so a non-zero
            # value means something has regressed.
            "deadlocked": longest_stuck >= max(20, max_steps // 10)}


def snake_steps(runner, state, controller, max_steps, rng, temperature=0.0):
    """One decision per ``next()``. See ``maze_steps``."""
    while state.alive and state.steps < max_steps:
        legal = state.candidate_actions()
        if controller == "reference":
            probs, _ = sk.oracle_action_distribution(state)
            action = max(probs, key=probs.get) if probs else rng.choice(legal)
            distribution = probs
        elif controller == "random":
            action, distribution = rng.choice(legal), {}
        else:
            sample = snake_state_sample(state)
            action, distribution = model_action(runner, sample, legal, rng, temperature)
        head, food, body = state.head, state.food, list(state.body)
        state.step(action)
        yield {"head": list(head), "food": list(food),
               "body": [list(c) for c in body], "action": action,
               "probabilities": {k: round(float(v), 4)
                                 for k, v in (distribution or {}).items()}}
    return {"food": state.eaten, "steps": state.steps, "alive": state.alive,
            "death": state.death, "length": len(state.body)}


def default_max_steps(game: str, size: int) -> int:
    """The step budget an episode gets when none is given.

    Lives here rather than in each caller because the arcade's live server
    plays the same games through the same generators, and a budget that
    differed between the two would make a live episode and a recorded one
    disagree about what "failed" means -- the model looks worse or better
    purely by where it was run.
    """
    return (4 if game == "maze" else 2) * size * size


def _drain(steps, record):
    """Run a step generator to the end, recording as it goes.

    The wrappers below keep ``play.py``'s own callers unchanged: the generator
    is the shared thing, and running it to completion is what batch recording
    happens to want.
    """
    while True:
        try:
            step = next(steps)
        except StopIteration as stop:
            return stop.value
        if record is not None:
            record.append(step)


def play_maze(runner, state, controller, max_steps, rng, temperature=0.0,
              record=None):
    return _drain(maze_steps(runner, state, controller, max_steps, rng,
                             temperature), record)


def play_snake(runner, state, controller, max_steps, rng, temperature=0.0,
               record=None):
    return _drain(snake_steps(runner, state, controller, max_steps, rng,
                              temperature), record)


# ---------------------------------------------------------------- driver

def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--checkpoint", default="runs/jevon-final")
    # None means "prefer balanced.pt", the checkpoint every shipped number
    # comes from. See jevon.inference.resolve_weights.
    p.add_argument("--weights", default=None)
    p.add_argument("--game", choices=("maze", "snake", "both"), default="both")
    p.add_argument("--controller", default="model",
                   choices=("model", "model+memory", "model-field",
                            "reference", "random", "random+memory"))
    p.add_argument("--episodes", type=int, default=12)
    p.add_argument("--maze-sizes", default="11,21,31,51")
    p.add_argument("--maze-topologies", default="corridor,tree,loops,random_obstacle")
    p.add_argument("--snake-sizes", default="12")
    p.add_argument("--max-steps", type=int, default=0, help="0 => 4*size^2 for maze, 2*size^2 for snake")
    p.add_argument("--temperature", type=float, default=0.0)
    p.add_argument("--seed", type=int, default=17)
    p.add_argument("--device", default="mps")
    p.add_argument("--out", default=None)
    p.add_argument("--trace-out", default=None, help="write per-step replay JSON")
    p.add_argument("--trace-max-steps", type=int, default=0,
                   help="record at most this many steps per episode; the "
                        "episode still runs to completion, so summaries are "
                        "unaffected (0 => record everything)")
    args = p.parse_args()

    if args.controller == "model-field" and args.game != "maze":
        # Snake's field is distance to food, which is the wrong thing to
        # descend greedily -- survival dominates. Silently falling back to
        # the action head would mislabel a row in the results table.
        p.error("--controller model-field only applies to --game maze")

    rng = random.Random(args.seed)
    runner = None
    if args.controller.startswith("model"):
        runner = JevonRunner(args.checkpoint, args.device, args.weights)

    results = {"controller": args.controller, "episodes": []}
    traces = []

    if args.game in ("maze", "both"):
        for size in [int(s) for s in args.maze_sizes.split(",")]:
            for topology in args.maze_topologies.split(","):
                for i in range(args.episodes):
                    seed = 900_000 + i * 37 + size * 101
                    state = mz.generate_maze(size, topology, seed)
                    shortest = mz.bfs_distances(size, state.walls,
                                                state.goal)[state.position]
                    cap = args.max_steps or default_max_steps("maze", size)
                    record = Recorder(args.trace_max_steps) if args.trace_out else None
                    # Before play, because play mutates the agent's position --
                    # the field itself does not depend on it, but recording it
                    # first keeps that independence obvious rather than assumed.
                    field = episode_field(runner, state) if record is not None else None
                    out = play_maze(runner, state, args.controller, cap, rng,
                                    args.temperature, record)
                    out.update(game="maze", size=size, topology=topology, seed=seed,
                               shortest=shortest,
                               efficiency=shortest / out["steps"] if out["solved"] else 0.0)
                    results["episodes"].append(out)
                    if record is not None:
                        traces.append({"game": "maze", "size": size,
                                       "topology": topology, "seed": seed,
                                       "walls": sorted(map(list, state.walls)),
                                       "goal": list(state.goal), "result": out,
                                       "field": field, "steps": list(record),
                                       "truncated": record.truncated})
                    print(f"maze {size:3d} {topology:16s} seed={seed} "
                          f"solved={out['solved']} steps={out['steps']:5d} "
                          f"shortest={shortest} collisions={out['collisions']}",
                          flush=True)

    if args.game in ("snake", "both"):
        for size in [int(s) for s in args.snake_sizes.split(",")]:
            for i in range(args.episodes):
                seed = 610_000 + i * 53 + size * 7
                state = sk.new_game(size, seed)
                cap = args.max_steps or default_max_steps("snake", size)
                record = Recorder(args.trace_max_steps) if args.trace_out else None
                out = play_snake(runner, state, args.controller, cap, rng,
                                 args.temperature, record)
                out.update(game="snake", size=size, seed=seed)
                results["episodes"].append(out)
                if record is not None:
                    traces.append({"game": "snake", "size": size, "seed": seed,
                                   "result": out, "steps": list(record),
                                   "truncated": record.truncated})
                print(f"snake {size:3d} seed={seed} food={out['food']:3d} "
                      f"steps={out['steps']:4d} alive={out['alive']} "
                      f"death={out['death']}", flush=True)

    maze_eps = [e for e in results["episodes"] if e["game"] == "maze"]
    snake_eps = [e for e in results["episodes"] if e["game"] == "snake"]
    summary = {}
    if maze_eps:
        solved = [e for e in maze_eps if e["solved"]]
        summary["maze"] = {
            "episodes": len(maze_eps),
            "solve_rate": len(solved) / len(maze_eps),
            "mean_steps_when_solved": (sum(e["steps"] for e in solved) / len(solved)
                                       if solved else None),
            "mean_efficiency": (sum(e["efficiency"] for e in solved) / len(solved)
                                if solved else 0.0),
            "mean_collisions": sum(e["collisions"] for e in maze_eps) / len(maze_eps),
            "deadlock_rate": sum(e.get("deadlocked", False)
                                 for e in maze_eps) / len(maze_eps),
            "mean_distinct_cells": sum(e.get("distinct_cells", 0)
                                       for e in maze_eps) / len(maze_eps),
            "by_size": {},
        }
        for size in sorted({e["size"] for e in maze_eps}):
            rows = [e for e in maze_eps if e["size"] == size]
            ok = [e for e in rows if e["solved"]]
            summary["maze"]["by_size"][str(size)] = {
                "episodes": len(rows),
                "solve_rate": len(ok) / len(rows),
                "mean_steps": sum(e["steps"] for e in ok) / len(ok) if ok else None,
                # Over solved episodes only, matching the aggregate above: an
                # unsolved episode has efficiency 0 by definition, so including
                # them would make this a second, noisier solve rate rather than
                # a statement about path length. It is here because the
                # model-vs-control comparison separates on efficiency and not
                # on arrival, and that comparison is per size.
                "mean_efficiency": (sum(e["efficiency"] for e in ok) / len(ok)
                                    if ok else 0.0),
                "mean_shortest": sum(e["shortest"] for e in rows) / len(rows)}
    if snake_eps:
        summary["snake"] = {
            "episodes": len(snake_eps),
            "mean_food": sum(e["food"] for e in snake_eps) / len(snake_eps),
            "max_food": max(e["food"] for e in snake_eps),
            "mean_steps": sum(e["steps"] for e in snake_eps) / len(snake_eps),
            "survival_rate": sum(e["alive"] for e in snake_eps) / len(snake_eps)}
    results["summary"] = summary
    print(json.dumps(summary, indent=2))
    if args.out:
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.out).write_text(json.dumps(results, indent=2))
    if args.trace_out:
        Path(args.trace_out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.trace_out).write_text(json.dumps(traces))


if __name__ == "__main__":
    main()