train.py

"""GRPO training entrypoint for DriftShield. Helpers live in ``train_*.py`` modules."""

from __future__ import annotations

import argparse
import csv
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TRL_EXPERIMENTAL_SILENCE", "1")

try:
    from support_ops_env import SupportOpsEnv
    from support_ops_env.tasks import get_curriculum_task_ids
except ImportError:
    from client import SupportOpsEnv
    from tasks import get_curriculum_task_ids

from train_action_json import parse_tool_calls, to_action
from train_obs_format import (
    apply_chat_template,
    format_history,
    format_observation,
    observation_coach_lines,
)
from train_reward_chart import plot_rewards
from train_reward_signal import (
    milestone_reward_from_history,
    reward_fields,
    reward_grounding,
    reward_reply,
    reward_total,
)
from train_rollout_episode import rollout_once
from train_system_prompt import SYSTEM_PROMPT
from train_tool_bonuses import TOOL_TIERS
from train_vllm_compat import (
    _require_vllm_trl_colocate_safe,
    patch_trl_vllm_compat,
    require_vllm_trl_colocate_safe,
)

logger = logging.getLogger(__name__)

__all__ = [
    "SYSTEM_PROMPT",
    "TOOL_TIERS",
    "apply_chat_template",
    "format_history",
    "format_observation",
    "milestone_reward_from_history",
    "observation_coach_lines",
    "parse_tool_calls",
    "patch_trl_vllm_compat",
    "plot_rewards",
    "require_vllm_trl_colocate_safe",
    "reward_fields",
    "reward_grounding",
    "reward_reply",
    "reward_total",
    "rollout_once",
    "to_action",
    "main",
    "parse_args",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="GRPO training for DriftShield")
    parser.add_argument(
        "--model-id",
        default=None,
        help="HF model id (default: Qwen3-1.7B stable, or Qwen3.5-4B with --use-unsloth).",
    )
    parser.add_argument(
        "--use-unsloth",
        action="store_true",
        help="Unsloth bf16 LoRA (needs pip install -e '.[unsloth]').",
    )
    parser.add_argument("--env-url", default="http://localhost:8000")
    parser.add_argument("--dataset-size", type=int, default=50)
    parser.add_argument("--max-turns", type=int, default=15)
    parser.add_argument("--num-generations", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=2e-6)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=4)
    parser.add_argument("--num-epochs", type=int, default=1)
    parser.add_argument("--save-steps", type=int, default=10)
    parser.add_argument("--output-dir", default=None)
    parser.add_argument("--use-vllm", action="store_true")
    parser.add_argument("--vllm-mode", choices=("colocate", "server"), default="colocate")
    parser.add_argument("--vllm-gpu-memory-utilization", type=float, default=0.5)
    parser.add_argument("--load-in-4bit", action="store_true")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-p", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=20)
    parser.add_argument("--logging-steps", type=int, default=1)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument("--reward-log", default="reward_log.csv")
    parser.add_argument("--difficulty", default="driftshield_easy")
    parser.add_argument("--curriculum-state", default=None)
    parser.add_argument("--mastery-min-repeat", type=int, default=1)
    parser.add_argument("--mastery-max-repeat", type=int, default=4)
    parser.add_argument("--mastery-cold-start", type=int, default=3)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    args = parse_args()

    if args.model_id is None:
        args.model_id = "unsloth/Qwen3.5-4B" if args.use_unsloth else "Qwen/Qwen3-1.7B"

    from datasets import Dataset
    from trl import GRPOConfig, GRPOTrainer

    patch_trl_vllm_compat()

    loaded_model: Any = None
    if args.use_unsloth:
        try:
            from unsloth import FastLanguageModel  # type: ignore
        except ImportError as exc:
            raise SystemExit("--use-unsloth requires: pip install -e '.[unsloth]'") from exc

        logger.info("[unsloth] loading %s", args.model_id)
        loaded_model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=args.model_id,
            max_seq_length=4096,
            load_in_4bit=False,
            load_in_16bit=True,
            full_finetuning=False,
            fast_inference=False,
        )
        loaded_model = FastLanguageModel.get_peft_model(
            loaded_model,
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            target_modules=[
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj",
            ],
            bias="none",
            use_gradient_checkpointing="unsloth",
            random_state=3407,
            max_seq_length=4096,
        )
    else:
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(args.model_id)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    env = SupportOpsEnv(base_url=args.env_url).sync()
    dataset = Dataset.from_dict(
        {"prompt": ["Triage and resolve this support operations case."] * args.dataset_size}
    )

    ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    output_dir = Path(args.output_dir or f"outputs/driftshield-grpo-{ts}")
    output_dir.mkdir(parents=True, exist_ok=True)

    reward_log = output_dir / args.reward_log
    with open(reward_log, "w", newline="") as fh:
        csv.writer(fh).writerow([
            "episode", "task_id",
            "total_reward",
            "investigation", "routing", "reply_quality", "groundedness", "submission",
            "penalty_total",
            "parse_ok_ratio",
            "timestamp",
        ])

    curriculum = get_curriculum_task_ids(
        args.difficulty,
        curriculum_state_path=args.curriculum_state,
        mastery_min_repeat=args.mastery_min_repeat,
        mastery_max_repeat=args.mastery_max_repeat,
        mastery_cold_start_episodes=args.mastery_cold_start,
    )
    logger.info("curriculum [%s] -> %s", args.difficulty, curriculum)
    if args.curriculum_state:
        logger.info(
            "mastery state=%s min=%s max=%s cold_start=%s",
            args.curriculum_state,
            args.mastery_min_repeat,
            args.mastery_max_repeat,
            args.mastery_cold_start,
        )
    task_cursor = [0]
    episode_counter = [0]

    def next_task_id() -> str:
        tid = curriculum[task_cursor[0] % len(curriculum)]
        task_cursor[0] += 1
        return tid

    diag_every = 5

    def log_episode_row(ep: Dict[str, Any]) -> None:
        episode_counter[0] += 1
        n = episode_counter[0]
        parse_ok = float(ep.get("parse_ok_ratio", 0.0))
        with open(reward_log, "a", newline="") as fh:
            csv.writer(fh).writerow([
                n,
                ep.get("task_id", ""),
                round(ep["total_reward"], 4),
                round(ep["investigation_reward"], 4),
                round(ep["routing_reward"], 4),
                round(ep["reply_reward"], 4),
                round(ep["grounding_reward"], 4),
                round(ep["submission_reward"], 4),
                round(ep["penalty_total"], 4),
                round(parse_ok, 4),
                datetime.now().isoformat(),
            ])

        if n % diag_every == 0:
            failures = int(ep.get("parse_failures", 0))
            attempts = int(ep.get("parse_attempts", 0))
            err = ep.get("last_parse_error") or "(none)"
            raw = (ep.get("last_raw_completion") or "").replace("\n", " ")[:160]
            logger.info(
                "[diag ep=%d] parse_ok=%.2f (%d/%d) | total=%+.2f inv=%.2f route=%.2f reply=%.2f "
                "ground=%.2f sub=%.2f | last_err=%s | last_raw[:160]=%r",
                n, parse_ok, attempts - failures, attempts,
                float(ep.get("total_reward", 0.0)),
                float(ep.get("investigation_reward", 0.0)),
                float(ep.get("routing_reward", 0.0)),
                float(ep.get("reply_reward", 0.0)),
                float(ep.get("grounding_reward", 0.0)),
                float(ep.get("submission_reward", 0.0)),
                err, raw,
            )
            if parse_ok < 0.5:
                logger.warning(
                    "[diag ep=%d] parse_ok_ratio<0.5 — check prompt, max_completion_length, base model",
                    n,
                )
            elif (
                parse_ok >= 0.9
                and float(ep.get("reply_reward", 0.0)) == 0.0
                and float(ep.get("grounding_reward", 0.0)) == 0.0
                and float(ep.get("submission_reward", 0.0)) == 0.0
            ):
                logger.warning("[diag ep=%d] parses OK but stalls before reply/submit", n)

    def rollout_func(prompts: List[str], trainer: Any) -> Dict[str, List[Any]]:
        out: Dict[str, List[Any]] = {
            "prompt_ids": [], "completion_ids": [], "logprobs": [],
            "total_reward": [], "field_reward": [], "reply_reward": [], "grounding_reward": [],
        }
        if hasattr(trainer, "vllm_generation") and trainer.vllm_generation is not None:
            current_step = int(getattr(getattr(trainer, "state", None), "global_step", -1))
            if getattr(trainer, "_driftshield_vllm_synced_step", None) != current_step:
                trainer.vllm_generation.sync_weights()
                trainer._driftshield_vllm_synced_step = current_step
        for _ in prompts:
            ep = rollout_once(
                trainer, env, tokenizer, SYSTEM_PROMPT, args.max_turns,
                task_id=next_task_id(),
                temperature=args.temperature,
                top_p=args.top_p,
                top_k=args.top_k,
            )
            log_episode_row(ep)
            for key in out:
                value = ep.get(key)
                if value is None:
                    raise KeyError(
                        f"rollout_once() missing key {key!r}. Available: {sorted(ep.keys())}"
                    )
                out[key].append(value)
        return out

    grpo_kwargs: Dict[str, Any] = dict(
        output_dir=str(output_dir),
        num_train_epochs=args.num_epochs,
        learning_rate=args.learning_rate,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        per_device_train_batch_size=args.num_generations,
        num_generations=args.num_generations,
        max_completion_length=768,
        logging_steps=args.logging_steps,
        save_strategy="steps",
        save_steps=args.save_steps,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        report_to="none",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        save_total_limit=3,
    )
    if args.use_vllm and not args.use_unsloth:
        _require_vllm_trl_colocate_safe()
        grpo_kwargs.update(
            use_vllm=True,
            vllm_mode=args.vllm_mode,
            vllm_gpu_memory_utilization=args.vllm_gpu_memory_utilization,
        )
    elif args.use_vllm and args.use_unsloth:
        logger.warning("--use-vllm ignored in Unsloth mode")

    if args.load_in_4bit and not args.use_unsloth:
        import torch
        from transformers import BitsAndBytesConfig

        compute_dtype = (
            torch.bfloat16
            if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
            else torch.float16
        )
        grpo_kwargs["model_init_kwargs"] = {
            "torch_dtype": compute_dtype,
            "quantization_config": BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=compute_dtype,
                bnb_4bit_use_double_quant=True,
            ),
        }
    elif not args.use_unsloth:
        import torch

        grpo_kwargs["model_init_kwargs"] = {"torch_dtype": torch.bfloat16}
    elif args.load_in_4bit and args.use_unsloth:
        logger.warning("--load-in-4bit ignored in Unsloth mode")

    grpo_config = GRPOConfig(**grpo_kwargs)

    if args.use_unsloth:
        model_arg: Any = loaded_model
        peft_config: Optional[Any] = None
    else:
        from peft import LoraConfig

        model_arg = args.model_id
        peft_config = LoraConfig(
            r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout,
            bias="none", task_type="CAUSAL_LM",
            target_modules=[
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj",
            ],
        )

    trainer = GRPOTrainer(
        model=model_arg,
        processing_class=tokenizer,
        reward_funcs=[reward_total, reward_fields, reward_reply, reward_grounding],
        train_dataset=dataset,
        args=grpo_config,
        rollout_func=rollout_func,
        peft_config=peft_config,
    )

    if args.dry_run:
        logger.info("--dry-run: skipping trainer.train()")
        return

    try:
        trainer.train()
    finally:
        env.close() if hasattr(env, "close") else None

    trainer.save_model(str(output_dir))
    plot_rewards(reward_log, output_dir / "reward_curve.png")
    logger.info("done — %s", output_dir)


if __name__ == "__main__":
    main()