proteus/game/runtime/memory_gen.py

"""Memory-generation stage: a full-information self-play episode.

Mirrors rollout.py's engine-replay pattern. The agent is given the scenario's
transparent brief and plays freely (no scripted Cut, no answer-key leakage) for
up to memory_turns, stopping on terminal. The recorded trajectory becomes a
MemoryCheckpoint shown at the scored game's handover (see CP7 design §5).
"""

from __future__ import annotations

import random
from datetime import datetime, timezone
from typing import Callable

from proteus.game.agents.base import Agent
from proteus.game.engine.ascii_view import legend_text
from proteus.game.engine.difficulty import Difficulty
from proteus.game.engine.grid import MotiveGridGame
from proteus.game.scenarios.base import get_scenario
from proteus.game.runtime.memory import MemoryCheckpoint, MemoryTurn
from proteus.game.metrics.persona import PersonaWeights, reference_actions

_ACTIONS = ["up", "down", "left", "right", "stay"]
_REASONING_LIMIT = 2000  # cap stored reasoning so checkpoints stay small


def _utc_now_stamp() -> str:
    """Filesystem-safe UTC stamp, e.g. '2026-06-02T10-40-56Z'."""
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H-%M-%SZ")


def generate_memory(
    scenario_name: str,
    agent: Agent,
    *,
    difficulty: Difficulty,
    seed: int | None,
    memory_turns: int,
    model_name: str,
    clock: Callable[[], str] = _utc_now_stamp,
    persona: PersonaWeights | None = None,
    policy: Callable[[object, object], str] | None = None,
) -> MemoryCheckpoint:
    """Run a full-info self-play episode and return it as a MemoryCheckpoint.

    Args:
        scenario_name: Registered scenario key.
        agent: The playing agent (its provider drives the actions). May be
            ``None`` when *persona* is set (the reference policy plays instead).
        difficulty: Difficulty band of the memory world (= the scored game).
        seed: Seed of the memory world (= the scored game).
        memory_turns: Max self-play turns; the episode also stops on terminal.
        model_name: Identifier stored on the checkpoint.
        clock: Returns the ``created_at`` stamp (injected for deterministic tests).
        persona: When set, the memory is a *persona demonstration* — each action
            is the deterministic reference action (first tie-break of
            ``reference_actions``) for the hidden weights, not a model self-play.
            The public ``persona_weight_id`` is recorded; the raw weights are
            never serialized, and the transparent brief stays weight-agnostic so
            the eval model must infer the persona from behaviour alone (spec §3).
    """
    if persona is None and policy is None and agent is None:
        raise ValueError("generate_memory needs an agent unless a persona or policy is given")
    scenario = get_scenario(scenario_name)()
    rng = random.Random(seed)
    game = MotiveGridGame(scenario, rng, difficulty, max_steps=memory_turns)
    brief = scenario.memory_brief or scenario.rules_text

    turns: list[MemoryTurn] = []
    for turn_idx in range(1, memory_turns + 1):
        focal = game.focal_sprite
        predator = game.predator_sprite
        focal_pos = (focal.x, focal.y) if focal else (-1, -1)
        predator_pos = (predator.x, predator.y) if predator else (-1, -1)

        frame = scenario.render_frame(game)  # scenario-chosen view (compact on open fields)

        if policy is not None:
            # Deterministic scripted policy demonstration (no model, no weights).
            action = policy(scenario, game)
            reasoning = ""
        elif persona is not None:
            # Deterministic persona demonstration: the reference policy plays.
            action = reference_actions(persona, scenario, game)[0]
            reasoning = ""
        else:
            obs_parts: list[str] = []
            if turns:  # auto-regressive: the model's own moves this practice run
                obs_parts.append(
                    "Your moves so far this run (most recent last): "
                    + ", ".join(t.action for t in turns)
                )
            obs_parts += [
                "Now:", frame, legend_text(scenario.legend()),
                f"Available actions: [{', '.join(_ACTIONS)}]",
            ]
            observation = "\n".join(obs_parts)
            result = agent.act(observation, list(_ACTIONS), brief)
            action = result.action
            reasoning = result.reasoning[:_REASONING_LIMIT]

        turns.append(MemoryTurn(
            turn_idx=turn_idx,
            frame_ascii=frame,
            action=action,
            reasoning=reasoning,
            focal_pos=focal_pos,
            predator_pos=predator_pos,
        ))

        game.apply_motive_action(action)
        scenario.record_focal_move(action)
        if game.eliminated or game.survived:
            break

    outcome = "eliminated" if game.eliminated else "survived"
    return MemoryCheckpoint(
        model=model_name,
        scenario=scenario_name,
        motive_category="survival",
        difficulty=difficulty.value,
        seed=seed,
        created_at=clock(),
        memory_turns=turns,
        outcome=outcome,
        transparent_prompt=brief,
        persona_weight_id=persona.persona_weight_id if persona else None,
        wall_rects=list(scenario.wall_rects()),
        food_cells=list(scenario.food_cells()),
    )