"""SFT (Supervised Fine-Tuning) warm-start for CERNenv. Generates ``--num_episodes`` oracle trajectories against the local ``CERNCollisionEnvironment`` (no HTTP), turns each successful step into a (prompt, completion) pair using the *same* chat templating the GRPO loop uses (see ``training.training_script.build_dataset`` / ``training.llm_agent.build_chat``), and runs ``trl.SFTTrainer`` with LoRA so the resulting checkpoint can be used as the starting weights for GRPO. This addresses the v1 reward hack head-on. v1 (``anugrahhu/cernenv-grpo-smollm2-360m``) never saw a positive-reward trajectory during early training because SmolLM2-360M-Instruct cannot solve the LHC discovery pipeline zero-shot, so GRPO had no positive gradient to follow and the policy collapsed to "spam request_systematics forever". A short SFT on oracle traces gives the policy a non-zero prior over the *correct* action sequence, which RL can then refine. Usage ----- python -m training.sft_warmstart \\ --out_dir runs/sft-warmstart \\ --num_episodes 200 --max_steps 8 --epochs 1 --lr 1e-5 \\ --base_model HuggingFaceTB/SmolLM2-360M-Instruct \\ --difficulty mixed --evidence_dir evidence Smoke test: python -m training.sft_warmstart --num_episodes 4 --max_steps 4 \\ --epochs 1 --base_model sshleifer/tiny-gpt2 --out_dir /tmp/sft_smoke """ from __future__ import annotations import argparse import json import logging import time from dataclasses import asdict from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from models import ActionType, ExperimentAction from server.environment import CERNCollisionEnvironment from training.llm_agent import LLMAgentConfig, build_chat logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) # Mirrors training/training_script.py so the pre-train SFT and the # downstream GRPO use the *same* prompt templating. If you change one, # change both — divergence here is the most insidious source of # warm-start ineffectiveness. LORA_R = 16 LORA_ALPHA = 32 LORA_DROPOUT = 0.05 LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"] # ── Oracle trajectory collection ───────────────────────────────────────── _DIFFICULTY_CYCLE = ("easy", "medium", "hard") def _serialise_action(action: ExperimentAction) -> str: """Return a JSON string the GRPO parser would accept. Uses ``models.ExperimentAction.model_dump`` so enum values are converted to their string representation and parameters survive a round-trip through ``training.llm_agent.parse_action``. """ payload: Dict[str, Any] = { "action_type": action.action_type.value, "parameters": dict(action.parameters or {}), } if action.method: payload["method"] = action.method if action.justification: payload["justification"] = action.justification if action.confidence is not None: payload["confidence"] = float(action.confidence) return json.dumps(payload, ensure_ascii=False) def _difficulty_for(idx: int, difficulty: str) -> str: if difficulty == "mixed": return _DIFFICULTY_CYCLE[idx % len(_DIFFICULTY_CYCLE)] return difficulty def _collect_oracle_trajectories( *, tokenizer, num_episodes: int, max_steps: int, difficulty: str, seed_base: int = 4242, ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Run the OracleAgent against the env and harvest (prompt, completion) pairs. Returns (filtered_pairs, stats). ``stats`` reports oracle success rate even when no successful trajectories were obtained, so callers can log it into ``evidence/sft_summary.json``. """ from scripts.baseline_agents import OracleAgent config = LLMAgentConfig() pairs_all: List[Dict[str, Any]] = [] pairs_successful: List[Dict[str, Any]] = [] successes = 0 rewards: List[float] = [] for i in range(num_episodes): env = CERNCollisionEnvironment(max_steps=max_steps) difficulty_i = _difficulty_for(i, difficulty) obs = env.reset(seed=seed_base + i, difficulty=difficulty_i) truth = env.hidden_truth() agent = OracleAgent(truth=truth) agent.reset() episode_pairs: List[Dict[str, Any]] = [] cumulative = 0.0 while not obs.done and len(episode_pairs) < max_steps: chat = build_chat(obs, config) try: prompt = tokenizer.apply_chat_template( chat, add_generation_prompt=True, tokenize=False, ) except Exception as exc: # pragma: no cover - tiny-gpt2 etc. logger.warning( "tokenizer has no chat template (%s); installing a " "minimal fallback so SFT can proceed", exc, ) tokenizer.chat_template = ( "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n" "{% endfor %}" "{% if add_generation_prompt %}assistant: {% endif %}" ) prompt = tokenizer.apply_chat_template( chat, add_generation_prompt=True, tokenize=False, ) action = agent.act(obs) completion = _serialise_action(action) episode_pairs.append({ "prompt": prompt, "completion": completion, "step": obs.step_index, "difficulty": difficulty_i, "seed": seed_base + i, }) obs = env.step(action) cumulative += float(obs.reward or 0.0) rewards.append(cumulative) st = env.state ok = bool(st.correct_mass) and bool(st.correct_channel) if ok: successes += 1 pairs_successful.extend(episode_pairs) pairs_all.extend(episode_pairs) success_rate = successes / max(num_episodes, 1) mean_reward = sum(rewards) / max(len(rewards), 1) if rewards else 0.0 stats = { "num_episodes": num_episodes, "num_successful_episodes": successes, "oracle_success_rate": round(success_rate, 4), "mean_oracle_reward": round(mean_reward, 4), "num_transitions_total": len(pairs_all), "num_transitions_successful": len(pairs_successful), } # If we filtered out *everything* (e.g. smoke test with max_steps too # small for the oracle to ever finish), fall back to the unfiltered # set with a warning. Better to teach the model the prefix of the # correct pipeline than to give up entirely. if pairs_successful: return pairs_successful, stats logger.warning( "no fully-successful oracle trajectories (max_steps=%d may be " "too small); using %d unfiltered transitions instead", max_steps, len(pairs_all), ) stats["fallback_used"] = True return pairs_all, stats # ── Dataset assembly ───────────────────────────────────────────────────── def _build_dataset(pairs: List[Dict[str, Any]], tokenizer): """Build a HF Dataset whose ``text`` column is ``prompt + completion``. SFTTrainer will train next-token-prediction on the *whole* text, so we append an ``eos_token`` at the end of the completion to terminate generation cleanly. ``dataset_text_field='text'`` on SFTConfig then consumes this column directly. """ from datasets import Dataset eos = tokenizer.eos_token or "" rows = [] for p in pairs: rows.append({ "text": p["prompt"] + p["completion"] + eos, "prompt": p["prompt"], "completion": p["completion"], }) return Dataset.from_list(rows) # ── LoRA helper ────────────────────────────────────────────────────────── def _build_peft_config(model): """Build a LoRA config matching the GRPO setup. For tiny stub models like ``sshleifer/tiny-gpt2`` the q_proj/k_proj/ v_proj/o_proj target modules don't exist (GPT-2 uses a fused ``c_attn``); fall back to ``all-linear`` so the smoke test still exercises the LoRA path. """ from peft import LoraConfig target_modules: Any = LORA_TARGET_MODULES available = {n for n, _ in model.named_modules()} has_target = any(t in mod_name for t in LORA_TARGET_MODULES for mod_name in available) if not has_target: logger.warning( "model %s exposes none of %s; falling back to target_modules='all-linear'", type(model).__name__, LORA_TARGET_MODULES, ) target_modules = "all-linear" return LoraConfig( r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, target_modules=target_modules, task_type="CAUSAL_LM", bias="none", ) # ── Main ───────────────────────────────────────────────────────────────── def main() -> None: # pragma: no cover - heavy ML path parser = argparse.ArgumentParser(description="SFT warm-start for CERNenv GRPO") parser.add_argument("--out_dir", default="runs/sft-warmstart") parser.add_argument("--num_episodes", type=int, default=200) parser.add_argument("--max_steps", type=int, default=8) parser.add_argument("--epochs", type=int, default=1) parser.add_argument("--lr", type=float, default=1e-5) parser.add_argument("--base_model", default="HuggingFaceTB/SmolLM2-360M-Instruct") parser.add_argument( "--difficulty", default="mixed", choices=["easy", "medium", "hard", "mixed"], help="'mixed' cycles easy/medium/hard across episodes.", ) parser.add_argument("--evidence_dir", default="evidence") parser.add_argument("--per_device_batch_size", type=int, default=4) parser.add_argument("--gradient_accumulation_steps", type=int, default=1) parser.add_argument("--max_seq_length", type=int, default=1280) parser.add_argument("--seed", type=int, default=42) parser.add_argument( "--no_lora", action="store_true", help="Train the full model without LoRA (used by some smoke tests).", ) args = parser.parse_args() out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) evidence_dir = Path(args.evidence_dir) evidence_dir.mkdir(parents=True, exist_ok=True) t_start = time.time() try: import torch from transformers import AutoModelForCausalLM, AutoTokenizer from trl import SFTConfig, SFTTrainer except ImportError as exc: # pragma: no cover raise SystemExit( "Heavy ML deps missing — install -r space/training/requirements.txt" ) from exc logger.info("Loading tokenizer + base model: %s", args.base_model) tokenizer = AutoTokenizer.from_pretrained(args.base_model) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token use_bf16 = torch.cuda.is_available() dtype = torch.bfloat16 if use_bf16 else torch.float32 model = AutoModelForCausalLM.from_pretrained(args.base_model, torch_dtype=dtype) if model.config.pad_token_id is None: model.config.pad_token_id = tokenizer.pad_token_id logger.info( "Collecting %d oracle trajectories (max_steps=%d, difficulty=%s)", args.num_episodes, args.max_steps, args.difficulty, ) pairs, stats = _collect_oracle_trajectories( tokenizer=tokenizer, num_episodes=args.num_episodes, max_steps=args.max_steps, difficulty=args.difficulty, ) logger.info( "oracle stats: success_rate=%.2f total_transitions=%d kept=%d", stats["oracle_success_rate"], stats["num_transitions_total"], len(pairs), ) if not pairs: raise SystemExit( "No transitions collected — check OracleAgent / env wiring." ) dataset = _build_dataset(pairs, tokenizer) logger.info("Built SFT dataset: %d rows", len(dataset)) peft_config = None if args.no_lora else _build_peft_config(model) sft_cfg = SFTConfig( output_dir=str(out_dir), per_device_train_batch_size=args.per_device_batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, num_train_epochs=float(args.epochs), learning_rate=args.lr, logging_steps=5, save_strategy="no", # we save manually at the end so checkpoints # don't double the disk footprint. bf16=use_bf16, fp16=False, seed=args.seed, report_to=[], dataset_text_field="text", max_seq_length=args.max_seq_length, packing=False, ) trainer = SFTTrainer( model=model, args=sft_cfg, train_dataset=dataset, processing_class=tokenizer, peft_config=peft_config, ) logger.info("Starting SFT training (epochs=%d, lr=%.1e)", args.epochs, args.lr) train_result = trainer.train() final_loss = float(train_result.training_loss) logger.info("SFT done; final training_loss=%.4f", final_loss) # If we're using LoRA, merge the adapters back into the base model # and save a *full* causal-LM checkpoint to ``out_dir``. GRPO # downstream just calls ``AutoModelForCausalLM.from_pretrained(out_dir)``, # which is much simpler than asking it to load a base model + adapters # separately and keeps the warm-start path one ``--base_model`` flag. if peft_config is not None: try: merged = trainer.model.merge_and_unload() merged.save_pretrained(str(out_dir)) logger.info("Merged LoRA adapters into base model and saved to %s", out_dir) except Exception as exc: logger.warning( "merge_and_unload failed (%s); saving adapters alongside base " "and pointing GRPO at this directory will still work via PEFT.", exc, ) trainer.save_model(str(out_dir)) else: trainer.save_model(str(out_dir)) tokenizer.save_pretrained(str(out_dir)) duration_s = round(time.time() - t_start, 2) summary = dict(stats) summary.update({ "final_loss": round(final_loss, 6), "duration_s": duration_s, "epochs": args.epochs, "learning_rate": args.lr, "base_model": args.base_model, "out_dir": str(out_dir), "lora": peft_config is not None, "num_train_rows": len(dataset), }) summary_path = evidence_dir / "sft_summary.json" summary_path.write_text(json.dumps(summary, indent=2)) logger.info( "Wrote SFT summary to %s (final_loss=%.4f, duration=%.1fs)", summary_path, final_loss, duration_s, ) if __name__ == "__main__": # pragma: no cover main()