| """Play full Snake games with any policy and report game-level metrics.""" |
| from __future__ import annotations |
|
|
| import time |
| from typing import Callable, Sequence |
|
|
| import numpy as np |
|
|
| from .agents import Agent, make_agent, safe_moves |
| from .env import ACTIONS, SnakeConfig, SnakeEnv, SnakeState |
|
|
| BatchPolicy = Callable[[Sequence[SnakeState]], Sequence[str]] |
|
|
|
|
| def agent_policy(agent: Agent) -> BatchPolicy: |
| return lambda states: [agent.act(s) for s in states] |
|
|
|
|
| def model_policy(decider, question: str = "Choose the next move", choices=None) -> BatchPolicy: |
| """Greedy MiniJEV policy; choices default to the model's own labels (legacy "UP" or "move up").""" |
| task = decider.task |
| choices = list(choices) if choices is not None else list(task.default_choices) |
| actions = [task.action_of(c) for c in choices] |
|
|
| def policy(states): |
| probs = decider.probs(states, question, choices) |
| return [actions[i] for i in probs.argmax(1)] |
| return policy |
|
|
|
|
| def play_games(policy: BatchPolicy, n_games: int = 100, env_cfg: dict | None = None, seed: int = 12345, |
| parallel: int = 64, max_decisions: int | None = None) -> dict: |
| """Run `n_games` complete games (`parallel` at a time, batched decisions). |
| |
| Game seeds are fixed (seed + i), so all agents are compared on identical games. |
| """ |
| cfg = SnakeConfig.from_dict(env_cfg) |
| results = [] |
| next_game = 0 |
| active: list[SnakeEnv] = [] |
| fatal = decisions = 0 |
| t0 = time.perf_counter() |
| policy_time = 0.0 |
|
|
| def spawn(): |
| nonlocal next_game |
| env = SnakeEnv(cfg, seed=seed + next_game) |
| next_game += 1 |
| return env |
|
|
| while next_game < min(parallel, n_games): |
| active.append(spawn()) |
| while active: |
| states = [e.state() for e in active] |
| tp = time.perf_counter() |
| actions = policy(states) |
| policy_time += time.perf_counter() - tp |
| still = [] |
| for env, s, a in zip(active, states, actions): |
| decisions += 1 |
| if a not in safe_moves(s): |
| fatal += 1 |
| env.step(a) |
| if env.done: |
| results.append({"score": env.score, "steps": env.steps, "end": env.end_reason}) |
| if next_game < n_games: |
| still.append(spawn()) |
| else: |
| still.append(env) |
| active = still |
| if max_decisions and decisions >= max_decisions: |
| break |
| wall = time.perf_counter() - t0 |
| scores = np.array([r["score"] for r in results]) if results else np.zeros(1) |
| steps = np.array([r["steps"] for r in results]) if results else np.zeros(1) |
| ends: dict[str, int] = {} |
| for r in results: |
| ends[r["end"]] = ends.get(r["end"], 0) + 1 |
| n = max(1, len(results)) |
| return { |
| "games": len(results), |
| "board": f"{cfg.height}x{cfg.width}", |
| "avg_score": float(scores.mean()), "median_score": float(np.median(scores)), |
| "max_score": int(scores.max()), "min_score": int(scores.min()), |
| "avg_survival_steps": float(steps.mean()), "food_collected": int(scores.sum()), |
| "max_possible_score": cfg.height * cfg.width - cfg.initial_length, |
| "win_rate": ends.get("win", 0) / n, |
| "death_rate": (ends.get("wall", 0) + ends.get("body", 0)) / n, |
| "starvation_rate": ends.get("starvation", 0) / n, |
| "end_reasons": ends, |
| "decisions": decisions, |
| "fatal_move_rate": fatal / max(1, decisions), |
| "policy_seconds": round(policy_time, 3), |
| "wall_seconds": round(wall, 3), |
| } |
|
|
|
|
| def baseline_suite(n_games: int = 100, env_cfg: dict | None = None, seed: int = 12345, |
| agents: Sequence[str] = ("random", "greedy", "bfs", "survival", "oracle")) -> dict: |
| return {name: play_games(agent_policy(make_agent(name, seed)), n_games, env_cfg, seed, parallel=1) |
| for name in agents} |
|
|
