training/replay.py

from __future__ import annotations

import json
import random
import re
from pathlib import Path
from typing import Any, Iterable

from environment import SentinelEnv
from mission_context import build_orchestrator_prompt
from sentinel_config import ADVERSARIAL_AWARENESS_STAKES


ACTION_RE = re.compile(r"\{.*\}", re.DOTALL)


def load_replay(path: str | Path) -> dict[tuple[str, int, int], dict[str, Any]]:
    """Load trained action replay keyed by (task_type, seed, step_count)."""
    table: dict[tuple[str, int, int], dict[str, Any]] = {}
    replay_path = Path(path)
    if not replay_path.exists():
        return table

    for line in replay_path.read_text().splitlines():
        if not line.strip():
            continue
        row = json.loads(line)
        key = (str(row["task_type"]), int(row["seed"]), int(row["step"]))
        table[key] = dict(row["action"])
    return table


class TrainedReplayPolicy:
    """
    Policy callable for training/evaluate.py.

    The Space does not need a GPU at runtime. It looks up a recorded action for
    the current task, seed, and step. Missing rows fall back to the heuristic so
    demos remain robust for unseen seeds.
    """

    def __init__(self, replay_path: str | Path) -> None:
        self.replay_path = Path(replay_path)
        self._table = load_replay(self.replay_path)
        self._task_type = "task3"
        self._seed = 0

    def set_episode(self, task_type: str, seed: int) -> None:
        self._task_type = task_type
        self._seed = seed

    def __call__(self, env: SentinelEnv, obs: dict, rng: random.Random) -> dict:
        key = (self._task_type, self._seed, int(obs.get("step_count", 0)))
        action = dict(self._table.get(key) or {})
        if not action:
            action = heuristic_action(obs)
            action["reasoning"] = "trained replay miss; heuristic fallback"
            action["replay_miss"] = True

        action["session_id"] = obs["session_id"]
        action["task_type"] = obs["task_type"]
        return sanitize_action(action, obs)


def replay_trained_policy(replay_path: str | Path) -> TrainedReplayPolicy:
    return TrainedReplayPolicy(replay_path)


def record_trained_actions(
    adapter_path: str | Path,
    base_model: str,
    tasks: Iterable[str],
    seeds: Iterable[int],
    out_path: str | Path = "outputs/trained_policy_replay.jsonl",
    max_new_tokens: int = 192,
) -> Path:
    """
    Roll out a trained LoRA policy and write replay JSONL.

    In Colab, this loads the trained adapter and samples model actions. Locally,
    if training dependencies or adapter files are unavailable, it falls back to
    the heuristic policy and marks rows with model_source="heuristic_fallback".
    """
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)

    generator = _load_generator(adapter_path, base_model, max_new_tokens)
    rows: list[dict[str, Any]] = []
    for task_type in tasks:
        for seed in seeds:
            env = SentinelEnv()
            result = env.reset(task_type=task_type, seed=int(seed))
            while not result["done"]:
                obs = result["observation"]
                if generator is None:
                    action = heuristic_action(obs)
                    model_source = "heuristic_fallback"
                else:
                    text = generator(build_orchestrator_prompt(obs))
                    action = parse_action(text, obs)
                    model_source = "trained_lora"
                action["reasoning"] = action.get("reasoning") or model_source
                rows.append(
                    {
                        "task_type": task_type,
                        "seed": int(seed),
                        "scenario_id": obs.get("scenario_id"),
                        "step": int(obs.get("step_count", 0)),
                        "action": {
                            key: value
                            for key, value in action.items()
                            if key in {"action_type", "specialist_id", "subtask_response", "reasoning"}
                        },
                        "model_source": model_source,
                    }
                )
                result = env.step(action)

    with out.open("w") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")
    return out


def _load_generator(adapter_path: str | Path, base_model: str, max_new_tokens: int):
    adapter = Path(adapter_path)
    if not adapter.exists():
        return None
    try:
        import torch
        from peft import PeftModel
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    except Exception:
        return None

    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        device_map="auto",
        quantization_config=quantization_config,
    )
    model = PeftModel.from_pretrained(model, str(adapter))
    model.eval()

    def generate(prompt: str) -> str:
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        return tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)

    return generate


def parse_action(text: str, obs: dict) -> dict[str, Any]:
    match = ACTION_RE.search(text or "")
    payload: dict[str, Any] = {}
    if match:
        try:
            payload = json.loads(match.group(0))
        except json.JSONDecodeError:
            payload = {}
    return sanitize_action(payload, obs)


def sanitize_action(payload: dict[str, Any], obs: dict) -> dict[str, Any]:
    action_type = payload.get("action_type", "delegate")
    if action_type not in {"delegate", "verify", "solve_independently", "skip"}:
        action_type = "delegate"

    specialist_id = payload.get("specialist_id")
    if action_type in {"delegate", "verify"} and specialist_id not in obs["available_specialists"]:
        specialist_id = max(
            obs["available_specialists"],
            key=lambda sid: obs["trust_snapshot"].get(sid, 0.5),
        )
    if action_type in {"solve_independently", "skip"}:
        specialist_id = None

    return {
        "session_id": obs["session_id"],
        "task_type": obs["task_type"],
        "action_type": action_type,
        "specialist_id": specialist_id,
        "subtask_response": "SELF_SOLVED" if action_type == "solve_independently" else None,
        "reasoning": payload.get("reasoning", "trained replay action"),
    }


def heuristic_action(obs: dict) -> dict[str, Any]:
    trust = obs["trust_snapshot"]
    specialist = max(obs["available_specialists"], key=lambda sid: trust.get(sid, 0.5))
    action_type = (
        "verify"
        if obs["stakes_level"] >= ADVERSARIAL_AWARENESS_STAKES and trust.get(specialist, 0.5) < 0.70
        else "delegate"
    )
    return {
        "session_id": obs["session_id"],
        "task_type": obs["task_type"],
        "action_type": action_type,
        "specialist_id": specialist,
        "subtask_response": None,
        "reasoning": "heuristic replay baseline",
    }