#!/usr/bin/env python # train_battleground_rlaif_gamehistory.py # # SFT + GRPO (RLAIF style) on Hearthstone Battlegrounds "game_history" data. # # Expected data format per JSON file (per game): # { # "game_metadata": {...}, # "turns": [ # { # "turn_number": 0, # "phase": "PlayerTurn", # "state": { # nested game_state / player_hero / resources / board_state # "game_state": {...}, # "player_hero": {...}, # "resources": {...}, # "board_state": {...} # }, # "candidates": [ # RLAIF annotations you add # {"role": "expert", "actions": [{...}, {...}], "reward": 1.0}, # {"role": "medium", "actions": [{...}], "reward": 0.0}, # {"role": "bad", "actions": [{...}], "reward": -0.5} # ], # ... other fields like battle_result, reward, etc. ... # }, # ... # ] # } # # Each candidate's "actions" field is a SEQUENCE (list) of atomic Battlegrounds # actions, where each atomic action dict uses the schema from the original # RLAIF pipeline: # { # "type": "BUY_FROM_TAVERN" | "PLAY_FROM_HAND" | "SELL_FROM_BOARD" | # "HERO_POWER" | "ROLL" | "UPGRADE_TAVERN" | "FREEZE" | "END_TURN", # "tavern_index": int or null, # "hand_index": int or null, # "board_index": int or null, # "card_name": string or null # } # # The loader flattens all labeled turns (those with "candidates") into per-step # records while preserving the nested "state" structure. import argparse import json import os import sys from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Dict, Any import torch from datasets import Dataset from transformers import AutoTokenizer, AutoModelForCausalLM from peft import LoraConfig from trl import SFTTrainer, SFTConfig, GRPOTrainer, GRPOConfig _SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) if _SCRIPT_DIR not in sys.path: sys.path.append(_SCRIPT_DIR) from battleground_nl_utils import ( game_state_to_natural_language, ) # ================== Model paths & defaults ================== LOCAL_INSTRUCT_PATH = "models/qwen3-4b-instruct-2507/Qwen/Qwen3-4B-Instruct-2507" def _resolve_default_model_id() -> str: env_override = os.environ.get("QWEN_INSTRUCT_MODEL") if env_override: return env_override if os.path.isdir(LOCAL_INSTRUCT_PATH): return LOCAL_INSTRUCT_PATH return "Qwen/Qwen3-4B-Instruct" DEFAULT_MODEL_ID = _resolve_default_model_id() DEFAULT_OUTPUT_DIR = "./battleground_rlaif_qwen_gamehistory" # By default, point to a single game-history style file. You can override # with a directory containing many such JSONs. DEFAULT_DATA_FILE = "RL/datasets/game_history_fixed.json" DEFAULT_TARGET_MODULES = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] # ================== Config dataclass ================== @dataclass class PipelineConfig: model_name_or_path: str = DEFAULT_MODEL_ID output_dir: str = DEFAULT_OUTPUT_DIR data_file: str = DEFAULT_DATA_FILE input_mode: str = "json" # "json" uses nested game_history state; "nl" uses natural language max_seq_length: int = 1024 sft_epochs: int = 3 grpo_epochs: int = 3 bf16: bool = True per_device_batch_size: int = 4 grad_accum_steps: int = 4 sft_learning_rate: float = 1e-5 grpo_learning_rate: float = 5e-6 max_completion_length: int = 128 num_generations: int = 3 steps_per_generation: int = 1 target_modules: Optional[List[str]] = None skip_sft: bool = False skip_grpo: bool = False train_on_all_data: bool = False def parse_args() -> PipelineConfig: parser = argparse.ArgumentParser( description="Run SFT + GRPO (RLAIF) on Battlegrounds game_history dataset." ) parser.add_argument( "--model", default=DEFAULT_MODEL_ID, help="Model id or local path for the Qwen instruct checkpoint.", ) parser.add_argument( "--output-dir", default=DEFAULT_OUTPUT_DIR, help="Directory for checkpoints and logs.", ) parser.add_argument( "--data-file", default=DEFAULT_DATA_FILE, help=( "Path to a game_history-style JSON file or a directory of such files. " "Each file should have {game_metadata, turns[...]} and each labeled turn " "must contain a 'candidates' list." ), ) parser.add_argument( "--input-mode", choices=["json", "nl"], default="json", help=( "Input format for game state: 'json' uses nested game_history JSON; " "'nl' converts the nested state to natural language." ), ) parser.add_argument("--max-seq-length", type=int, default=1024) parser.add_argument("--sft-epochs", type=int, default=20) parser.add_argument("--grpo-epochs", type=int, default=3) parser.add_argument( "--per-device-batch-size", type=int, default=4, help="Batch size per device (default: 4 for A800 80GB)", ) parser.add_argument("--grad-accum-steps", type=int, default=4) parser.add_argument("--sft-learning-rate", type=float, default=1e-5) parser.add_argument("--grpo-learning-rate", type=float, default=5e-6) parser.add_argument("--max-completion-length", type=int, default=128) parser.add_argument("--num-generations", type=int, default=3) parser.add_argument( "--target-modules", default=None, help="Comma-separated list of module names for LoRA (defaults to Qwen attn/FFN blocks).", ) parser.add_argument( "--disable-bf16", action="store_true", help="Force fp16/fp32 training if bf16 is not desired or unsupported.", ) parser.add_argument("--skip-sft", action="store_true", help="Skip the SFT phase.") parser.add_argument("--skip-grpo", action="store_true", help="Skip the GRPO phase.") parser.add_argument( "--train-on-all-data", action="store_true", help="Use all rows as training data (no hold-out split); SFT eval runs on the same data.", ) args = parser.parse_args() target_modules = ( [m.strip() for m in args.target_modules.split(",") if m.strip()] if args.target_modules else None ) return PipelineConfig( model_name_or_path=args.model, output_dir=args.output_dir, data_file=args.data_file, input_mode=args.input_mode, max_seq_length=args.max_seq_length, sft_epochs=args.sft_epochs, grpo_epochs=args.grpo_epochs, bf16=not args.disable_bf16, per_device_batch_size=args.per_device_batch_size, grad_accum_steps=args.grad_accum_steps, sft_learning_rate=args.sft_learning_rate, grpo_learning_rate=args.grpo_learning_rate, max_completion_length=args.max_completion_length, num_generations=args.num_generations, target_modules=target_modules, skip_sft=args.skip_sft, skip_grpo=args.skip_grpo, train_on_all_data=args.train_on_all_data, ) # ================== Data: Battlegrounds formatting ================== INSTRUCTION_PREFIX = """You are a Hearthstone Battlegrounds AI. Given the current game state as a JSON object, choose the best full-turn sequence of actions and respond with a single JSON object in this exact format: {"actions":[{"type":"","tavern_index":,"hand_index":,"board_index":,"card_name":}, ...]} Rules: 1. Respond with JSON only. Do not add explanations or any extra text. 2. The top-level object must have exactly one key: "actions". 3. "actions" must be a JSON array (possibly empty, but usually 1+ steps) of atomic action objects. 4. Use 0-based integers for indices or null when not used. 5. "type" must be one of: "BUY_FROM_TAVERN","PLAY_FROM_HAND","SELL_FROM_BOARD", "HERO_POWER","ROLL","UPGRADE_TAVERN","FREEZE","END_TURN". 6. "card_name" must exactly match a card name from the game state when required, otherwise null. Now here is the game state JSON: """ INSTRUCTION_PREFIX_NL = """You are a Hearthstone Battlegrounds AI. Given the following natural language description of the current game state, choose the best full-turn sequence of actions and respond with a single JSON object in this exact format: {"actions":[{"type":"","tavern_index":,"hand_index":,"board_index":,"card_name":}, ...]} Rules: 1. Respond with JSON only. Do not add explanations or any extra text. 2. The top-level object must have exactly one key: "actions". 3. "actions" must be a JSON array (possibly empty, but usually 1+ steps) of atomic action objects. 4. Use 0-based integers for indices or null when not used. 5. "type" must be one of: "BUY_FROM_TAVERN","PLAY_FROM_HAND","SELL_FROM_BOARD", "HERO_POWER","ROLL","UPGRADE_TAVERN","FREEZE","END_TURN". 6. "card_name" must exactly match a card name from the game state when required, otherwise null. Now here is the description of the game state: """ def _build_prompt(example: Dict[str, Any], input_mode: str = "json") -> str: """Build a prompt from a flattened game_history example. The example has: - phase: string (e.g., "PlayerTurn") - turn: int - state: nested dict with keys: game_state, player_hero, resources, board_state """ state = example.get("state", {}) if input_mode == "nl": # state is already in the game_state / player_hero / resources / board_state shape nl_state = game_state_to_natural_language(state) prefix = INSTRUCTION_PREFIX_NL state_text = nl_state else: # JSON mode: wrap the nested state in a small task object. gs = state.get("game_state", {}) or {} phase = example.get("phase", gs.get("phase", "Unknown")) turn = example.get("turn", gs.get("turn_number", 0)) obj = { "task": "battlegrounds_policy_v1", "phase": phase, "turn": turn, "state": state, } state_text = json.dumps(obj, separators=(",", ":"), ensure_ascii=False) prefix = INSTRUCTION_PREFIX return prefix + "\n" + state_text def _build_completion_from_actions(actions: List[Dict[str, Any]]) -> str: """Pack a sequence of atomic actions into the expected JSON completion. {"actions": [ {...}, {...}, ... ]} """ return json.dumps({"actions": actions}, separators=(",", ":"), ensure_ascii=False) def load_gamehistory_rlaif( data_file: str, test_size: float = 0.1, seed: int = 42, train_on_all_data: bool = False, input_mode: str = "json", ): """Load game_history-style JSON data and build SFT & RL datasets. - data_file can be: * a single JSON file with {game_metadata, turns: [...]} structure; * a JSON file containing a list of such game objects; * a directory containing multiple .json files in either of the above forms. - Each labeled turn must contain a "candidates" list; turns without candidates are skipped. """ path = Path(data_file) if not path.exists(): raise FileNotFoundError(f"Data file or directory not found: {data_file}") rows: List[Dict[str, Any]] = [] def _consume_game_obj(game_obj: Dict[str, Any], game_id_hint: str) -> None: meta = game_obj.get("game_metadata", {}) or {} turns = game_obj.get("turns", []) or [] for t in turns: state = t.get("state", {}) or {} candidates = t.get("candidates") if not candidates: # Skip unlabeled turns (no RLAIF annotations yet) continue gs = state.get("game_state", {}) or {} phase = t.get("phase") or gs.get("phase", "PlayerTurn") turn = gs.get("turn_number", t.get("turn_number", 0)) row_meta = { "game_metadata": meta, "battle_result": t.get("battle_result"), "health_before_battle": t.get("health_before_battle"), "health_after_battle": t.get("health_after_battle"), "health_change": t.get("health_change"), "action_taken": t.get("action_taken"), } rows.append( { "game_id": meta.get("game_id") or game_id_hint, "step_id": t.get("turn_number", turn), "turn": turn, "phase": phase, "state": state, "candidates": candidates, "meta": row_meta, } ) def _load_one_json_file(p: Path) -> None: with p.open("r", encoding="utf-8") as f: data = json.load(f) # Case 1: single game_history object with turns if isinstance(data, dict) and "turns" in data: _consume_game_obj(data, game_id_hint=p.stem) # Case 2: already-flattened per-turn rows in a list elif isinstance(data, list) and data and isinstance(data[0], dict) and "state" in data[0]: for idx, row in enumerate(data): if not isinstance(row, dict): raise ValueError( f"Unsupported JSON row at index {idx} in file {p}: expected dict with 'state'." ) candidates = row.get("candidates") if not candidates: # Skip unlabeled rows (no RLAIF annotations yet) continue state = row.get("state", {}) or {} gs = state.get("game_state", {}) or {} if "phase" not in row: row["phase"] = gs.get("phase", "PlayerTurn") if "turn" not in row: row["turn"] = gs.get("turn_number", row.get("step_id", 0)) # Ensure at least the keys expected downstream are present; keep any # extra metadata fields as-is. rows.append(row) # Case 3: list of game_history objects with turns elif isinstance(data, list): for idx, item in enumerate(data): if isinstance(item, dict) and "turns" in item: game_id_hint = item.get("game_metadata", {}).get("game_id") or f"{p.stem}_{idx}" _consume_game_obj(item, game_id_hint=game_id_hint) else: raise ValueError( f"Unsupported JSON object at index {idx} in file {p}: expected game_history with 'turns' or flat rows with 'state'." ) else: raise ValueError( f"Unsupported JSON structure in file {p}: expected dict with 'turns', list of such dicts, or list of flat rows with 'state'." ) if path.is_dir(): json_files = sorted(path.glob("*.json")) if not json_files: raise ValueError(f"No .json files found in directory: {data_file}") for p in json_files: _load_one_json_file(p) else: _load_one_json_file(path) if not rows: raise ValueError( "No labeled turns (with 'candidates') were found in the provided data. " "Make sure each turn you want to train on has a non-empty 'candidates' list." ) raw = Dataset.from_list(rows) # Train / eval split if train_on_all_data: raw_train = raw raw_eval = raw else: split = raw.train_test_split(test_size=test_size, seed=seed) raw_train = split["train"] raw_eval = split["test"] def to_sft(example: Dict[str, Any]) -> Dict[str, Any]: # Pick the expert candidate; if not present, fall back to max reward. candidates = example["candidates"] expert = None for c in candidates: if c.get("role") == "expert": expert = c break if expert is None: expert = max(candidates, key=lambda x: float(x.get("reward", 0.0))) prompt = _build_prompt(example, input_mode=input_mode) # In the game_history pipeline, each candidate carries a SEQUENCE of # atomic actions under the "actions" key. completion = _build_completion_from_actions(expert["actions"]) return { "prompt": prompt, "completion": completion, } def to_rl(example: Dict[str, Any]) -> Dict[str, Any]: prompt = _build_prompt(example, input_mode=input_mode) return { "prompt": prompt, "candidates": example["candidates"], } sft_train = raw_train.map(to_sft, remove_columns=raw_train.column_names) sft_eval = raw_eval.map(to_sft, remove_columns=raw_eval.column_names) rl_train = raw_train.map(to_rl, remove_columns=raw_train.column_names) return sft_train, sft_eval, rl_train # ================== Reward function for GRPO (RLAIF style) ================== def _parse_actions_from_completion(text: str) -> Optional[List[Dict[str, Any]]]: """Parse a model completion into a sequence of atomic action dicts. Expected formats: - {"actions": [ {...}, {...}, ... ]} - {"action": [ {...}, {...}, ... ]} # tolerated fallback """ text = text.strip() # Try to locate a JSON object within the text (in case of extra chatter # before/after the JSON), similar to the eval-time parser. start_idx = text.find("{") if start_idx == -1: return None end_idx = text.rfind("}") if end_idx == -1: return None json_str = text[start_idx : end_idx + 1] try: obj = json.loads(json_str) except Exception: return None if not isinstance(obj, dict): return None seq = None # Preferred key from the instruction if "actions" in obj: if isinstance(obj["actions"], list): seq = obj["actions"] elif isinstance(obj["actions"], dict): # Tolerate a single dict instead of a list seq = [obj["actions"]] # Fallback key for older/variant outputs elif "action" in obj: if isinstance(obj["action"], list): seq = obj["action"] elif isinstance(obj["action"], dict): seq = [obj["action"]] if seq is None: return None # Ensure we have a list of dicts. actions: List[Dict[str, Any]] = [] for step in seq: if not isinstance(step, dict): return None actions.append(step) return actions def _action_sequences_equal( a: List[Dict[str, Any]], b: List[Dict[str, Any]] ) -> bool: """Strict equality for sequences of atomic actions. Both length and each per-step dict must match exactly. This relies on a canonical action representation in the data and model outputs. """ if len(a) != len(b): return False for s1, s2 in zip(a, b): if s1 != s2: return False return True def battleground_rlaif_reward( completions: List[str], candidates: List[List[Dict[str, Any]]], **kwargs, ) -> List[float]: """RLAIF-style reward function for GRPOTrainer. For each completion (one JSON text): 1. Parse into a sequence of atomic actions. 2. Compare with the example's candidates[i].actions. 3. If it exactly matches a candidate.actions sequence, return that candidate's reward. 4. Otherwise reward = 0.0. """ rewards: List[float] = [] for comp_text, cand_list in zip(completions, candidates): seq = _parse_actions_from_completion(comp_text) if seq is None: rewards.append(0.0) continue best_reward = 0.0 for cand in cand_list: cand_actions = cand.get("actions") if not isinstance(cand_actions, list): continue if _action_sequences_equal(seq, cand_actions): r = float(cand.get("reward", 0.0)) if r > best_reward: best_reward = r rewards.append(best_reward) return rewards # ================== SFT phase ================== def run_sft(train_ds, eval_ds, tokenizer, cfg: PipelineConfig): """Run a supervised fine-tuning pass with LoRA adapters (prompt→action JSON).""" target_modules = cfg.target_modules or DEFAULT_TARGET_MODULES peft_config = LoraConfig( r=16, lora_alpha=32, lora_dropout=0.05, bias="none", target_modules=target_modules, task_type="CAUSAL_LM", ) sft_config = SFTConfig( output_dir=os.path.join(cfg.output_dir, "sft"), per_device_train_batch_size=cfg.per_device_batch_size, per_device_eval_batch_size=cfg.per_device_batch_size, gradient_accumulation_steps=cfg.grad_accum_steps, learning_rate=cfg.sft_learning_rate, num_train_epochs=cfg.sft_epochs, logging_steps=10, save_steps=200, eval_steps=200, eval_strategy="steps", save_total_limit=2, max_length=cfg.max_seq_length, bf16=cfg.bf16, fp16=not cfg.bf16, report_to=["none"], ) trainer = SFTTrainer( model=cfg.model_name_or_path, # model id / path; SFTTrainer loads it args=sft_config, train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tokenizer, peft_config=peft_config, ) trainer.train() save_path = os.path.join(cfg.output_dir, "sft_model") trainer.save_model(save_path) return trainer.model # PEFT-wrapped model instance # ================== GRPO phase ================== def run_grpo(rl_dataset, base_model, tokenizer, cfg: PipelineConfig): """Run a GRPO RLAIF loop on top of the (optionally) SFT-initialized model.""" target_modules = cfg.target_modules or DEFAULT_TARGET_MODULES if hasattr(base_model, "peft_config"): peft_config = None else: peft_config = LoraConfig( r=8, lora_alpha=16, lora_dropout=0.05, bias="none", target_modules=target_modules, task_type="CAUSAL_LM", ) generation_batch_size = cfg.per_device_batch_size * cfg.num_generations grpo_config = GRPOConfig( output_dir=os.path.join(cfg.output_dir, "grpo"), num_train_epochs=cfg.grpo_epochs, per_device_train_batch_size=cfg.per_device_batch_size, gradient_accumulation_steps=cfg.grad_accum_steps, logging_steps=10, save_strategy="epoch", save_total_limit=cfg.grpo_epochs, bf16=cfg.bf16, fp16=not cfg.bf16, learning_rate=cfg.grpo_learning_rate, max_prompt_length=cfg.max_seq_length, max_completion_length=cfg.max_completion_length, num_generations=cfg.num_generations, generation_batch_size=generation_batch_size, report_to=["none"], ) if peft_config is not None: trainer = GRPOTrainer( model=base_model, args=grpo_config, processing_class=tokenizer, reward_funcs=battleground_rlaif_reward, train_dataset=rl_dataset, peft_config=peft_config, ) else: trainer = GRPOTrainer( model=base_model, args=grpo_config, processing_class=tokenizer, reward_funcs=battleground_rlaif_reward, train_dataset=rl_dataset, ) trainer.train() trainer.save_model(os.path.join(cfg.output_dir, "grpo_model")) # ================== Main ================== def main(): cfg = parse_args() os.makedirs(cfg.output_dir, exist_ok=True) print(f"Using model: {cfg.model_name_or_path}") print("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained( cfg.model_name_or_path, use_fast=True, trust_remote_code=True, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # For GRPO, we want left padding tokenizer.padding_side = "left" print(f"Loading Battlegrounds game_history dataset from: {cfg.data_file}") sft_train, sft_eval, rl_train = load_gamehistory_rlaif( cfg.data_file, train_on_all_data=cfg.train_on_all_data, input_mode=cfg.input_mode, ) # ----- SFT ----- if cfg.skip_sft: print("Skipping SFT phase; loading base model directly.") dtype = ( torch.bfloat16 if cfg.bf16 and torch.cuda.is_available() else (torch.float16 if torch.cuda.is_available() else torch.float32) ) model_kwargs: Dict[str, Any] = { "torch_dtype": dtype, "trust_remote_code": True, } if torch.cuda.is_available(): model_kwargs["device_map"] = "auto" base_model = AutoModelForCausalLM.from_pretrained( cfg.model_name_or_path, **model_kwargs ) else: print("Running SFT phase...") base_model = run_sft(sft_train, sft_eval, tokenizer, cfg) # ----- GRPO ----- if cfg.skip_grpo: print("Skipping GRPO phase; only SFT outputs (if any) were produced.") else: print("Running GRPO (RLAIF) phase...") run_grpo(rl_train, base_model, tokenizer, cfg) print("All done. Check outputs under:", cfg.output_dir) if __name__ == "__main__": main()