gen_sample_replay.py

"""Generate showcase/data/replay.json from a local MCTS episode."""

from __future__ import annotations

import json
from pathlib import Path

from env import Game2048Env
from policies import mcts_policy

OUT = Path("showcase/data/replay.json")


def main() -> None:
    env = Game2048Env()
    policy = mcts_policy(__import__("numpy").random.default_rng(42), simulations=10)
    obs = env.reset(seed=42)
    turns: list[dict] = []
    step = 0
    while True:
        action = policy(env)
        result = env.step(action)
        step += 1
        turns.append(
            {
                "step": step,
                "observation": obs,
                "reasoning": f"MCTS chose {action} (valid: {json.loads(result.info['valid_moves'])})",
                "action": action,
                "reward": result.reward,
                "terminated": result.terminated,
                "truncated": result.truncated,
                "info": result.info,
            }
        )
        obs = result.observation
        if result.terminated or result.truncated:
            break

    payload = {
        "run": {"scores": {"mean_reward": sum(t["reward"] for t in turns)}},
        "episodes": [{"id": "sample-ep-1", "seed": 42, "total_reward": sum(t["reward"] for t in turns)}],
        "replay": {"sample-ep-1": turns},
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    js = Path("showcase/data/replay.js")
    js.write_text("window.REPLAY = " + json.dumps(payload) + ";\n", encoding="utf-8")
    print(f"wrote {OUT} ({len(turns)} turns)")


if __name__ == "__main__":
    main()