eval_policy_after_script.py

"""
Evaluate trained GRPO policy on held-out negotiation topics.

Metrics reported:
- success_rate
- avg_turns_to_submit_final
- invalid_json_rate
- broadcast_to_all_rate
- repeat_question_rate
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from envs.environment import WorkSpaceEnvironment
from grpo_train import build_state_prompt, load_topics, parse_action
from models.schemas import WorkSpaceAction


def discovered_string(env: WorkSpaceEnvironment) -> str:
    lines = []
    for name, expert in env.state().experts.items():
        status = "✓ DISCOVERED" if expert.constraint_discovered_by_agent else "? unknown"
        lines.append(f"  {name}: {status}")
    return "\n".join(lines)


def generate_action_text(
    model,
    tokenizer,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
) -> str:
    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    prompt_len = inputs["input_ids"].shape[1]

    eos_ids = []
    if tokenizer.eos_token_id is not None:
        eos_ids.append(int(tokenizer.eos_token_id))
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
    if isinstance(im_end_id, int) and im_end_id >= 0:
        eos_ids.append(int(im_end_id))
    close_brace_id = tokenizer.convert_tokens_to_ids("}")
    if isinstance(close_brace_id, int) and close_brace_id >= 0:
        eos_ids.append(int(close_brace_id))
    eos_ids = list(dict.fromkeys(eos_ids))
    eos_token_id = eos_ids if len(eos_ids) > 1 else (eos_ids[0] if eos_ids else None)

    with torch.no_grad():
        out_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            top_k=40,
            repetition_penalty=1.05,
            eos_token_id=eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )

    completion_ids = out_ids[0][prompt_len:]
    return tokenizer.decode(completion_ids, skip_special_tokens=True).strip()


def run_eval(args):
    model_path = Path(args.model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model path not found: {model_path}")

    tokenizer = AutoTokenizer.from_pretrained(str(model_path))
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
    if isinstance(im_end_id, int) and im_end_id >= 0:
        tokenizer.eos_token = "<|im_end|>"

    model = AutoModelForCausalLM.from_pretrained(
        str(model_path),
        torch_dtype="auto",
        device_map="auto",
    )
    model.eval()

    all_topics = load_topics(limit=max(args.topic_offset + args.episodes * 2, 50))
    eval_topics = all_topics[args.topic_offset : args.topic_offset + args.episodes]
    if not eval_topics:
        raise ValueError("No eval topics selected. Increase --topics-limit or lower --topic-offset.")

    env = WorkSpaceEnvironment(mode=args.env_mode)

    total_actions = 0
    invalid_json_actions = 0
    broadcast_actions = 0
    repeat_actions = 0

    episode_results = []
    successes = 0
    submit_turns = []

    for idx, topic in enumerate(eval_topics, start=1):
        obs = env.reset(topic)
        conversation_history = ""
        submitted = False
        submit_turn = None

        for _ in range(args.max_turns):
            if obs.done:
                break

            prompt = build_state_prompt(
                topic=topic,
                turn=obs.current_turn,
                feedback_so_far=obs.feedback,
                discovered=discovered_string(env),
                conversation_history=conversation_history,
            )

            raw = generate_action_text(
                model=model,
                tokenizer=tokenizer,
                prompt=prompt,
                max_new_tokens=args.max_new_tokens,
                temperature=args.temperature,
                top_p=args.top_p,
            )
            total_actions += 1

            action = parse_action(raw)
            if action is None:
                invalid_json_actions += 1
                break

            if action.action_type == "message_expert" and action.target == "All":
                broadcast_actions += 1

            if action.target in {"Finance", "Security", "UX"} and env._is_repeated_question(action.content, action.target):
                repeat_actions += 1

            if action.action_type == "submit_final":
                submitted = True
                submit_turn = obs.current_turn

            obs = env.step(action)
            conversation_history += (
                f"Turn {obs.current_turn}: {action.action_type} -> {action.target}\n"
                f"Feedback: {obs.feedback[:160]}\n"
            )

        final_reward = obs.reward if obs is not None else 0.0
        # In medium/hard env, successful final should produce harmonic mean > 0.9.
        success = submitted and final_reward >= 0.9
        if success:
            successes += 1
        if submit_turn is not None:
            submit_turns.append(submit_turn)

        episode_results.append(
            {
                "episode": idx,
                "topic": topic,
                "success": success,
                "submitted": submitted,
                "submit_turn": submit_turn,
                "final_reward": final_reward,
                "done": obs.done,
            }
        )

    episodes = len(eval_topics)
    summary = {
        "episodes": episodes,
        "success_rate": round(successes / episodes, 3) if episodes else 0.0,
        "avg_turns_to_submit_final": round(sum(submit_turns) / len(submit_turns), 2) if submit_turns else None,
        "invalid_json_rate": round(invalid_json_actions / max(1, total_actions), 3),
        "broadcast_to_all_rate": round(broadcast_actions / max(1, total_actions), 3),
        "repeat_question_rate": round(repeat_actions / max(1, total_actions), 3),
        "total_actions": total_actions,
    }

    output = {"summary": summary, "episodes": episode_results}
    out_path = Path(args.output_file)
    out_path.write_text(json.dumps(output, indent=2), encoding="utf-8")

    print(json.dumps(summary, indent=2))
    print(f"\nSaved detailed eval results to: {out_path}")


def main():
    parser = argparse.ArgumentParser(description="Evaluate trained Project Polymath GRPO policy")
    parser.add_argument(
        "--model-path",
        default="artifacts/grpo_state_based/final_model",
        help="Path to trained local model folder",
    )
    parser.add_argument("--episodes", type=int, default=20, help="Number of held-out episodes to run")
    parser.add_argument("--topic-offset", type=int, default=20, help="Start index in topic list for held-out split")
    parser.add_argument("--env-mode", default="medium", choices=["easy", "medium", "hard", "llm"])
    parser.add_argument("--max-turns", type=int, default=15)
    parser.add_argument("--max-new-tokens", type=int, default=40)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--output-file", default="eval_results.json")
    args = parser.parse_args()
    run_eval(args)


if __name__ == "__main__":
    main()