""" GRPO Training with LogisticsShipmentRL Environment =================================================== Trains an LLM to act as an AI Logistics Coordinator using Group Relative Policy Optimization (TRL + OpenEnv). Usage: # Option 1: HF Space environment python train_grpo.py --model Qwen/Qwen3-1.7B --env-url https://YOUR_USERNAME-logistics-shipment-env.hf.space # Option 2: Local environment (run server first) # PYTHONPATH=src uvicorn envs.logistics_shipment_env.server.app:app --port 8000 --reload python train_grpo.py --model Qwen/Qwen3-1.7B --env-url http://localhost:8000 --vllm-mode colocate Install: pip install trl vllm datasets transformers pip install git+https://github.com/meta-pytorch/OpenEnv.git """ from __future__ import annotations import argparse import json import sys from pathlib import Path # -------------------------------------------------------------------------- # Args # -------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="GRPO training for Logistics Shipment RL") p.add_argument("--model", default="Qwen/Qwen3-1.7B", help="Model id (HF hub or local)") p.add_argument("--env-url", default="http://localhost:8000", help="URL for the logistics environment server") p.add_argument("--dataset-size", type=int, default=500, help="Training dataset size") p.add_argument("--max-turns", type=int, default=5, help="Max turns per episode") p.add_argument("--epochs", type=int, default=1) p.add_argument("--lr", type=float, default=5e-6) p.add_argument("--grad-accum", type=int, default=32) p.add_argument("--num-gen", type=int, default=2, help="Rollout generations per prompt") p.add_argument("--output-dir", default="outputs/logistics-grpo") p.add_argument("--vllm-mode", choices=["colocate","server"], default="colocate") p.add_argument("--vllm-url", default="http://localhost:8000", help="vLLM server URL (if --vllm-mode=server)") p.add_argument("--push-to-hub", action="store_true") return p.parse_args() # -------------------------------------------------------------------------- # System prompt the LLM sees as the "logistics coordinator" # -------------------------------------------------------------------------- SYSTEM_PROMPT = """ You are an AI Logistics Coordinator managing a fleet of shipments under active disruption. Your goal is to minimise delivery delays, maintain SLA compliance, communicate proactively with customers, and avoid unnecessary escalations. ## YOUR TOOLS You interact with the environment via MCP tool calls. Available tools: - get_network_status() → See all shipments, disruptions, route options - reroute_shipment(id, route, carrier, reason) → Switch a delayed shipment to a better route - set_priority([ids]) → Fast-track up to 3 critical shipments - communicate_eta(id, message) → Send customer ETA update (graded for quality) - escalate(id, reason) → Flag for human dispatcher (-0.1 reward each!) - end_turn() → Commit all decisions and receive your reward ## STRATEGY 1. Always call get_network_status() first. 2. Re-route shipments with negative sla_buffer_hours away from congested routes. 3. Prioritise high-value or perishable cargo. 4. Send clear, empathetic ETA updates to delayed customers. 5. Only escalate if truly unresolvable — each escalation costs -0.1 reward. 6. Always end with end_turn() to commit your actions. - Delay Reduction: 40% - SLA Compliance: 30% - Communication Quality: 20% - Escalation Control: 10% ## MULTI-AGENT WARNING Routes now have limited capacity (Theme #1). Check `route_load` in the network status. If a route is > 85% loaded, rerouting to it will fail. """.strip() # -------------------------------------------------------------------------- # Helpers # -------------------------------------------------------------------------- def build_user_prompt(turn: int, network_status: dict) -> str: """Format the current turn's situation as a user message.""" delayed = [s for s in network_status.get("shipments", []) if s["sla_buffer_hours"] < 0] disruptions = network_status.get("disruptions", []) load = network_status.get("route_load", {}) load_str = ", ".join([f"{r}: {v*100:.0f}%" for r, v in load.items()]) return ( f"Turn {turn}/{network_status.get('max_turns', 5)}. " f"Network Load: {load_str}. " f"Disruptions: {'; '.join(disruptions[:2])}. " f"Delayed shipments: {[s['id'] for s in delayed]}. " f"Cumulative reward so far: {network_status.get('cumulative_reward', 0.0):.3f}. " "What are your next actions? Use the MCP tools to respond." ) def rollout_once(trainer, env, tokenizer, system_prompt: str, max_turns: int) -> dict: """ Run one full episode of the Logistics environment. Returns token ids, logprobs, and reward signals for GRPO. """ from trl.experimental.openenv import generate_rollout_completions env.reset() prompt_ids_all, completion_ids_all, logprobs_all = [], [], [] turn_rewards, sla_rewards, comm_rewards = [], [], [] for turn in range(max_turns): # Get current state status = env.call_tool("get_network_status") if status.get("turns_remaining", 1) == 0: break user_prompt = build_user_prompt(turn + 1, status) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] prompt_text = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=False, enable_thinking=False ) # Generate model action rollout = generate_rollout_completions(trainer, [prompt_text])[0] prompt_ids_all.extend(rollout["prompt_ids"]) completion_ids_all.extend(rollout["completion_ids"]) logprobs_all.extend(rollout["logprobs"]) action_text = rollout.get("text") or tokenizer.decode( rollout["completion_ids"], skip_special_tokens=True ) # Parse tool calls from the LLM's text (simple heuristic) _execute_tool_calls(env, action_text) # Commit the turn result = env.call_tool("end_turn") breakdown = result.get("reward_breakdown", {}) turn_rewards.append(float(result.get("turn_reward", 0.0))) sla_rewards.append(float(breakdown.get("sla_compliance", 0.0))) comm_rewards.append(float(breakdown.get("communication_quality", 0.0))) if result.get("done"): break return { "prompt_ids": prompt_ids_all, "completion_ids": completion_ids_all, "logprobs": logprobs_all, "delay_reward": turn_rewards[-1] if turn_rewards else 0.0, "sla_reward": sla_rewards[-1] if sla_rewards else 0.0, "comm_reward": comm_rewards[-1] if comm_rewards else 0.0, } def _execute_tool_calls(env, text: str) -> None: """ Naive tool-call parser: looks for known tool names in the LLM's output and calls them. For production, use the MCP structured output. """ text_lower = text.lower() # If LLM mentions rerouting, attempt it for first delayed shipment if "reroute" in text_lower or "r2" in text_lower: try: env.call_tool( "reroute_shipment", shipment_id="SHIP-001", new_route="R2", new_carrier="SpeedLane", reason="Extracted from model output: avoid congested R1", ) except Exception: pass # If LLM mentions customer communication if any(w in text_lower for w in ["eta", "reschedul", "sorry", "apologis", "delay"]): try: # Find first delayed shipment from the status env.call_tool( "communicate_eta", shipment_id="SHIP-001", message=text[:300], # Use the model's own words ) except Exception: pass # If LLM mentions priority if "priority" in text_lower or "ship-003" in text_lower: try: env.call_tool("set_priority", shipment_ids=["SHIP-001", "SHIP-003"]) except Exception: pass # -------------------------------------------------------------------------- # Reward functions (TRL GRPO format) # -------------------------------------------------------------------------- def reward_delay(completions: list, **kwargs) -> list[float]: """Primary reward: delay hours saved.""" return [float(r) for r in kwargs.get("delay_reward", [0.0] * len(completions))] def reward_sla(completions: list, **kwargs) -> list[float]: """Secondary reward: SLA compliance rate.""" return [float(r) for r in kwargs.get("sla_reward", [0.0] * len(completions))] def reward_communication(completions: list, **kwargs) -> list[float]: """Tertiary reward: communication quality.""" return [float(r) for r in kwargs.get("comm_reward", [0.0] * len(completions))] # -------------------------------------------------------------------------- # Main # -------------------------------------------------------------------------- def main() -> None: args = parse_args() # Lazy imports so the file can be read without GPU deps from datasets import Dataset from transformers import AutoTokenizer from trl import GRPOConfig, GRPOTrainer sys.path.insert(0, str(Path(__file__).parent.parent / "src")) from envs.logistics_shipment_env import LogisticsShipmentEnv print(f"🚛 Logistics Shipment GRPO Training") print(f" Model: {args.model}") print(f" Env: {args.env_url}") print(f" Dataset: {args.dataset_size} prompts × {args.num_gen} rollouts") # Tokenizer tokenizer = AutoTokenizer.from_pretrained(args.model) tokenizer.pad_token = tokenizer.eos_token # Environment client (sync wrapper) raw_env = LogisticsShipmentEnv(base_url=args.env_url) env = raw_env.sync().__enter__() # Dataset — each row is a prompt seeded with the task description dataset_prompt = ( "You are managing a logistics network with 4 active shipments under disruption. " "Minimise delays, maintain SLA, and communicate clearly with customers." ) dataset = Dataset.from_dict({"prompt": [dataset_prompt] * args.dataset_size}) # GRPO config grpo_config = GRPOConfig( num_train_epochs=args.epochs, learning_rate=args.lr, gradient_accumulation_steps=args.grad_accum, per_device_train_batch_size=1, warmup_steps=10, num_generations=args.num_gen, max_completion_length=512, max_prompt_length=1024, use_vllm=True, vllm_mode=args.vllm_mode, vllm_server_base_url=args.vllm_url if args.vllm_mode == "server" else None, output_dir=args.output_dir, report_to="none", logging_steps=1, save_steps=20, push_to_hub=args.push_to_hub, ) # Rollout function def rollout_func(prompts: list[str], trainer: GRPOTrainer) -> dict: all_prompt_ids, all_comp_ids, all_logprobs = [], [], [] d_rewards, s_rewards, c_rewards = [], [], [] for _ in prompts: ep = rollout_once(trainer, env, tokenizer, SYSTEM_PROMPT, args.max_turns) all_prompt_ids.append(ep["prompt_ids"]) all_comp_ids.append(ep["completion_ids"]) all_logprobs.append(ep["logprobs"]) d_rewards.append(ep["delay_reward"]) s_rewards.append(ep["sla_reward"]) c_rewards.append(ep["comm_reward"]) return { "prompt_ids": all_prompt_ids, "completion_ids": all_comp_ids, "logprobs": all_logprobs, "delay_reward": d_rewards, "sla_reward": s_rewards, "comm_reward": c_rewards, } # Trainer trainer = GRPOTrainer( model=args.model, processing_class=tokenizer, reward_funcs=[reward_delay, reward_sla, reward_communication], train_dataset=dataset, args=grpo_config, rollout_func=rollout_func, ) print("\n🏋️ Starting GRPO training...\n") try: trainer.train() trainer.save_model(args.output_dir) if args.push_to_hub: trainer.push_to_hub() print(f"\n✅ Training complete. Model saved to {args.output_dir}") finally: env.__exit__(None, None, None) if __name__ == "__main__": main()