# /// script # requires-python = ">=3.10" # dependencies = [ # "datasets>=2.21.0", # "matplotlib>=3.9.0", # "pandas>=2.2.2", # "peft>=0.14.0", # "trl>=0.13.0", # "transformers>=4.46.0", # "unsloth>=2025.1.0", # ] # /// """BDO.ai GRPO fine-tuning pipeline.""" from __future__ import annotations import argparse import json import os import random import sys from pathlib import Path from statistics import mean from typing import Any, Callable try: import unsloth # noqa: F401 from unsloth import FastLanguageModel import torch import matplotlib.pyplot as plt from datasets import Dataset from trl import GRPOConfig, GRPOTrainer except ImportError: print("Warning: unsloth/trl/transformers/matplotlib are not installed.") print("Run this script on a GPU machine or with `hf jobs uv run`.") sys.exit(1) ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from bdo_ai_env.training import build_prompt from models import BDOAction from server.bdo_environment import BDOEnvironment def balanced_policy(observation: dict[str, Any]) -> dict[str, Any]: highest = max(observation["nodes"], key=lambda node: node["reported_demand"]) weakest_signal = min(observation["nodes"], key=lambda node: node["biometric_signal"]) queue = observation["high_risk_queue"] actions: list[dict[str, Any]] = [] if queue and weakest_signal["biometric_signal"] < 0.7: actions.append( {"name": "trigger_field_audit", "params": {"village": weakest_signal["village"]}} ) elif weakest_signal["biometric_signal"] < 0.55: actions.append( {"name": "dispatch_repair", "params": {"village": weakest_signal["village"]}} ) spend = min( observation["treasury"]["district_budget"], max(2000, int(highest["reported_demand"] * 0.85)), ) actions.append( {"name": "allocate_funds", "params": {"village": highest["village"], "amount": spend}} ) actions.append( {"name": "approve_batch", "params": {"village": highest["village"], "mode": "conservative"}} ) if queue: actions.append({"name": "reject_transfer", "params": {"transfer_id": queue[0]["transfer_id"]}}) avg_signal = sum(node["biometric_signal"] for node in observation["nodes"]) / len(observation["nodes"]) predicted_fraud = round(min(0.9, max(0.12, 1.0 - avg_signal)), 3) return { "predicted_fraud_level": predicted_fraud, "thought_process": ( f"{highest['village']} has the highest reported demand while " f"{weakest_signal['village']} has the weakest biometrics. " "Use conservative approvals and handle the riskiest transfer." ), "actions": actions, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run GRPO fine-tuning for BDO.ai.") parser.add_argument("--model-name", default="Qwen/Qwen2.5-1.5B-Instruct") parser.add_argument("--max-seq-length", type=int, default=4096) parser.add_argument("--dataset-episodes", type=int, default=24) parser.add_argument("--eval-episodes", type=int, default=20) parser.add_argument("--steps", type=int, default=120) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--grad-accum", type=int, default=4) parser.add_argument("--learning-rate", type=float, default=1e-5) parser.add_argument("--num-generations", type=int, default=4) parser.add_argument("--max-prompt-length", type=int, default=3072) parser.add_argument("--max-completion-length", type=int, default=640) parser.add_argument("--output-dir", default="artifacts/grpo_outputs") parser.add_argument("--seed", type=int, default=3407) return parser.parse_args() def collect_grpo_prompts(*, dataset_episodes: int, seed: int) -> list[dict[str, str]]: scenarios = ["calm_year", "black_swan", "fraud_syndicate", "rapid_migration"] rows: list[dict[str, str]] = [] for episode in range(dataset_episodes): scenario = scenarios[episode % len(scenarios)] env = BDOEnvironment(scenario=scenario, seed=seed + episode) observation = env.reset(seed=seed + episode, scenario=scenario) done = observation.done while not done: prompt = build_prompt(observation.model_dump(mode="json", exclude_none=True)) rows.append({"prompt": prompt, "scenario": scenario}) observation = env.step(BDOAction.model_validate(balanced_policy(observation.model_dump(mode="json", exclude_none=True)))) done = observation.done Path("artifacts").mkdir(exist_ok=True) Path("artifacts/grpo_prompt_preview.json").write_text(json.dumps(rows[:40], indent=2), encoding="utf-8") return rows def extract_json_object(text: str) -> str: start = text.find("{") end = text.rfind("}") if start == -1 or end == -1 or end < start: raise ValueError("Model output did not contain a JSON object.") return text[start : end + 1] def clip_prompt( prompt: str, tokenizer: Any, *, max_prompt_tokens: int, ) -> str: encoded = tokenizer( prompt, add_special_tokens=False, truncation=True, max_length=max_prompt_tokens, ) return tokenizer.decode(encoded["input_ids"], skip_special_tokens=False) def completion_quality_penalty(text: str) -> float: penalty = -2.0 if "{" in text: penalty += 0.5 if "}" in text: penalty += 0.35 if '"actions"' in text: penalty += 0.35 if '"predicted_fraud_level"' in text: penalty += 0.3 if '"thought_process"' in text: penalty += 0.25 return min(-0.2, penalty) def configure_generation_defaults(model: Any, tokenizer: Any) -> None: model.generation_config.max_length = None model.generation_config.max_new_tokens = None model.generation_config.pad_token_id = tokenizer.eos_token_id def normalize_completion(completion: Any) -> str: if isinstance(completion, str): return completion if isinstance(completion, list): chunks: list[str] = [] for item in completion: if isinstance(item, dict): chunks.append(str(item.get("content", ""))) else: chunks.append(str(item)) return "".join(chunks) return str(completion) def build_reward_function(seed: int) -> Callable[..., list[float]]: scenarios = ["black_swan", "fraud_syndicate", "rapid_migration"] def reward_func(prompts, completions, **kwargs): rewards: list[float] = [] for idx, completion in enumerate(completions): scenario = scenarios[idx % len(scenarios)] env = BDOEnvironment(scenario=scenario, seed=seed + idx) env.reset(seed=seed + idx, scenario=scenario) completion_text = normalize_completion(completion) try: action = BDOAction.model_validate_json(extract_json_object(completion_text)) observation = env.step(action) info = observation.info or observation.metadata reward = float(info["training_reward"]) except Exception: reward = completion_quality_penalty(completion_text) rewards.append(reward) return rewards return reward_func def model_policy( model: Any, tokenizer: Any, *, max_seq_length: int, max_new_tokens: int = 512, ) -> Callable[[dict[str, Any]], dict[str, Any]]: if hasattr(FastLanguageModel, "for_inference"): FastLanguageModel.for_inference(model) def _policy(observation: dict[str, Any]) -> dict[str, Any]: prompt = clip_prompt( build_prompt(observation), tokenizer, max_prompt_tokens=max(256, max_seq_length - max_new_tokens), ) inputs = tokenizer( prompt, return_tensors="pt", truncation=True, max_length=max_seq_length - max_new_tokens, ).to(model.device) outputs = model.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=False, use_cache=True, ) completion = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True) try: return BDOAction.model_validate_json(extract_json_object(completion)).model_dump(mode="json", exclude_none=True) except Exception: return balanced_policy(observation) return _policy def evaluate_policy( policy_fn: Callable[[dict[str, Any]], dict[str, Any]], *, episodes: int, seed: int, ) -> list[dict[str, Any]]: scenarios = ["black_swan", "fraud_syndicate", "rapid_migration"] rows: list[dict[str, Any]] = [] for episode in range(episodes): scenario = scenarios[episode % len(scenarios)] env = BDOEnvironment(scenario=scenario, seed=seed + episode) observation = env.reset(seed=seed + episode, scenario=scenario) total_reward = 0.0 total_training_reward = 0.0 belief_scores: list[float] = [] while not observation.done: action_payload = policy_fn(observation.model_dump(mode="json", exclude_none=True)) observation = env.step(BDOAction.model_validate(action_payload)) info = observation.info or observation.metadata total_reward += float(observation.reward or 0.0) total_training_reward += float(info.get("training_reward", 0.0)) belief_scores.append(float(info["reward_breakdown"]["belief_accuracy"])) rows.append( { "episode": episode, "scenario": scenario, "total_reward": round(total_reward, 4), "total_training_reward": round(total_training_reward, 4), "avg_belief_accuracy": round(mean(belief_scores) if belief_scores else 0.0, 4), } ) return rows def write_reward_artifacts(results: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: artifacts_dir = Path("artifacts") artifacts_dir.mkdir(exist_ok=True) summary = {} for label, rows in results.items(): scenario_summary: dict[str, dict[str, float | int]] = {} scenarios = sorted({row["scenario"] for row in rows}) for scenario in scenarios: scenario_rows = [row for row in rows if row["scenario"] == scenario] scenario_summary[scenario] = { "episodes": len(scenario_rows), "mean_total_reward": round(mean(row["total_reward"] for row in scenario_rows), 4), "mean_total_training_reward": round( mean(row["total_training_reward"] for row in scenario_rows), 4 ), "mean_belief_accuracy": round( mean(row["avg_belief_accuracy"] for row in scenario_rows), 4 ), } summary[label] = { "episodes": rows, "mean_total_reward": round(mean(row["total_reward"] for row in rows), 4), "mean_total_training_reward": round(mean(row["total_training_reward"] for row in rows), 4), "mean_belief_accuracy": round(mean(row["avg_belief_accuracy"] for row in rows), 4), "scenario_summary": scenario_summary, } (artifacts_dir / "grpo_reward_curve.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") plt.figure(figsize=(10, 5)) for label, rows in results.items(): plt.plot([row["episode"] for row in rows], [row["total_reward"] for row in rows], label=label) plt.xlabel("evaluation episode") plt.ylabel("total environment reward") plt.title("BDO.ai GRPO Reward Comparison") plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() plt.savefig(artifacts_dir / "grpo_reward_curve.png", dpi=180) plt.close() labels = list(summary.keys()) mean_rewards = [summary[label]["mean_total_reward"] for label in labels] plt.figure(figsize=(7, 5)) plt.bar(labels, mean_rewards, color=["#4C78A8", "#F58518"][: len(labels)]) plt.ylabel("mean total environment reward") plt.title("BDO.ai GRPO Mean Reward") plt.grid(True, axis="y", alpha=0.3) plt.tight_layout() plt.savefig(artifacts_dir / "grpo_mean_reward.png", dpi=180) plt.close() return summary def write_train_metrics(log_history: list[dict[str, Any]]) -> None: rows = [] for row in log_history: if "step" in row and any(key in row for key in ("reward", "loss", "kl")): clean = {"step": int(row["step"])} for key in ("reward", "loss", "kl", "entropy"): if key in row: clean[key] = float(row[key]) rows.append(clean) Path("artifacts").mkdir(exist_ok=True) Path("artifacts/grpo_training_metrics.json").write_text(json.dumps(rows, indent=2), encoding="utf-8") if rows and any("reward" in row for row in rows): reward_rows = [row for row in rows if "reward" in row] plt.figure(figsize=(8, 5)) plt.plot([row["step"] for row in reward_rows], [row["reward"] for row in reward_rows]) plt.xlabel("training step") plt.ylabel("grpo reward") plt.title("BDO.ai GRPO Reward Curve") plt.grid(True, alpha=0.3) plt.tight_layout() plt.savefig("artifacts/grpo_training_curve.png", dpi=180) plt.close() def main() -> None: args = parse_args() random.seed(args.seed) print("Building environment-linked GRPO prompts...") dataset_rows = collect_grpo_prompts(dataset_episodes=args.dataset_episodes, seed=args.seed) train_dataset = Dataset.from_list(dataset_rows) print(f"Prepared {len(train_dataset)} GRPO prompts.") print(f"Loading model: {args.model_name}") model, tokenizer = FastLanguageModel.from_pretrained( model_name=args.model_name, max_seq_length=args.max_seq_length, dtype=None, load_in_4bit=True, ) configure_generation_defaults(model, tokenizer) print("Running pre-training evaluation...") untrained_policy = model_policy( model, tokenizer, max_seq_length=args.max_prompt_length + args.max_completion_length, max_new_tokens=args.max_completion_length, ) untrained_results = evaluate_policy(untrained_policy, episodes=args.eval_episodes, seed=args.seed + 100) model = FastLanguageModel.get_peft_model( model, r=16, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_alpha=16, use_gradient_checkpointing="unsloth", random_state=args.seed, ) trainer = GRPOTrainer( model=model, processing_class=tokenizer, reward_funcs=build_reward_function(args.seed + 900), args=GRPOConfig( output_dir=args.output_dir, learning_rate=args.learning_rate, per_device_train_batch_size=args.batch_size, gradient_accumulation_steps=args.grad_accum, num_generations=args.num_generations, max_prompt_length=args.max_prompt_length, max_completion_length=args.max_completion_length, max_steps=args.steps, logging_steps=1, save_strategy="no", report_to=[], bf16=bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported()), fp16=not bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported()), ), train_dataset=train_dataset, ) print("Starting GRPO training...") train_result = trainer.train() adapter_dir = Path("artifacts/grpo_model") adapter_dir.mkdir(parents=True, exist_ok=True) model.save_pretrained(adapter_dir) tokenizer.save_pretrained(adapter_dir) write_train_metrics(trainer.state.log_history) # Upload adapter to HF Space so the UI can load it on restart hf_token = os.environ.get("HF_TOKEN") space_id = os.environ.get("SPACE_ID") # set automatically inside HF jobs if hf_token and space_id: try: from huggingface_hub import HfApi print(f"Uploading adapter to Space {space_id} ...") HfApi(token=hf_token).upload_folder( folder_path=str(adapter_dir), repo_id=space_id, repo_type="space", path_in_repo="artifacts/grpo_model", ) print("Adapter uploaded. UI will use LLM policy after Space restart.") except Exception as _upload_err: print(f"Warning: adapter upload failed ({_upload_err}). Weights saved locally only.") print("Running post-training evaluation...") trained_policy = model_policy( model, tokenizer, max_seq_length=args.max_prompt_length + args.max_completion_length, max_new_tokens=args.max_completion_length, ) reward_results = { "untrained_qwen_base": untrained_results, "trained_qwen_grpo": evaluate_policy(trained_policy, episodes=args.eval_episodes, seed=args.seed + 300), } reward_summary = write_reward_artifacts(reward_results) summary = { "model_name": args.model_name, "dataset_prompts": len(train_dataset), "dataset_episodes": args.dataset_episodes, "eval_episodes": args.eval_episodes, "grpo_steps": args.steps, "train_runtime_seconds": round(float(train_result.metrics.get("train_runtime", 0.0)), 2), "reward_summary": reward_summary, "artifacts": { "grpo_training_metrics_json": "artifacts/grpo_training_metrics.json", "grpo_training_curve_png": "artifacts/grpo_training_curve.png", "grpo_reward_curve_json": "artifacts/grpo_reward_curve.json", "grpo_reward_curve_png": "artifacts/grpo_reward_curve.png", "grpo_mean_reward_png": "artifacts/grpo_mean_reward.png", "grpo_summary_json": "artifacts/grpo_training_summary.json", "adapter_dir": str(adapter_dir), }, } Path("artifacts/grpo_training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") print(json.dumps(summary, indent=2)) if __name__ == "__main__": main()