scripts/hf_train_supplymind_roles.py

# /// script
# dependencies = [
#   "torch",
#   "transformers>=4.45.0",
#   "trl>=0.12.0",
#   "peft>=0.13.0",
#   "accelerate",
#   "datasets",
#   "huggingface_hub",
#   "trackio",
#   "pydantic",
#   "pyyaml",
#   "matplotlib",
# ]
# ///
from __future__ import annotations

import argparse
import inspect
import json
import os
import re
import sys
import time
from pathlib import Path
from statistics import mean
from typing import Any

from datasets import Dataset
from huggingface_hub import snapshot_download
import torch
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import GRPOConfig, GRPOTrainer


REPO_ID = "rishavutk/supplymind"
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
REWARD_SCALE = 10.0
REWARD_CLIP = 20.0
REWARD_LOG_EVERY = 10
MALFORMED_REWARD = -2.0
INVALID_ACTION_PENALTY = 0.5


def log(message: str, **fields: Any) -> None:
    payload = {"message": message, **fields}
    print(json.dumps(payload, sort_keys=True), flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--role", choices=["center", "warehouse", "joint"], default="center")
    parser.add_argument("--task-id", default="v2_train_easy")
    parser.add_argument("--seeds", default="101,113,127")
    parser.add_argument("--max-steps", type=int, default=30)
    parser.add_argument("--max-completion-length", type=int, default=256)
    parser.add_argument("--hub-model-id", default="")
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--init-adapter-id", default="")
    parser.add_argument("--model-id", default=MODEL_ID)
    parser.add_argument("--load-in-4bit", action="store_true")
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--num-generations", type=int, default=2)
    parser.add_argument("--generation-batch-size", type=int, default=2)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=2)
    return parser.parse_args()


def prepare_repo() -> Path:
    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
    os.environ["PYTHONUTF8"] = "1"
    log("downloading_supplymind_space", repo_id=REPO_ID)
    target_dir = Path("supplymind_snapshot").resolve()
    local_dir = Path(
        snapshot_download(
            repo_id=REPO_ID,
            repo_type="space",
            allow_patterns=["src/**", "configs/**"],
            local_dir=target_dir,
        )
    )
    config_path = local_dir / "configs" / "supplymind_v2_rewards.yaml"
    os.environ["SUPPLYMIND_REWARD_CONFIG"] = str(config_path)
    sys.path.insert(0, str(local_dir / "src"))
    log(
        "downloaded_supplymind_space",
        path=str(local_dir),
        src_exists=(local_dir / "src").exists(),
        config_path=str(config_path),
        config_exists=config_path.exists(),
    )
    return local_dir


def completion_to_text(completion: Any) -> str:
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list) and completion:
        last = completion[-1]
        if isinstance(last, dict):
            return str(last.get("content", ""))
    return str(completion)


def extract_json(completion: Any) -> dict[str, Any] | None:
    text = completion_to_text(completion)
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        return None


def compact_observation(observation: Any, role: str, warehouse_id: str | None = None) -> dict[str, Any]:
    data = observation.model_dump(mode="json")
    data["scenario_info"].pop("public_rules", None)
    if role == "center":
        return {
            "role": "center",
            "round_index": data["round_index"],
            "remaining_rounds": data["remaining_rounds"],
            "task_id": data["task_id"],
            "scenario_info": data["scenario_info"],
            "center": data["center"],
            "warehouse_summaries": data["center"].get("warehouse_summaries", []),
            "market_signals": data["center"].get("market_signals", []),
            "feedback": data.get("feedback", {}),
        }
    if role == "warehouse":
        if warehouse_id:
            warehouse = data["warehouses"][warehouse_id]
            return {
                "role": "warehouse",
                "warehouse_id": warehouse_id,
                "round_index": data["round_index"],
                "remaining_rounds": data["remaining_rounds"],
                "task_id": data["task_id"],
                "scenario_info": data["scenario_info"],
                "warehouse": warehouse,
                "pending_transfer_proposals": warehouse.get("pending_transfer_proposals", []),
                "feedback": data.get("feedback", {}),
            }
        return {
            "role": "warehouse",
            "round_index": data["round_index"],
            "remaining_rounds": data["remaining_rounds"],
            "task_id": data["task_id"],
            "scenario_info": data["scenario_info"],
            "warehouses": data["warehouses"],
            "pending_transfer_proposals": data["center"].get("pending_transfer_proposals", []),
            "feedback": data.get("feedback", {}),
        }
    return data


def system_prompt(role: str) -> str:
    if role == "center":
        return (
            "You are the center policy in SupplyMind. Return only strict JSON matching CenterAction: "
            "central_procurements, central_liquidations, central_replenishments, inventory_transfer_proposals, offer_matches. "
            "Warehouses are controlled by a fixed heuristic. Earn margin and a small share of realized service profit, "
            "but avoid waste, stockouts, overpriced actions, and needless shipments."
        )
    if role == "warehouse":
        return (
            "You are the shared warehouse policy in SupplyMind, copied across all warehouses. "
            "You control exactly one warehouse from the user observation. Return only strict JSON matching WarehouseAction: "
            "order_decisions, inventory_offers, inventory_requests, transfer_responses, and local_priority. "
            "The center is controlled by a fixed heuristic. Accept orders you can serve, request needed stock, "
            "and reject bad or impossible commitments. Only use order_id and proposal_id values visible in this observation. "
            "Do not invent IDs, do not use markdown, and prefer fewer high-confidence actions over broad noisy actions."
        )
    return (
        "You are playing SupplyMind. Return only strict JSON with top-level keys warehouse_actions and central_action. "
        "Optimize global welfare while avoiding invalid actions, missed accepted orders, stockouts, waste, and pointless transfers."
    )


def build_rows(role: str, task_id: str, seeds: list[int]) -> list[dict[str, Any]]:
    from supplymind_env_v2.environment import V2SupplyMindEnv
    from supplymind_env_v2.policies import heuristic_joint_policy

    rows: list[dict[str, Any]] = []
    for seed in seeds:
        env = V2SupplyMindEnv(default_task_id=task_id)
        observation = env.reset_internal(task_id, seed)
        while not env.done:
            rows.append(
                {
                    "prompt": [
                        {"role": "system", "content": system_prompt(role)},
                        {"role": "user", "content": json.dumps(compact_observation(observation, role), separators=(",", ":"))},
                    ],
                    "task_id": task_id,
                    "seed": seed,
                    "round_index": observation.round_index,
                }
            )
            result = env.step(heuristic_joint_policy(observation), grade_terminal=False)
            observation = result.observation
    return rows


def build_warehouse_rows(task_id: str, seeds: list[int]) -> list[dict[str, Any]]:
    from supplymind_env_v2.environment import V2SupplyMindEnv
    from supplymind_env_v2.policies import heuristic_joint_policy

    rows: list[dict[str, Any]] = []
    for seed in seeds:
        env = V2SupplyMindEnv(default_task_id=task_id)
        observation = env.reset_internal(task_id, seed)
        while not env.done:
            for warehouse_id in observation.warehouses:
                rows.append(
                    {
                        "prompt": [
                            {"role": "system", "content": system_prompt("warehouse")},
                            {
                                "role": "user",
                                "content": json.dumps(compact_observation(observation, "warehouse", warehouse_id), separators=(",", ":")),
                            },
                        ],
                        "task_id": task_id,
                        "seed": seed,
                        "round_index": observation.round_index,
                        "warehouse_id": warehouse_id,
                    }
                )
            result = env.step(heuristic_joint_policy(observation), grade_terminal=False)
            observation = result.observation
    return rows


def make_reward_fn(role: str):
    from supplymind_env_v2.environment import V2SupplyMindEnv
    from supplymind_env_v2.models import CenterAction, V2JointAction, V2WarehouseRoleAction
    from supplymind_env_v2.policies import fixed_center_action, fixed_warehouse_actions, heuristic_joint_policy

    call_count = 0

    def replay_to_round(task_id: str, seed: int, round_index: int):
        env = V2SupplyMindEnv(default_task_id=task_id)
        observation = env.reset_internal(task_id, seed)
        while observation.round_index < round_index and not env.done:
            result = env.step(heuristic_joint_policy(observation), grade_terminal=False)
            observation = result.observation
        return env, observation

    def reward_completions(prompts, completions, task_id, seed, round_index, **kwargs) -> list[float]:
        nonlocal call_count
        call_count += 1
        rewards: list[float] = []
        invalid_payloads = 0
        invalid_actions = 0
        for completion, current_task, current_seed, current_round in zip(completions, task_id, seed, round_index, strict=True):
            payload = extract_json(completion)
            if payload is None:
                rewards.append(MALFORMED_REWARD)
                invalid_payloads += 1
                continue
            env, observation = replay_to_round(current_task, int(current_seed), int(current_round))
            before = dict(env.agent_rewards)
            try:
                if role == "center":
                    center_action = CenterAction.model_validate(payload)
                    action = V2JointAction(warehouse_actions=fixed_warehouse_actions(observation), central_action=center_action)
                elif role == "warehouse":
                    target_warehouse = kwargs.get("warehouse_id", [None] * len(completions))[len(rewards)]
                    if target_warehouse:
                        from supplymind_env_v2.models import WarehouseAction

                        warehouse_action = WarehouseAction.model_validate(payload)
                        role_action = V2WarehouseRoleAction(warehouse_actions={str(target_warehouse): warehouse_action})
                    else:
                        role_action = V2WarehouseRoleAction.model_validate(payload)
                    action = V2JointAction(
                        warehouse_actions=role_action.warehouse_actions,
                        central_action=fixed_center_action(observation, role_action.warehouse_actions),
                    )
                else:
                    action = V2JointAction.model_validate(payload)
            except Exception:
                rewards.append(MALFORMED_REWARD)
                invalid_payloads += 1
                continue
            result = env.step(action, grade_terminal=False)
            invalid_count = len(result.observation.feedback.get("invalid_action_details", []))
            invalid_actions += invalid_count
            if role == "center":
                role_delta = env.agent_rewards["center"] - before.get("center", 0.0)
            elif role == "warehouse":
                target_warehouse = kwargs.get("warehouse_id", [None] * len(completions))[len(rewards)]
                if target_warehouse:
                    role_delta = env.agent_rewards[str(target_warehouse)] - before.get(str(target_warehouse), 0.0)
                else:
                    warehouse_ids = [key for key in env.agent_rewards if key != "center"]
                    role_delta = mean(env.agent_rewards[key] - before.get(key, 0.0) for key in warehouse_ids)
            else:
                role_delta = float(result.reward.step_reward)
            scaled = max(-REWARD_CLIP, min(REWARD_CLIP, role_delta / REWARD_SCALE))
            rewards.append(scaled + 1.0 - INVALID_ACTION_PENALTY * invalid_count)
        if call_count == 1 or call_count % REWARD_LOG_EVERY == 0:
            log(
                "reward_batch",
                role=role,
                call=call_count,
                count=len(rewards),
                mean_reward=round(mean(rewards), 4) if rewards else None,
                min_reward=round(min(rewards), 4) if rewards else None,
                max_reward=round(max(rewards), 4) if rewards else None,
                invalid_payloads=invalid_payloads,
                invalid_actions=invalid_actions,
            )
        return rewards

    return reward_completions


def baseline_role_probe(role: str, task_id: str, seeds: list[int]) -> None:
    from supplymind_env_v2.environment import V2SupplyMindEnv
    from supplymind_env_v2.policies import heuristic_joint_policy, no_op_policy

    for policy_name, policy_fn in (("no_op", no_op_policy), ("heuristic", heuristic_joint_policy)):
        summaries: list[dict[str, Any]] = []
        for seed in seeds:
            env = V2SupplyMindEnv(default_task_id=task_id)
            observation = env.reset_internal(task_id, seed)
            while not env.done:
                result = env.step(policy_fn(observation))
                observation = result.observation
            summaries.append(env.last_episode_summary or {})
        log(
            "baseline_probe",
            role=role,
            policy=policy_name,
            task_id=task_id,
            seeds=seeds,
            mean_score=round(mean(float(row.get("graded_score", 0.0)) for row in summaries), 4),
            mean_center_role_score=round(mean(float(row.get("center_role_score", 0.0)) for row in summaries), 4),
            mean_warehouse_role_score=round(mean(float(row.get("warehouse_role_score", 0.0)) for row in summaries), 4),
            mean_raw_reward=round(mean(float(row.get("raw_reward", 0.0)) for row in summaries), 3),
            mean_center_reward=round(mean(float(row.get("center_reward", 0.0)) for row in summaries), 3),
            mean_average_warehouse_reward=round(mean(float(row.get("average_warehouse_reward", 0.0)) for row in summaries), 3),
        )


def make_grpo_config(
    output_dir: str,
    max_steps: int,
    hub_model_id: str,
    role: str,
    max_completion_length: int,
    num_generations: int,
    generation_batch_size: int,
    gradient_accumulation_steps: int,
) -> GRPOConfig:
    requested = {
        "output_dir": output_dir,
        "max_steps": max_steps,
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": gradient_accumulation_steps,
        "num_generations": num_generations,
        "generation_batch_size": generation_batch_size,
        "max_prompt_length": 2048,
        "max_completion_length": max_completion_length,
        "logging_steps": 1,
        "report_to": [],
        "project": "supplymind",
        "run_name": f"{role}-grpo-smoke",
        "push_to_hub": False,
        "hub_model_id": hub_model_id,
        "save_strategy": "no",
    }
    signature = inspect.signature(GRPOConfig)
    supported = {key: value for key, value in requested.items() if key in signature.parameters}
    dropped = sorted(set(requested) - set(supported))
    if dropped:
        log("grpo_config_dropped_unsupported_keys", dropped=dropped)
    return GRPOConfig(**supported)


def main() -> None:
    started = time.time()
    args = parse_args()
    prepare_repo()
    seeds = [int(value.strip()) for value in args.seeds.split(",") if value.strip()]
    output_dir = args.output_dir or f"outputs/supplymind-{args.role}-qwen-grpo"
    hub_model_id = args.hub_model_id or f"rishavutk/supplymind-{args.role}-qwen-0.5b-grpo"

    rows = build_warehouse_rows(args.task_id, seeds) if args.role == "warehouse" else build_rows(args.role, args.task_id, seeds)
    dataset = Dataset.from_list(rows)
    log(
        "dataset_ready",
        role=args.role,
        rows=len(rows),
        task_id=args.task_id,
        seeds=seeds,
        max_steps=args.max_steps,
        max_completion_length=args.max_completion_length,
        hub_model_id=hub_model_id,
    )
    baseline_role_probe(args.role, args.task_id, seeds)

    quantization_config = None
    if args.load_in_4bit:
        compute_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=compute_dtype,
        )
    log("loading_model", model_id=args.model_id, load_in_4bit=args.load_in_4bit)
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        torch_dtype="auto",
        device_map="auto",
        quantization_config=quantization_config,
    )
    if args.load_in_4bit:
        model = prepare_model_for_kbit_training(model)
    peft_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=max(args.lora_r * 2, 8),
        lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        task_type="CAUSAL_LM",
    )
    if args.init_adapter_id:
        log("loading_initial_adapter", adapter_id=args.init_adapter_id)
        model = PeftModel.from_pretrained(model, args.init_adapter_id, is_trainable=True)
        peft_config = None
    log("model_loaded", model_id=args.model_id, init_adapter_id=args.init_adapter_id or None)

    trainer_kwargs = {
        "model": model,
        "processing_class": tokenizer,
        "reward_funcs": make_reward_fn(args.role),
        "train_dataset": dataset,
        "args": make_grpo_config(
            output_dir,
            args.max_steps,
            hub_model_id,
            args.role,
            args.max_completion_length,
            args.num_generations,
            args.generation_batch_size,
            args.gradient_accumulation_steps,
        ),
    }
    if peft_config is not None:
        trainer_kwargs["peft_config"] = peft_config
    trainer = GRPOTrainer(**trainer_kwargs)
    log("training_start", role=args.role, max_steps=args.max_steps, max_completion_length=args.max_completion_length)
    trainer.train()
    log("training_done", elapsed_seconds=round(time.time() - started, 2))
    log("pushing_model", hub_model_id=hub_model_id)
    trainer.push_to_hub()
    log("job_done", hub_model_id=hub_model_id, elapsed_seconds=round(time.time() - started, 2))


if __name__ == "__main__":
    main()