training/scripts/kaggle_rl_1p5b_methods.py

#!/usr/bin/env python3
"""
Kaggle GPU RL trainer for AutoDataLab++ CoS routing on Qwen2.5-1.5B.

This script is intentionally practical for hackathon iteration:

Methods
-------
  grpo
      Group-relative preference optimization over candidate JSON actions.
      For each environment state, score multiple candidate actions, normalize
      rewards within the group, and update LoRA weights toward high-advantage
      actions.

  grpo_rlvr
      Same GRPO update, but the reward is "verifiable" from the environment:
      correct next oracle action + strict environment step reward + penalties
      for premature summarize/submit/repeats. This is the best RLVR-style
      option for our setup.

  ppo
      Lightweight PPO-style clipped policy update over candidate JSON actions.
      It uses old/new action log-prob ratios and clipped advantages. This is
      not a full value-head PPO trainer; it is a compact bandit-PPO variant
      tailored for CoS action routing.

Recommended flow
----------------
1. Train SFT/DPO first with:
     training/kaggle_train_1p5b_methods.py --method sft_then_dpo ...

2. Then run RL from that adapter:
     !python3 training/kaggle_rl_1p5b_methods.py \\
       --method grpo_rlvr \\
       --init-adapter /kaggle/working/cos_1p5b_runs/qwen15b_sft_then_dpo_v1/adapter \\
       --run-name qwen15b_grpo_rlvr_v1

3. Check:
     /kaggle/working/cos_1p5b_rl_runs/<run-name>/eval/evidence.md

If trajectory is:
  analyst -> finance -> strategy -> hr -> summarize -> submit
with fallback=False, keep it. Otherwise skip and try another method/seed.
"""
from __future__ import annotations

import argparse
import gc
import json
import math
import os
import random
import re
import shutil
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

REPO = Path(__file__).resolve().parents[2]
if str(REPO) not in sys.path:
    sys.path.insert(0, str(REPO))

from ceo_brief_env.environment import CEOBriefEnvironment, oracle_action_for_observation, required_experts_for_task
from ceo_brief_env.models import CoSAction, CoSObservation

TASKS = ["easy_brief", "medium_brief", "hard_brief", "expert_brief", "risk_brief", "crisis_brief"]
VALID_ACTIONS = {"consult", "ask", "summarize", "submit", "noop"}
VALID_EXPERTS = {"analyst", "finance", "hr", "strategy"}
JSON_RE = re.compile(r"\{[^{}]*\}", re.S)

SYSTEM_PROMPT = (
    "You are the Chief of Staff in AutoDataLab++. You orchestrate four specialists: "
    "analyst, finance, strategy, hr. Reply with STRICT JSON only.\n"
    'Schema: {"action_type": one of [consult, ask, summarize, submit, noop], '
    '"expert_id": one of [analyst, finance, hr, strategy] or null}.\n'
    "Rules: consult each required expert exactly once when required, then summarize, then submit. "
    "Never summarize while required experts are missing. Never repeat summarize."
)


def set_seed(seed: int) -> None:
    random.seed(seed)
    try:
        import numpy as np
        import torch

        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except Exception:
        pass


def action_json(action: CoSAction) -> str:
    return json.dumps(action.model_dump(exclude_none=True), separators=(",", ":"), sort_keys=True)


def action_label(action: CoSAction) -> str:
    if action.action_type in {"consult", "ask"}:
        return f"{action.action_type}:{action.expert_id or 'null'}"
    return action.action_type


def parse_action(text: str) -> CoSAction:
    m = JSON_RE.search(text or "")
    if not m:
        return CoSAction(action_type="noop")
    try:
        payload = json.loads(m.group(0))
    except Exception:
        return CoSAction(action_type="noop")
    action_type = payload.get("action_type")
    if action_type not in VALID_ACTIONS:
        return CoSAction(action_type="noop")
    expert_id = payload.get("expert_id")
    if expert_id is not None and expert_id not in VALID_EXPERTS:
        expert_id = None
    return CoSAction(action_type=action_type, expert_id=expert_id)


def render_obs(obs: CoSObservation, variant: int = 0) -> str:
    required = required_experts_for_task(obs.task_name)
    missing = [e for e in required if e not in obs.consulted_experts]
    if variant == 0:
        return json.dumps(
            {
                "task": obs.task_name,
                "step": obs.step_count,
                "max_steps": obs.max_steps,
                "required_experts": required,
                "consulted_experts": list(obs.consulted_experts),
                "missing_required_experts": missing,
                "brief_ready": obs.current_brief is not None,
                "rag_enabled": obs.rag_enabled,
            },
            separators=(",", ":"),
        )
    return (
        f"task={obs.task_name}; step={obs.step_count}/{obs.max_steps}; "
        f"required={required}; consulted={obs.consulted_experts}; missing={missing}; "
        f"brief_ready={obs.current_brief is not None}; return next strict JSON action."
    )


def format_chat(tokenizer, messages: list[dict[str, str]], add_generation_prompt: bool = False) -> str:
    if hasattr(tokenizer, "apply_chat_template"):
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt)
    text = ""
    for m in messages:
        text += f"{m['role'].upper()}:\n{m['content']}\n"
    if add_generation_prompt:
        text += "ASSISTANT:\n"
    return text


def prompt_for_obs(tokenizer, obs: CoSObservation, variant: int = 0) -> str:
    return format_chat(
        tokenizer,
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": render_obs(obs, variant=variant)},
        ],
        add_generation_prompt=True,
    )


def candidate_actions(obs: CoSObservation) -> list[CoSAction]:
    """Fixed action set for candidate-action RL."""
    required = required_experts_for_task(obs.task_name)
    out: list[CoSAction] = [CoSAction(action_type="consult", expert_id=e) for e in ["analyst", "finance", "strategy", "hr"]]
    out += [
        CoSAction(action_type="summarize"),
        CoSAction(action_type="submit"),
        CoSAction(action_type="noop"),
    ]
    # Keep action set stable, but put oracle-ish actions first for easier inspection.
    oracle = oracle_action_for_observation(obs)
    oracle_s = action_json(oracle)
    ordered = [oracle] + [a for a in out if action_json(a) != oracle_s]
    # If task does not require strategy, strategy remains a candidate but will be
    # lower reward; useful for learning "don't over-consult".
    return ordered


def oracle_reward(obs: CoSObservation, action: CoSAction) -> float:
    oracle = oracle_action_for_observation(obs)
    if action.model_dump(exclude_none=True) == oracle.model_dump(exclude_none=True):
        return 1.0
    required = required_experts_for_task(obs.task_name)
    missing = [e for e in required if e not in obs.consulted_experts]
    if action.action_type == "consult" and action.expert_id in missing:
        return 0.65
    if missing and action.action_type in {"summarize", "submit"}:
        return -1.0
    if action.action_type == "noop":
        return -0.8
    if action.action_type == "consult" and action.expert_id in obs.consulted_experts:
        return -0.6
    return -0.2


def env_step_reward(obs: CoSObservation, action: CoSAction) -> float:
    """Verifiable one-step reward from strict env, without auto-fill."""
    env = CEOBriefEnvironment(shaping="strict", auto_fill_required=False)
    sim = env.reset(task=obs.task_name, use_rag=obs.rag_enabled)
    # Recreate state approximately by replaying history. Histories are action JSONs.
    for h in obs.history:
        try:
            payload = json.loads(h)
            sim = env.step(CoSAction.model_validate(payload))
        except Exception:
            pass
    sim = env.step(action)
    return float(sim.reward)


def rlvr_reward(obs: CoSObservation, action: CoSAction) -> float:
    """Combined RLVR reward: oracle correctness + strict env signal."""
    return oracle_reward(obs, action) + 2.0 * env_step_reward(obs, action)


def collect_states(tasks: list[str], rag_modes: list[bool], variants: int) -> list[dict[str, Any]]:
    states: list[dict[str, Any]] = []
    for task in tasks:
        for use_rag in rag_modes:
            env = CEOBriefEnvironment(auto_fill_required=False)
            obs = env.reset(task=task, use_rag=use_rag)
            while not obs.done and obs.step_count < obs.max_steps:
                for variant in range(variants):
                    states.append(
                        {
                            "task": task,
                            "use_rag": use_rag,
                            "variant": variant,
                            "obs": obs,
                        }
                    )
                obs = env.step(oracle_action_for_observation(obs))
    random.shuffle(states)
    return states


@dataclass
class Loaded:
    tokenizer: Any
    model: Any


def load_policy_model(args) -> Loaded:
    import torch
    from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    tok_source = args.init_adapter if args.init_adapter else args.model_id
    tok = AutoTokenizer.from_pretrained(tok_source, token=args.hf_token or None)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "right"

    bnb = None
    if not args.no_4bit:
        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        token=args.hf_token or None,
        device_map="auto",
        quantization_config=bnb,
        torch_dtype=torch.float16,
    )
    model.resize_token_embeddings(len(tok))
    if not args.no_4bit:
        model = prepare_model_for_kbit_training(model)
    if args.init_adapter:
        model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True)
    else:
        cfg = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        )
        model = get_peft_model(model, cfg)
    model.config.use_cache = False
    return Loaded(tok, model)


def sequence_logprob(model, tokenizer, prompt: str, completion: str, max_length: int):
    """Mean log-prob of completion tokens under CausalLM."""
    import torch

    full = prompt + completion + (tokenizer.eos_token or "")
    enc = tokenizer(full, return_tensors="pt", truncation=True, max_length=max_length).to(model.device)
    prompt_ids = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length).to(model.device)
    input_ids = enc["input_ids"]
    labels = input_ids.clone()
    prompt_len = min(prompt_ids["input_ids"].shape[1], labels.shape[1])
    labels[:, :prompt_len] = -100
    out = model(**enc)
    logits = out.logits[:, :-1, :]
    target = labels[:, 1:]
    mask = target.ne(-100)
    logp = torch.log_softmax(logits, dim=-1)
    safe_target = target.masked_fill(~mask, 0)
    tok_logp = logp.gather(-1, safe_target.unsqueeze(-1)).squeeze(-1)
    return (tok_logp * mask).sum() / mask.sum().clamp_min(1)


def train_candidate_rl(args, out_dir: Path) -> Path:
    import torch

    loaded = load_policy_model(args)
    model = loaded.model
    tok = loaded.tokenizer
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    states = collect_states(TASKS, [False, True], args.variants)
    if args.max_train_states:
        states = states[: args.max_train_states]
    print(f"[data] states={len(states)} method={args.method}", flush=True)

    metrics: list[dict[str, float]] = []
    tb_writer = None
    wandb_run = None
    if args.report_to == "tensorboard":
        try:
            from torch.utils.tensorboard import SummaryWriter

            tb_writer = SummaryWriter(log_dir=str(out_dir / "tb_logs"))
            print(f"[tracking] tensorboard logs -> {out_dir / 'tb_logs'}", flush=True)
        except Exception as e:
            print(f"[tracking] tensorboard unavailable: {e}", flush=True)
    elif args.report_to == "wandb":
        try:
            import wandb

            wandb_run = wandb.init(
                project=args.wandb_project,
                name=args.run_name or f"{args.method}_qwen15b",
                config={k: str(v) for k, v in vars(args).items()},
            )
            print(f"[tracking] wandb run -> {wandb_run.url}", flush=True)
        except Exception as e:
            print(f"[tracking] wandb unavailable: {e}", flush=True)
    global_step = 0
    for epoch in range(args.epochs):
        random.shuffle(states)
        for i, item in enumerate(states, start=1):
            obs: CoSObservation = item["obs"]
            prompt = prompt_for_obs(tok, obs, variant=int(item["variant"]))
            cands = candidate_actions(obs)
            rewards = []
            for action in cands:
                r = rlvr_reward(obs, action) if args.method == "grpo_rlvr" else oracle_reward(obs, action)
                rewards.append(float(r))
            rewards_t = torch.tensor(rewards, dtype=torch.float32, device=model.device)

            logps = []
            with torch.no_grad():
                old_logps = []
                for action in cands:
                    old_logps.append(sequence_logprob(model, tok, prompt, action_json(action), args.max_length).detach())
                old_logps_t = torch.stack(old_logps)

            for action in cands:
                logps.append(sequence_logprob(model, tok, prompt, action_json(action), args.max_length))
            logps_t = torch.stack(logps)

            if args.method in {"grpo", "grpo_rlvr"}:
                adv = (rewards_t - rewards_t.mean()) / rewards_t.std().clamp_min(1e-4)
                # Group-relative update: improve high-advantage candidates and
                # suppress low-advantage candidates. The detached log-softmax
                # normalization keeps this stable on tiny candidate groups.
                log_policy = torch.log_softmax(logps_t, dim=0)
                loss = -(adv.detach() * log_policy).mean()
            else:
                # PPO-style clipped update over candidate actions.
                adv = (rewards_t - rewards_t.mean()) / rewards_t.std().clamp_min(1e-4)
                ratio = torch.exp(logps_t - old_logps_t)
                clipped = torch.clamp(ratio, 1.0 - args.ppo_clip, 1.0 + args.ppo_clip)
                loss = -torch.minimum(ratio * adv.detach(), clipped * adv.detach()).mean()

            if args.sft_anchor > 0:
                oracle = oracle_action_for_observation(obs)
                oracle_s = action_json(oracle)
                oracle_logp = sequence_logprob(model, tok, prompt, oracle_s, args.max_length)
                # Positive anchor: keep the known-good routing behavior while
                # RL nudges preferences. This is critical when continuing from
                # an already-good SFT/DPO adapter.
                loss = loss - args.sft_anchor * oracle_logp

            loss = loss / args.grad_accum
            loss.backward()
            if i % args.grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1
                if global_step % args.log_every == 0:
                    best_idx = int(torch.argmax(logps_t.detach()).item())
                    reward_best = float(rewards_t[best_idx].item())
                    oracle_idx = int(torch.argmax(rewards_t).item())
                    chosen_ok = best_idx == oracle_idx
                    row = {
                        "step": global_step,
                        "epoch": float(epoch),
                        "loss": float(loss.detach().cpu().item() * args.grad_accum),
                        "mean_reward": float(rewards_t.mean().detach().cpu().item()),
                        "best_reward": reward_best,
                        "chosen_ok": float(chosen_ok),
                    }
                    metrics.append(row)
                    if tb_writer is not None:
                        for key in ("loss", "mean_reward", "best_reward", "chosen_ok"):
                            tb_writer.add_scalar(f"train/{key}", row[key], global_step)
                    if wandb_run is not None:
                        wandb_run.log({f"train/{k}": v for k, v in row.items()}, step=global_step)
                    print(
                        f"[train] step={global_step} loss={row['loss']:.4f} "
                        f"mean_reward={row['mean_reward']:.3f} best_reward={row['best_reward']:.3f} "
                        f"chosen_ok={chosen_ok}",
                        flush=True,
                    )
        optimizer.zero_grad(set_to_none=True)

    adapter_dir = out_dir / "adapter"
    if tb_writer is not None:
        tb_writer.flush()
        tb_writer.close()
    if wandb_run is not None:
        wandb_run.finish()
    model.save_pretrained(adapter_dir)
    tok.save_pretrained(adapter_dir)
    (out_dir / "train_metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    try:
        import matplotlib.pyplot as plt

        if metrics:
            xs = [m["step"] for m in metrics]
            ys = [m["best_reward"] for m in metrics]
            ok = [m["chosen_ok"] for m in metrics]
            losses = [m["loss"] for m in metrics]
            plt.figure(figsize=(8, 4))
            plt.plot(xs, ys, marker="o", label="best candidate reward")
            plt.plot(xs, ok, marker=".", label="policy picks max-reward candidate")
            plt.xlabel("optimizer step")
            plt.ylabel("metric")
            plt.grid(True, alpha=0.3)
            plt.legend()
            plt.tight_layout()
            plt.savefig(out_dir / "train_curve.png", dpi=160)
            plt.close()

            plt.figure(figsize=(8, 4))
            plt.plot(xs, losses, marker="o", color="#9467bd", label="RL objective loss")
            plt.axhline(0.0, color="black", linewidth=0.8)
            plt.title(f"{args.method} training loss")
            plt.xlabel("logged train step")
            plt.ylabel("loss")
            plt.grid(True, alpha=0.3)
            plt.legend()
            plt.tight_layout()
            plt.savefig(out_dir / "loss_curve.png", dpi=160)
            plt.close()
    except Exception as e:
        print(f"[plot] train curve skipped: {e}", flush=True)
    del loaded, model, tok
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return adapter_dir


def load_for_eval(args, adapter_dir: Path):
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    tok = AutoTokenizer.from_pretrained(adapter_dir, token=args.hf_token or None)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    bnb = None
    if not args.no_4bit:
        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        token=args.hf_token or None,
        device_map="auto",
        quantization_config=bnb,
        torch_dtype=torch.float16,
    )
    model.resize_token_embeddings(len(tok))
    model = PeftModel.from_pretrained(model, adapter_dir)
    model.eval()
    return tok, model


def generate_action(tok, model, obs: CoSObservation, max_new_tokens: int) -> tuple[CoSAction, str]:
    import torch

    prompt = prompt_for_obs(tok, obs, variant=0)
    ids = tok(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tok.pad_token_id or tok.eos_token_id,
        )
    text = tok.decode(out[0][ids["input_ids"].shape[1] :], skip_special_tokens=True)
    return parse_action(text), text.strip()


def deterministic_next(obs: CoSObservation, task: str) -> CoSAction:
    missing = [e for e in required_experts_for_task(task) if e not in obs.consulted_experts]
    if missing:
        return CoSAction(action_type="consult", expert_id=missing[0])
    if obs.current_brief is None:
        return CoSAction(action_type="summarize")
    return CoSAction(action_type="submit")


def evaluate(args, adapter_dir: Path, eval_dir: Path) -> list[dict[str, Any]]:
    import matplotlib.pyplot as plt
    import torch

    tok, model = load_for_eval(args, adapter_dir)
    rows = []
    rag_modes = [s.strip().lower() in {"1", "true", "yes", "rag"} for s in args.eval_rag_modes.split(",")]
    for use_rag in rag_modes:
        for task in [t.strip() for t in args.eval_tasks.split(",") if t.strip()]:
            env = CEOBriefEnvironment(shaping="strict", auto_fill_required=False)
            obs = env.reset(task=task, use_rag=use_rag)
            trace = []
            rewards = []
            routed = []
            for _ in range(args.policy_steps):
                if obs.done:
                    break
                action, completion = generate_action(tok, model, obs, args.eval_new_tokens)
                obs = env.step(action)
                rewards.append(float(obs.reward))
                if action.expert_id in required_experts_for_task(task) and action.expert_id not in routed:
                    routed.append(action.expert_id)
                trace.append(
                    {
                        "step": obs.step_count,
                        "action": action.model_dump(exclude_none=True),
                        "action_label": action_label(action),
                        "completion_preview": completion[:300],
                        "reward": round(float(obs.reward), 4),
                        "consulted_after": list(obs.consulted_experts),
                        "model_routed_required": list(routed),
                    }
                )
            fallback = []
            while not obs.done and obs.step_count < obs.max_steps:
                act = deterministic_next(obs, task)
                obs = env.step(act)
                rewards.append(float(obs.reward))
                fallback.append(action_label(act))
            rows.append(
                {
                    "task": task,
                    "rag": bool(use_rag),
                    "action_sequence": [t["action_label"] for t in trace],
                    "model_routed_required": routed,
                    "required_experts": required_experts_for_task(task),
                    "fallback": fallback,
                    "needed_fallback": bool(fallback),
                    "policy_reward": round(sum(t["reward"] for t in trace), 4),
                    "total_reward": round(sum(rewards), 4),
                    "terminal_score": round(float(obs.terminal_grader_score or 0.0), 4),
                    "trace": trace,
                }
            )
    eval_dir.mkdir(parents=True, exist_ok=True)
    (eval_dir / "evidence.json").write_text(json.dumps(rows, indent=2, default=str), encoding="utf-8")
    md = [
        "# AutoDataLab++ RL Evidence",
        "",
        "| Task | RAG | Action sequence | Routed required experts | Needed fallback | Policy reward | Terminal |",
        "|---|---:|---|---|---:|---:|---:|",
    ]
    for row in rows:
        md.append(
            f"| {row['task']} | {row['rag']} | `{' -> '.join(row['action_sequence'])}` | "
            f"{', '.join(row['model_routed_required']) or '-'} | {row['needed_fallback']} | "
            f"{row['policy_reward']} | {row['terminal_score']} |"
        )
    (eval_dir / "evidence.md").write_text("\n".join(md), encoding="utf-8")
    plt.figure(figsize=(8, 4))
    for row in rows:
        total = 0.0
        curve = []
        for t in row["trace"]:
            total += float(t["reward"])
            curve.append(total)
        if curve:
            plt.plot(range(1, len(curve) + 1), curve, marker="o", label=f"{row['task']} rag={row['rag']}")
    plt.title("RL policy reward on model-controlled steps")
    plt.xlabel("policy step")
    plt.ylabel("cumulative strict reward")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(eval_dir / "reward_curve.png", dpi=160)
    plt.close()
    del tok, model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return rows


def evidence_score(rows: list[dict[str, Any]]) -> float:
    """Single scalar for quick continue/skip decisions."""
    if not rows:
        return -999.0
    vals = []
    for row in rows:
        required = set(row.get("required_experts") or [])
        routed = set(row.get("model_routed_required") or [])
        coverage = len(required & routed) / max(len(required), 1)
        no_fallback = 1.0 if not row.get("needed_fallback") else 0.0
        policy_reward = float(row.get("policy_reward") or 0.0)
        vals.append(coverage + no_fallback + 0.1 * policy_reward)
    return sum(vals) / len(vals)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--method", choices=("grpo", "grpo_rlvr", "ppo"), required=True)
    ap.add_argument("--model-id", default="Qwen/Qwen2.5-1.5B-Instruct")
    ap.add_argument("--init-adapter", default="", help="optional SFT/DPO adapter directory or HF repo")
    ap.add_argument("--run-name", default="")
    ap.add_argument("--out-root", type=Path, default=Path("/kaggle/working/cos_1p5b_rl_runs") if Path("/kaggle/working").is_dir() else Path("cos_1p5b_rl_runs"))
    ap.add_argument("--hf-token", default=os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "")
    ap.add_argument("--no-4bit", action="store_true")
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument(
        "--lr",
        type=float,
        default=5e-6,
        help="conservative default; higher LR can collapse a good SFT/DPO adapter to noop",
    )
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--max-length", type=int, default=1024)
    ap.add_argument("--variants", type=int, default=2)
    ap.add_argument("--max-train-states", type=int, default=0)
    ap.add_argument("--log-every", type=int, default=5)
    ap.add_argument("--max-grad-norm", type=float, default=0.3)
    ap.add_argument("--ppo-clip", type=float, default=0.2)
    ap.add_argument(
        "--sft-anchor",
        type=float,
        default=0.2,
        help="oracle logprob anchor; keep >0 when continuing from a good SFT/DPO adapter",
    )
    ap.add_argument("--lora-r", type=int, default=16)
    ap.add_argument("--lora-alpha", type=int, default=32)
    ap.add_argument("--policy-steps", type=int, default=6)
    ap.add_argument("--eval-new-tokens", type=int, default=48)
    ap.add_argument("--eval-tasks", default="expert_brief,risk_brief,crisis_brief")
    ap.add_argument("--eval-rag-modes", default="false,true", help="comma list: false,true")
    ap.add_argument(
        "--report-to",
        choices=("none", "tensorboard", "wandb"),
        default=os.environ.get("REPORT_TO", "tensorboard"),
        help="experimental tracking backend for RL loss/reward logs; default tensorboard for judging",
    )
    ap.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "autodatalab-plus"))
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()
    if args.max_train_states == 0:
        args.max_train_states = None
    set_seed(args.seed)

    run_name = args.run_name or f"{args.method}_qwen15b_seed{args.seed}"
    out_dir = args.out_root / run_name
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "config.json").write_text(json.dumps(vars(args), indent=2, default=str), encoding="utf-8")
    print(f"[run] {run_name} -> {out_dir}", flush=True)

    before_rows = []
    if args.init_adapter:
        print("[baseline] evaluating init adapter before RL...", flush=True)
        before_rows = evaluate(args, Path(args.init_adapter), out_dir / "eval_before")
        print(f"[baseline] evidence_score={evidence_score(before_rows):.4f}", flush=True)

    adapter_dir = train_candidate_rl(args, out_dir)
    rows = evaluate(args, adapter_dir, out_dir / "eval")
    after_score = evidence_score(rows)
    before_score = evidence_score(before_rows) if before_rows else None
    print("\n=== RL EVIDENCE SUMMARY ===", flush=True)
    for row in rows:
        print(
            f"{row['task']}: {' -> '.join(row['action_sequence'])} | "
            f"routed={row['model_routed_required']} | fallback={row['needed_fallback']} | "
            f"policy_reward={row['policy_reward']} terminal={row['terminal_score']}",
            flush=True,
        )
    print(f"\n[evidence_score] after={after_score:.4f}", flush=True)
    if before_score is not None:
        print(f"[evidence_score] before={before_score:.4f}", flush=True)
        if after_score + 1e-6 < before_score:
            print(
                "[decision] SKIP this RL adapter: it regressed versus the init adapter.",
                flush=True,
            )
        else:
            print(
                "[decision] KEEP this RL adapter: it matched or improved the init adapter.",
                flush=True,
            )
    print(f"\n[adapter] {adapter_dir}", flush=True)
    print(f"[eval] {out_dir / 'eval'}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())