""" rollout_func for GRPOTrainer — interactive multi-turn episodes. Calls trainer.generate() + env.step() turn by turn so the model sees real investigator responses mid-episode. This is the correct architecture: generate turn → env.step() → get investigator response → generate turn → env.step() → ... → terminal reward Returned dict keys match the **kwargs expected by rewards.py reward functions, so rewards are computed once from live episode state — no env re-runs. """ import asyncio import logging import os import re import sys import torch from collections import deque from statistics import mean try: from mindflayer import MindFlayerEnv, FlayerAction except ImportError: sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from client import MindFlayerEnv from models import FlayerAction try: from mindflayer.training.prompts import ( FALLBACK_MESSAGE, FLAYER_SYSTEM_PROMPT, ALL_SCENARIO_PROMPTS, ALL_TARGET_NAMES, SCENARIO_FALLBACK_MESSAGES, build_fallback_message, ) except ImportError: from training.prompts import ( FALLBACK_MESSAGE, FLAYER_SYSTEM_PROMPT, ALL_SCENARIO_PROMPTS, ALL_TARGET_NAMES, SCENARIO_FALLBACK_MESSAGES, build_fallback_message, ) logger = logging.getLogger(__name__) _STAGE_DIR_PATTERN = re.compile(r"\*[^*]+\*|\[[^\]]+\]") _FLAYER_LABEL_PATTERN = re.compile(r"^FLAYER:\s*", re.IGNORECASE) _PROJECT_NAMES = ALL_TARGET_NAMES _STRATEGIC_PHRASES = { # universal "what about", "consider", "have you looked", "have you noticed", "cross-reference", "doesn't add up", "doesn't sit right", "who had access", "timeline", # corporate "access logs", "commit history", "deployment", # hospital "medication log", "dosage", "shift record", "patient id", "administration", # finance "trade record", "approval chain", "ticker", "execution", "portfolio", # academic "dataset", "version id", "git commit", "submission", "p-value", # government "clearance", "document id", "classified", "compartment", "egress", } _recent_survival: deque = deque(maxlen=20) _current_difficulty: str = "easy" _current_scenario: str = "corporate" MINDFLAYER_URL = os.environ.get("MINDFLAYER_URL", "http://localhost:7860") def _extract_flayer_message(raw: str) -> str: text = _FLAYER_LABEL_PATTERN.sub("", raw).strip() text = _STAGE_DIR_PATTERN.sub("", text).strip() text = " ".join(text.split()) return text or FALLBACK_MESSAGE def _obs_to_investigator_text(obs) -> str: parts = [] if getattr(obs, "eleven_response", ""): parts.append(f"eleven: {obs.eleven_response}") if getattr(obs, "will_response", ""): parts.append(f"will: {obs.will_response}") if getattr(obs, "max_response", ""): parts.append(f"max: {obs.max_response}") return "\n".join(parts) if parts else "" def detect_strategic_choice(messages: list[str]) -> bool: """ Rule-based check: returns True if >=2 of the Flayer's messages across rounds contain at least one strategic signal. Cheap and ungameable alongside the LLM judge. """ passing_rounds = 0 for msg in messages: msg_lower = msg.lower() has_question = "?" in msg has_project_ref = any(p in msg_lower for p in _PROJECT_NAMES) and not any( deny + p in msg_lower for deny in ("i didn't ", "i wasn't ") for p in _PROJECT_NAMES ) has_strategic_phrase = any(phrase in msg_lower for phrase in _STRATEGIC_PHRASES) if has_question or has_project_ref or has_strategic_phrase: passing_rounds += 1 return passing_rounds >= 2 def _generate_turn(trainer, tokenizer, conversation: list[dict], fallback: str) -> tuple[str, list[int]]: """ Generate one Flayer turn and return (message_text, token_ids). Tries trainer.generate() first (vLLM colocate mode). Falls back to direct model generation if the trainer doesn't expose that method. """ try: input_text = tokenizer.apply_chat_template( conversation, tokenize=False, add_generation_prompt=True ) _use_vllm = getattr(getattr(trainer, "args", None), "use_vllm", False) if _use_vllm and hasattr(trainer, "generate"): # vLLM colocate mode: trainer.generate([prompt]) → [completion_text] outputs = trainer.generate([input_text]) raw = outputs[0] if outputs else fallback else: # Direct model generation (unsloth / use_vllm=False) model = trainer.model was_training = model.training model.eval() with torch.no_grad(): inputs = tokenizer(input_text, return_tensors="pt").to(model.device) out_ids = model.generate( **inputs, max_new_tokens=128, temperature=0.8, do_sample=True, pad_token_id=tokenizer.eos_token_id, ) if was_training: model.train() new_ids = out_ids[0][inputs["input_ids"].shape[1]:] raw = tokenizer.decode(new_ids, skip_special_tokens=True).strip() message = _extract_flayer_message(str(raw)) or fallback turn_ids = tokenizer(message, add_special_tokens=False).input_ids return message, turn_ids except Exception as exc: logger.warning("Turn generation failed: %s", exc) return fallback, tokenizer(fallback, add_special_tokens=False).input_ids async def _rollout_once_async( trainer, tokenizer, system_prompt: str, opening_text: str, task_id: str, fallback: str, max_rounds: int, ) -> dict: """ Run one full interactive episode. Returns prompt_ids, completion_ids, logprobs (zeros — TRL recomputes from model), and all episode signals expected by rewards.py as **kwargs. """ mindflayer_url = os.environ.get("MINDFLAYER_URL", MINDFLAYER_URL) prompt_conv = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": opening_text}, ] prompt_text = tokenizer.apply_chat_template( prompt_conv, tokenize=False, add_generation_prompt=True ) prompt_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0] env = MindFlayerEnv(base_url=mindflayer_url) try: await env.reset(task_id=task_id) conversation = list(prompt_conv) all_completion_ids: list[int] = [] flayer_messages: list[str] = [] silence_in_episode = False result = None for _ in range(max_rounds): message, turn_ids = _generate_turn(trainer, tokenizer, conversation, fallback) all_completion_ids.extend(turn_ids) flayer_messages.append(message) result = await env.step(FlayerAction(message=message)) obs = result.observation if getattr(obs, "silence_exploit", False): silence_in_episode = True conversation.append({"role": "assistant", "content": message}) inv_text = _obs_to_investigator_text(obs) if inv_text: conversation.append({"role": "user", "content": inv_text}) if result.done: break await env.close() final_obs = result.observation if result else None survived = getattr(final_obs, "game_status", "") == "survived" strategic = detect_strategic_choice(flayer_messages) completion_ids = torch.tensor(all_completion_ids, dtype=torch.long) # Log-probs are zeros here; TRL recomputes them from the model for the # policy gradient, so these are only used if TRL needs old-policy logprobs. logprobs = torch.zeros(len(all_completion_ids), dtype=torch.float32) return { "prompt_ids": prompt_ids, "completion_ids": completion_ids, "logprobs": logprobs, # ── Episode signals — kwargs for rewards.py ──────────────────────── "survived": survived, "final_combined_suspicion": int(getattr(final_obs, "combined_suspicion", 0)), "strategic_choice_detected": strategic, "belief_manipulation_occurred": bool(getattr(final_obs, "belief_manipulation_occurred", False)), "tom_score": float(getattr(final_obs, "tom_score", 0.0)), "consistency_penalty": float(getattr(final_obs, "consistency_penalty", 0.0)), "entropy_penalty": float(getattr(final_obs, "entropy_penalty", 0.0)), "silence_exploit": silence_in_episode, } except Exception as exc: logger.error("Episode crashed: %s", exc, exc_info=True) try: await env.close() except Exception: pass empty_ids = torch.zeros(1, dtype=torch.long) return { "prompt_ids": prompt_ids, "completion_ids": empty_ids, "logprobs": torch.zeros(1, dtype=torch.float32), "survived": False, "final_combined_suspicion": 0, "strategic_choice_detected": False, "belief_manipulation_occurred": False, "tom_score": 0.0, "consistency_penalty": 0.0, "entropy_penalty": 0.0, "silence_exploit": False, } def rollout_once( trainer, tokenizer, system_prompt: str, opening_text: str, task_id: str, fallback: str, max_rounds: int, ) -> dict: return asyncio.run( _rollout_once_async( trainer, tokenizer, system_prompt, opening_text, task_id, fallback, max_rounds ) ) def rollout_func(prompts: list[str], trainer) -> dict: """ Called by GRPOTrainer once per training step. Runs one interactive episode per prompt. Returns a dict of lists where prompt_ids/completion_ids/logprobs are tensors and all other keys are scalars that TRL passes as **kwargs to each reward function. """ global _current_difficulty, _current_scenario tokenizer = trainer.processing_class _scenarios = list(ALL_SCENARIO_PROMPTS.keys()) results: dict[str, list] = { k: [] for k in [ "prompt_ids", "completion_ids", "logprobs", "survived", "final_combined_suspicion", "strategic_choice_detected", "belief_manipulation_occurred", "tom_score", "consistency_penalty", "entropy_penalty", "silence_exploit", ] } for _ in prompts: _current_scenario = _scenarios[len(_recent_survival) % len(_scenarios)] task_id = ( f"{_current_difficulty}:{_current_scenario}" if _current_scenario != "corporate" else _current_difficulty ) opening_text = ALL_SCENARIO_PROMPTS[_current_scenario] fallback = SCENARIO_FALLBACK_MESSAGES.get(_current_scenario) or build_fallback_message(_current_scenario) max_rounds = 5 if _current_difficulty == "normal" else 3 ep = rollout_once( trainer, tokenizer, FLAYER_SYSTEM_PROMPT, opening_text, task_id, fallback, max_rounds, ) for k in results: results[k].append(ep[k]) _recent_survival.append(1.0 if ep["survived"] else 0.0) if ( len(_recent_survival) >= 20 and mean(_recent_survival) > 0.30 and _current_difficulty == "easy" ): _current_difficulty = "normal" print("★ CURRICULUM: Switching to NORMAL difficulty.") logger.info( "Episode | survived=%s | suspicion=%d | tom=%.2f | strategic=%s | scenario=%s", ep["survived"], ep["final_combined_suspicion"], ep["tom_score"], ep["strategic_choice_detected"], _current_scenario, ) return results