import re
import json
import uuid
import random
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, prepare_model_for_kbit_training
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
from server.simulation import MedchainSimulation
from server.tasks import make_task_config
from config import MODEL_ID, G, B, MAX_NEW_TOKENS, MAX_TURNS, MAX_ROUND_TURNS, MAX_STEPS, LR, SAVE_STEPS
SYSTEM_PROMPT = """You are the central supply coordinator for a hospital network running 8 rounds per episode.
Each round lasts 2 simulated days. After advance_round your conversation history is wiped — only the round brief carries over.
═══ MANDATORY ROUND SEQUENCE ═══
1. read_inbox(filter="unread") ← ALWAYS first; catches crises and recall alerts
2. view_requests() ← see what wards are asking for (may be padded)
3. query_erp(table="pipeline_orders") ← CHECK IN-TRANSIT ORDERS BEFORE ORDERING MORE
4. query_erp(table="inventory", location="central_pharmacy") ← current on-hand stock
[steps 3+4 can be called in parallel in a single turn]
5. query_ward_history / query_supplier ← additional context if needed (1-2 calls max)
[multiple query_ward_history calls can be parallelised across wards in one turn]
6. submit_po(...) [+ file_justification if expedited] ← order ONLY the net gap (see PO rules below)
7. quarantine_lot(...) ← ONLY if inbox contains a recall or cold-chain breach alert
8. submit_allocation_plan(plan_json=...) ← REQUIRED every round; see rules below
9. advance_round() ← LAST call; ends the round
═══ REFERENCE — VALID IDs (use EXACTLY these strings) ═══
Locations (for query_erp and submit_po destination_id):
central_pharmacy ward_icu ward_er ward_general
Suppliers (for query_supplier and submit_po):
MEDLINE (lead 2d, cost 1.0×, all SKUs)
BACKUP-B (lead 3d, cost 1.3×, all SKUs — fallback when MEDLINE disrupted)
FASTMED (lead 1d, cost 1.8×, all SKUs — use only for life-critical emergencies)
═══ PURCHASE ORDER RULES (critical for budget score) ═══
• ALWAYS query pipeline_orders BEFORE placing any PO. Your context is wiped after each round —
last round's orders are still in transit and will appear in pipeline_orders.
• Derive consumption from query_ward_history — do NOT invent or memorise fixed quantities.
• Compute for each SKU: net_qty = history_avg_consumption - in_transit_qty - surplus_on_hand.
*** If net_qty ≤ 0: DO NOT call submit_po. Calling submit_po with quantity ≤ 0 is FORBIDDEN. ***
*** If net_qty > 0: call submit_po ONCE with quantity = net_qty. ONE call per SKU total. ***
• Sum need across all wards before calling submit_po — never submit per-ward separately.
All stock routes through central_pharmacy regardless of destination.
• Do NOT call view_requests() more than once per round.
═══ WARD SKU REFERENCE ═══
ward_icu : BLOOD-RBC BLOOD-PLT BLOOD-FFP IV-SAL-500 ANTIBIO-01 OXY-MASK SYR-10 GLOVE-001
ward_er : BLOOD-RBC BLOOD-PLT BLOOD-FFP IV-SAL-500 ANTIBIO-01 OXY-MASK SYR-10 GLOVE-001 GAUZE-01
ward_general : IV-SAL-500 ANTIBIO-01 SYR-10 GLOVE-001 MASK-001 GAUZE-01
═══ ALLOCATION PLAN RULES (critical for score) ═══
• Include ALL THREE wards — ward_icu, ward_er, ward_general — every single round.
• Include EVERY SKU listed above for each ward. Omitting a SKU is treated as 0 (stockout penalty).
• Never allocate 0 for any SKU. Use ward_history average consumption when uncertain.
• Start from request quantities, then discount by the inflation factor (wards pad 10-60%).
ICU pads ~10%, ER pads ~20-50%, General pads ~25-60% — allocate closer to true need.
• Blood products (BLOOD-RBC, BLOOD-PLT, BLOOD-FFP) are CRITICAL for ICU and ER — never stockout.
• Allocation comes from central_pharmacy stock. Submit POs first if stock is low.
═══ QUERY RULES ═══
• query_erp: location and sku are optional — but NEVER pass null or empty string.
Omit the field entirely when you want all rows: query_erp(table="inventory", location="central_pharmacy")
• query_supplier: use only MEDLINE, BACKUP-B, or FASTMED — no other IDs exist.
• file_justification ticket_id: use the PO ticket ID returned by submit_po response, not a guessed ID.
═══ EVENTS TO WATCH FOR ═══
• MCI (mass casualty): blood demand ×2.8 at ICU+ER — pre-position via expedited FASTMED orders.
• Supplier disruption: switch to BACKUP-B for urgent items; FASTMED for life-critical.
• Product recall: quarantine_lot at every listed location using the EXACT lot_id string from the inbox message body — never invent or guess it.
• Cold-chain breach: lot is auto-quarantined by the system; place emergency replenishment order — no quarantine_lot call needed."""
TOOL_SCHEMAS = [
{"type": "function", "function": {
"name": "read_inbox",
"description": "Read inbox messages. Always call first each round.",
"parameters": {"type": "object", "properties": {
"filter": {"type": "string", "enum": ["unread", "all", "flagged"]}
}, "required": ["filter"]},
}},
{"type": "function", "function": {
"name": "view_requests",
"description": "View pending ward supply requests. Requests may be inflated.",
"parameters": {"type": "object", "properties": {}},
}},
{"type": "function", "function": {
"name": "query_ward_history",
"description": "Query historical requests and allocations for a ward.",
"parameters": {"type": "object", "properties": {
"ward_id": {"type": "string"},
"product_id": {"type": "string"},
"n_rounds": {"type": "integer"},
}, "required": ["ward_id"]},
}},
{"type": "function", "function": {
"name": "query_erp",
"description": (
"Query ERP system for inventory, expiry, or pipeline orders. "
"location must be one of: central_pharmacy, ward_icu, ward_er, ward_general. "
"sku must be a valid SKU string. OMIT location or sku entirely to get all rows — "
"never pass null or an empty string."
),
"parameters": {"type": "object", "properties": {
"table": {"type": "string", "enum": ["inventory", "expiry", "pipeline_orders"]},
"location": {"type": "string", "enum": ["central_pharmacy", "ward_icu", "ward_er", "ward_general"]},
"sku": {"type": "string"},
}, "required": ["table"]},
}},
{"type": "function", "function": {
"name": "query_supplier",
"description": "Query supplier lead time and disruption status. Valid supplier_id values: MEDLINE, BACKUP-B, FASTMED.",
"parameters": {"type": "object", "properties": {
"supplier_id": {"type": "string", "enum": ["MEDLINE", "BACKUP-B", "FASTMED"]},
}, "required": ["supplier_id"]},
}},
{"type": "function", "function": {
"name": "submit_po",
"description": "Submit a purchase order to a supplier. destination_id must be central_pharmacy (the central hub — stock is allocated from there to wards). Use FASTMED only for life-critical emergencies (1d lead, 1.8× cost).",
"parameters": {"type": "object", "properties": {
"supplier_id": {"type": "string", "enum": ["MEDLINE", "BACKUP-B", "FASTMED"]},
"product_id": {"type": "string"},
"destination_id": {"type": "string", "enum": ["central_pharmacy", "ward_icu", "ward_er", "ward_general"]},
"quantity": {"type": "integer"},
"priority": {"type": "string", "enum": ["standard", "expedited"]},
}, "required": ["supplier_id", "product_id", "destination_id", "quantity"]},
}},
{"type": "function", "function": {
"name": "file_justification",
"description": "File justification for an expedited purchase order.",
"parameters": {"type": "object", "properties": {
"ticket_id": {"type": "string"},
"reason": {"type": "string"},
}, "required": ["ticket_id", "reason"]},
}},
{"type": "function", "function": {
"name": "quarantine_lot",
"description": "Quarantine a lot due to recall or cold-chain breach.",
"parameters": {"type": "object", "properties": {
"location_id": {"type": "string"},
"sku": {"type": "string"},
"lot_id": {"type": "string"},
}, "required": ["location_id", "sku", "lot_id"]},
}},
{"type": "function", "function": {
"name": "submit_allocation_plan",
"description": "Submit stock allocation plan across all wards as JSON string {ward_id: {sku: qty}}.",
"parameters": {"type": "object", "properties": {
"plan_json": {"type": "string"},
}, "required": ["plan_json"]},
}},
{"type": "function", "function": {
"name": "advance_round",
"description": "Close the current round and advance to the next. Call after all actions are complete.",
"parameters": {"type": "object", "properties": {}},
}},
]
# Integer params that need type coercion (native format gives all values as strings)
_INT_PARAMS: dict[str, set[str]] = {
tool["function"]["name"]: {
k for k, v in tool["function"]["parameters"].get("properties", {}).items()
if v.get("type") == "integer"
}
for tool in TOOL_SCHEMAS
}
def parse_tool_calls(text: str) -> list[dict]:
"""Parse Qwen3.5 native XML tool call format."""
results = []
for block in re.finditer(r'(.*?)', text, re.DOTALL):
inner = block.group(1)
fn = re.search(r']+)>', inner)
if not fn:
continue
name = fn.group(1).strip()
args = {}
for p in re.finditer(r']+)>\n?(.*?)\n?', inner, re.DOTALL):
k, v = p.group(1).strip(), p.group(2).strip()
if k in _INT_PARAMS.get(name, set()):
try:
v = int(v)
except (ValueError, TypeError):
pass
args[k] = v
results.append({"name": name, "arguments": args})
return results
# Reward closure — rollout_func populates this, reward_func reads it
_current_rewards: list[float] = []
def reward_func(completions, prompts, **kwargs) -> list[float]:
return list(_current_rewards)
def medchain_rollout(prompts, trainer) -> dict:
"""
Custom rollout with fresh context windows per round.
Receives B=8 prompts from TRL. Creates G=4 episodes per seed = 32 total.
Each episode resets its context window after every advance_round() call.
All active episodes are batched into one forward pass per loop iteration.
Returns 32 results; TRL groups consecutive G=4 as one GRPO group.
"""
global _current_rewards
model = trainer.model
tok = trainer.processing_class
model.eval()
# ── Initialise 32 episodes (B=8 seeds × G=4 rollouts) ──────────────────
episodes = []
for prompt in prompts:
# seed and difficulty embedded in the last user message by the dataset
meta = prompt[-1]["content"]
seed = int(re.search(r"seed=(\d+)", meta).group(1))
diff = re.search(r"diff=(\w+)", meta).group(1)
for _ in range(G):
sim = MedchainSimulation(make_task_config(seed=seed, difficulty=diff))
brief = sim.reset(seed=seed, episode_id=str(uuid.uuid4()))
episodes.append({
"sim": sim,
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": brief},
],
"comp_ids": [], # per-turn cpu tensors
"logprobs": [], # per-turn cpu tensors
"first_prompt_ids": None,
"done": False,
"reward": 0.0,
"round_turns": 0,
})
print(f"[rollout] init {len(episodes)} episodes ({len(prompts)} seeds × G={G})")
# ── Batched fresh-window inference loop ─────────────────────────────────
for turn in range(MAX_TURNS):
active = [ep for ep in episodes if not ep["done"]]
if not active:
break
if turn % 10 == 0:
n_done = len(episodes) - len(active)
vram = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0
print(f"[rollout turn {turn:3d}] active={len(active)} done={n_done}/{len(episodes)} VRAM={vram:.2f}GB")
tok.padding_side = "left"
texts = [
tok.apply_chat_template(
ep["messages"],
tools=TOOL_SCHEMAS,
tokenize=False,
add_generation_prompt=True,
chat_template_kwargs={"enable_thinking": False},
)
for ep in active
]
inputs = tok(texts, return_tensors="pt", padding=True).to(model.device)
padded_len = inputs["input_ids"].shape[1]
with torch.no_grad():
gen = model.generate(
**inputs,
max_new_tokens=MAX_NEW_TOKENS,
temperature=0.7,
top_p=0.8,
top_k=20,
do_sample=True,
return_dict_in_generate=True,
output_scores=True,
pad_token_id=tok.eos_token_id,
)
for i, ep in enumerate(active):
comp = gen.sequences[i, padded_len:].cpu()
# Trim at first EOS
eos_pos = (comp == tok.eos_token_id).nonzero()
if len(eos_pos):
comp = comp[: eos_pos[0].item() + 1]
# Per-token logprobs from generation scores
T = len(comp)
if T > 0 and gen.scores:
scores = torch.stack([gen.scores[t][i].cpu() for t in range(T)])
lp = F.log_softmax(scores.float(), dim=-1)
token_lp = lp.gather(1, comp.unsqueeze(-1)).squeeze(-1)
else:
token_lp = torch.zeros(T)
ep["comp_ids"].append(comp)
ep["logprobs"].append(token_lp)
# Store first-round prompt_ids for TRL's loss mask
if ep["first_prompt_ids"] is None:
raw = gen.sequences[i, :padded_len].cpu()
non_pad = (raw != tok.pad_token_id).nonzero()
ep["first_prompt_ids"] = raw[non_pad[0].item():] if len(non_pad) else raw
# Decode and parse tool calls
text = tok.decode(comp, skip_special_tokens=True)
ep["messages"].append({"role": "assistant", "content": text})
tool_calls = parse_tool_calls(text)
if not tool_calls:
print(f" [ep{i}] turn={turn} round_turn={ep['round_turns']} — no tool call parsed text={text[:80]!r}")
ep["messages"].append({
"role": "user",
"content": "Use a tool. Call advance_round when done with this round.",
})
else:
for tc in tool_calls:
name = tc.get("name", "")
args = tc.get("arguments", {})
sim = ep["sim"]
try:
result = (
getattr(sim, name)(**args)
if hasattr(sim, name)
else f"ERROR: Unknown tool '{name}'"
)
except Exception as e:
result = f"ERROR: {e}"
print(f" [ep{i}] TOOL ERROR {name}({args}): {e}")
ep["messages"].append({"role": "user", "content": f"{result}"})
if name == "advance_round":
if sim._done:
ep["done"] = True
ep["reward"] = sim._last_reward
print(f" [ep{i}] DONE reward={ep['reward']:.4f}")
else:
ep["messages"] = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": result},
]
ep["round_turns"] = 0
print(f" [ep{i}] advance_round → next round round_turns reset")
break
ep["round_turns"] += 1
# Force advance if model gets stuck within a round
if not ep["done"] and ep["round_turns"] >= MAX_ROUND_TURNS:
print(f" [ep{i}] FORCE advance_round (stuck {ep['round_turns']} turns in round)")
result = ep["sim"].advance_round()
if ep["sim"]._done:
ep["done"] = True
ep["reward"] = ep["sim"]._last_reward
print(f" [ep{i}] DONE (forced) reward={ep['reward']:.4f}")
else:
ep["messages"] = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": result},
]
ep["round_turns"] = 0
# ── Package results for TRL ─────────────────────────────────────────────
_current_rewards = [ep["reward"] for ep in episodes]
n_done = sum(1 for ep in episodes if ep["done"])
rewards = _current_rewards
print(f"[rollout] done={n_done}/{len(episodes)} "
f"reward mean={sum(rewards)/len(rewards):.4f} "
f"min={min(rewards):.4f} max={max(rewards):.4f} "
f"zeros={rewards.count(0.0)}/{len(rewards)}")
return {
"prompt_ids": [
ep["first_prompt_ids"]
if ep["first_prompt_ids"] is not None
else torch.tensor([], dtype=torch.long)
for ep in episodes
],
"completion_ids": [
torch.cat(ep["comp_ids"])
if ep["comp_ids"]
else torch.tensor([], dtype=torch.long)
for ep in episodes
],
"logprobs": [
torch.cat(ep["logprobs"])
if ep["logprobs"]
else torch.tensor([], dtype=torch.float)
for ep in episodes
],
}
def load_model():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID, quantization_config=bnb_config, device_map="auto"
)
model = prepare_model_for_kbit_training(model)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
return model, tokenizer
peft_config = LoraConfig(
r=16,
lora_alpha=16,
target_modules="all-linear",
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
# Do NOT call get_peft_model here — TRL applies it via peft_config
# and uses disable_adapter() for the KL reference model.
def build_dataset() -> Dataset:
random.seed(42)
seeds_list = list(range(50)) * 20 # 1000 rows; seeds 0-49 each ×20
diffs_list = [random.choice(["light", "medium", "heavy"]) for _ in seeds_list]
return Dataset.from_dict({
"prompt": [
[
{"role": "system", "content": SYSTEM_PROMPT},
# Embed seed/difficulty so rollout_func can extract them
{"role": "user", "content": f"seed={s};diff={d}"},
]
for s, d in zip(seeds_list, diffs_list)
],
}).shuffle(seed=42)
def main():
import argparse
import wandb
parser = argparse.ArgumentParser()
parser.add_argument("--wandb-project", default="medchain-grpo")
parser.add_argument("--max-steps", type=int, default=MAX_STEPS)
parser.add_argument("--lr", type=float, default=LR)
parser.add_argument("--save-steps", type=int, default=SAVE_STEPS)
args = parser.parse_args()
wandb.init(project=args.wandb_project)
model, tokenizer = load_model()
train_dataset = build_dataset()
config = GRPOConfig(
num_generations=G, # TRL groups consecutive G=4 results as one GRPO group
per_device_train_batch_size=B, # B=8 seeds per step; rollout_func returns B*G=32 results
max_completion_length=16384, # ceiling on concatenated completion_ids per episode
learning_rate=args.lr,
max_steps=args.max_steps,
bf16=True,
output_dir="checkpoints",
save_steps=args.save_steps,
logging_steps=1,
report_to="wandb",
)
trainer = GRPOTrainer(
model=model,
processing_class=tokenizer,
reward_funcs=reward_func,
args=config,
train_dataset=train_dataset,
peft_config=peft_config,
rollout_func=medchain_rollout,
)
trainer.train()
if __name__ == "__main__":
main()