training script

training/train.py

@@ -0,0 +1,512 @@
+"""
+OrgOS GRPO Training Script
+Equivalent to training/grpo_orgos.ipynb but runs headlessly.
+
+Outputs:
+  training_log.txt       — structured training log for submission
+  before_after_curves.png — score improvement chart
+  orgos_lora_adapter/    — trained LoRA weights
+"""
+
+import datetime
+import json
+import os
+import re
+import subprocess
+import sys
+import time
+from typing import List
+
+import httpx
+import matplotlib
+matplotlib.use("Agg")   # headless — no display needed
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import numpy as np
+import torch
+from datasets import Dataset
+from transformers import TrainerCallback
+from trl import GRPOConfig, GRPOTrainer
+from unsloth import FastLanguageModel
+
+# ------------------------------------------------------------------
+# Config
+# ------------------------------------------------------------------
+
+MODEL_NAME             = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-3B-Instruct")
+ENV_URL                = "http://localhost:8000"
+LOG_FILE               = "training_log.txt"
+N_PROMPTS_PER_WORKFLOW = 20
+N_EVAL                 = 10
+NUM_EPOCHS             = 3
+BATCH_SIZE             = 4
+GRAD_ACCUM             = 2
+LR                     = 5e-5
+NUM_GEN                = 4
+TEMPERATURE            = 0.8
+BETA                   = 0.04
+LORA_R                 = 16
+MAX_SEQ_LEN            = 2048
+
+# ------------------------------------------------------------------
+# Logger
+# ------------------------------------------------------------------
+
+with open(LOG_FILE, "w") as f:
+    f.write(f"# OrgOS GRPO Training Log\n")
+    f.write(f"# Generated: {datetime.datetime.utcnow().isoformat()}Z\n\n")
+
+
+def tlog(line: str) -> None:
+    print(line, flush=True)
+    with open(LOG_FILE, "a") as f:
+        f.write(line + "\n")
+
+
+# ------------------------------------------------------------------
+# Start OrgOS environment server
+# ------------------------------------------------------------------
+
+def start_env_server():
+    print("Starting OrgOS environment server...", flush=True)
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "uvicorn", "server.app:app",
+         "--host", "0.0.0.0", "--port", "8000"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    # Wait until healthy
+    for _ in range(20):
+        time.sleep(2)
+        try:
+            health = httpx.get(f"{ENV_URL}/health", timeout=5).json()
+            if health.get("status") == "healthy":
+                tlog(f"[ENV] status=healthy version={health.get('version', '?')}")
+                return proc
+        except Exception:
+            pass
+    raise RuntimeError("OrgOS server failed to start after 40 seconds")
+
+
+# ------------------------------------------------------------------
+# Model
+# ------------------------------------------------------------------
+
+def load_model():
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name     = MODEL_NAME,
+        max_seq_length = MAX_SEQ_LEN,
+        dtype          = None,
+        load_in_4bit   = True,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r              = LORA_R,
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                          "gate_proj", "up_proj", "down_proj"],
+        lora_alpha     = LORA_R,
+        lora_dropout   = 0,
+        bias           = "none",
+        use_gradient_checkpointing = "unsloth",
+        random_state   = 42,
+    )
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    tlog(f"[TRAIN_CONFIG] model={MODEL_NAME} lora_r={LORA_R} "
+         f"max_seq_len={MAX_SEQ_LEN} trainable_params={trainable:,} quantization=4bit")
+    return model, tokenizer
+
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+SYSTEM_PROMPT = """\
+You are OrgOS Agent — an enterprise workflow automation agent.
+You operate across four SaaS applications: Jira, Zendesk, Salesforce, and Workday.
+
+Each turn you receive a JSON observation with:
+  - workflow_goal    : the task you must complete
+  - pending_steps    : remaining steps in the workflow
+  - app_states       : current state of each app
+  - schema_hints     : field renames in effect this episode (e.g. {"jira.priority": "severity"})
+  - active_rules     : current SLA / approval thresholds
+  - message          : feedback from the last action
+  - current_score    : your cumulative score (0.001-0.999)
+
+Respond ONLY with a valid JSON object — no markdown, no explanation.
+
+Action format:
+  {"app": "<app>", "operation": "<op>", "args": {...}}
+
+Available apps and key operations:
+  jira:       get_issue, create_issue, update_status, set_priority, assign_owner,
+              add_label, link_zendesk_ticket, close_issue, list_issues
+  zendesk:    get_ticket, acknowledge_ticket, set_urgency, assign_agent,
+              escalate_to_jira, resolve_ticket, add_note, list_tickets,
+              create_agent_profile
+  salesforce: get_account, list_accounts, update_deal_stage, flag_churn_risk,
+              assign_account_owner, log_interaction, get_opportunity
+  workday:    get_employee, list_employees, provision_access, log_sla_event,
+              request_budget_approval, create_onboarding_task, complete_task
+
+CRITICAL RULES:
+1. Read schema_hints FIRST — if "jira.priority" -> "severity", use "severity" not "priority" in args.
+2. Complete ALL pending_steps in order.
+3. Do not repeat a successful action.
+4. If an operation fails, read the message carefully and adapt.
+5. Use list_* operations to discover record IDs when needed.
+6. Stop when pending_steps is empty or done=true.
+"""
+
+
+def obs_to_text(obs: dict) -> str:
+    hints   = obs.get("schema_hints", {})
+    pending = obs.get("pending_steps", [])
+    lines = [
+        f"current_score: {obs['current_score']}",
+        f"step_count:    {obs['step_count']}",
+        f"workflow_id:   {obs['workflow_id']}",
+        "",
+        "=== WORKFLOW GOAL ===",
+        obs["workflow_goal"],
+        "",
+        "=== PENDING STEPS ===",
+        "\n".join(f"  - {s}" for s in pending) or "  (all steps complete!)",
+        "",
+        "=== SCHEMA HINTS (use these field names) ===",
+        json.dumps(hints, indent=2) if hints else "  (no drift — use canonical names)",
+        "",
+        "=== ACTIVE RULES ===",
+        json.dumps(obs.get("active_rules", {}), indent=2),
+        "",
+        "=== LAST MESSAGE ===",
+        obs["message"],
+        "",
+        "=== APP STATES ===",
+    ]
+    for app_name, view in obs.get("app_states", {}).items():
+        lines.append(f"  [{app_name.upper()}]")
+        lines.append(f"  {view}")
+        lines.append("")
+    return "\n".join(lines)
+
+
+def parse_action(text: str):
+    text = re.sub(r"```(?:json)?\s*", "", text.strip()).strip()
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        m = re.search(r"\{.*\}", text, re.DOTALL)
+        if m:
+            try:
+                return json.loads(m.group())
+            except Exception:
+                pass
+    return None
+
+
+def build_prompt(obs_text: str, tokenizer) -> str:
+    messages = [{"role": "user", "content": SYSTEM_PROMPT + "\n\n---\n\n" + obs_text}]
+    return tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+
+# ------------------------------------------------------------------
+# Prompt dataset
+# ------------------------------------------------------------------
+
+def build_prompt_dataset(tokenizer) -> Dataset:
+    rows = []
+    print("Collecting prompts from env resets...", flush=True)
+    for wf in ["A", "B", "C"]:
+        for _ in range(N_PROMPTS_PER_WORKFLOW):
+            result   = httpx.post(f"{ENV_URL}/reset", json={"workflow_id": wf}).json()
+            obs      = result["observation"]
+            obs_text = obs_to_text(obs)
+            rows.append({
+                "prompt":      build_prompt(obs_text, tokenizer),
+                "workflow_id": wf,
+                "obs_text":    obs_text,
+            })
+    tlog(f"[TRAIN_CONFIG] algorithm=GRPO prompts={len(rows)} "
+         f"workflows=A,B,C prompts_per_workflow={N_PROMPTS_PER_WORKFLOW}")
+    return Dataset.from_list(rows)
+
+
+# ------------------------------------------------------------------
+# Reward function
+# ------------------------------------------------------------------
+
+def orgos_reward_fn(completions: List[str], prompts: List[str], **kwargs) -> List[float]:
+    workflow_ids = kwargs.get("workflow_id", ["A"] * len(completions))
+    rewards = []
+    for completion, wf_id in zip(completions, workflow_ids):
+        action = parse_action(completion)
+        if action is None:
+            rewards.append(-0.1)
+            continue
+        try:
+            httpx.post(f"{ENV_URL}/reset", json={"workflow_id": wf_id}, timeout=10)
+            result = httpx.post(f"{ENV_URL}/step", json=action, timeout=10).json()
+            rewards.append(float(result["reward"]))
+        except Exception:
+            rewards.append(-0.1)
+    return rewards
+
+
+# ------------------------------------------------------------------
+# Episode evaluation
+# ------------------------------------------------------------------
+
+def run_episode_with_model(model, tokenizer, workflow_id: str, max_steps: int = 15) -> float:
+    result  = httpx.post(f"{ENV_URL}/reset", json={"workflow_id": workflow_id}).json()
+    obs     = result["observation"]
+    history = []
+
+    for _ in range(max_steps):
+        if obs["done"]:
+            break
+
+        obs_text = obs_to_text(obs)
+        history.append({"role": "user", "content": obs_text})
+
+        messages    = list(history)
+        messages[0] = {"role": "user",
+                       "content": SYSTEM_PROMPT + "\n\n---\n\n" + messages[0]["content"]}
+
+        text   = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+        with torch.no_grad():
+            out = model.generate(
+                **inputs,
+                max_new_tokens = 256,
+                temperature    = 0.0,
+                do_sample      = False,
+                pad_token_id   = tokenizer.eos_token_id,
+            )
+        action_str = tokenizer.decode(
+            out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
+        ).strip()
+
+        history.append({"role": "assistant", "content": action_str})
+
+        action = parse_action(action_str)
+        if action is None:
+            break
+
+        result = httpx.post(f"{ENV_URL}/step", json=action).json()
+        obs    = result["observation"]
+        if obs["done"]:
+            break
+
+    return obs.get("current_score", 0.001)
+
+
+def evaluate(model, tokenizer, phase: str) -> dict:
+    scores = {wf: [] for wf in ["A", "B", "C"]}
+    tlog(f"[EVAL_START] phase={phase}")
+    for wf in ["A", "B", "C"]:
+        for ep in range(N_EVAL):
+            score = run_episode_with_model(model, tokenizer, wf)
+            scores[wf].append(score)
+            tlog(f"[EVAL] phase={phase} workflow={wf} episode={ep+1} score={score:.4f}")
+        wf_mean = np.mean(scores[wf])
+        tlog(f"[EVAL_WORKFLOW] phase={phase} workflow={wf} "
+             f"mean={wf_mean:.4f} min={min(scores[wf]):.4f} max={max(scores[wf]):.4f}")
+    overall = np.mean([s for v in scores.values() for s in v])
+    tlog(f"[EVAL_END] phase={phase} overall_mean={overall:.4f}")
+    return scores
+
+
+# ------------------------------------------------------------------
+# Plot
+# ------------------------------------------------------------------
+
+def plot_results(baseline_scores: dict, post_scores: dict) -> None:
+    fig = plt.figure(figsize=(14, 8), facecolor="#0f172a")
+    fig.suptitle("OrgOS: Before vs After GRPO Training", fontsize=15,
+                 color="white", fontweight="bold", y=0.98)
+
+    gs     = gridspec.GridSpec(2, 3, figure=fig, hspace=0.45, wspace=0.35)
+    COLORS = {"before": "#f87171", "after": "#34d399", "bg": "#1e293b", "grid": "#334155"}
+    LABELS = {
+        "A": "Workflow A\nCustomer Bug Fix",
+        "B": "Workflow B\nEmployee Onboarding",
+        "C": "Workflow C\nChurn Risk Alert",
+    }
+
+    for col, wf in enumerate(["A", "B", "C"]):
+        ax    = fig.add_subplot(gs[0, col])
+        ax.set_facecolor(COLORS["bg"])
+        ax.grid(color=COLORS["grid"], linewidth=0.5, alpha=0.7)
+        before = baseline_scores[wf]
+        after  = post_scores[wf]
+        delta  = np.mean(after) - np.mean(before)
+        ax.plot(before, color=COLORS["before"], linewidth=1.5, alpha=0.8, label="Before GRPO")
+        ax.plot(after,  color=COLORS["after"],  linewidth=1.5, alpha=0.8, label="After GRPO")
+        ax.axhline(np.mean(before), color=COLORS["before"], linestyle="--", linewidth=1, alpha=0.5)
+        ax.axhline(np.mean(after),  color=COLORS["after"],  linestyle="--", linewidth=1, alpha=0.5)
+        ax.set_title(LABELS[wf] + f"\n(Δ = {delta:+.4f})", color="white", fontsize=9)
+        ax.set_xlabel("Episode", color="#94a3b8", fontsize=8)
+        ax.set_ylabel("Final Score", color="#94a3b8", fontsize=8)
+        ax.tick_params(colors="#64748b", labelsize=7)
+        ax.set_ylim(0, 1)
+        ax.legend(fontsize=7, facecolor="#1e293b", labelcolor="white",
+                  edgecolor="#475569", framealpha=0.8)
+        for spine in ax.spines.values():
+            spine.set_edgecolor("#334155")
+
+    ax_hist = fig.add_subplot(gs[1, :])
+    ax_hist.set_facecolor(COLORS["bg"])
+    ax_hist.grid(color=COLORS["grid"], linewidth=0.5, alpha=0.5, axis="x")
+    all_before = [s for v in baseline_scores.values() for s in v]
+    all_after  = [s for v in post_scores.values() for s in v]
+    bins = np.linspace(0, 1, 25)
+    ax_hist.hist(all_before, bins=bins, color=COLORS["before"], alpha=0.6,
+                 label=f"Before GRPO  (mean={np.mean(all_before):.4f})", edgecolor="none")
+    ax_hist.hist(all_after,  bins=bins, color=COLORS["after"],  alpha=0.6,
+                 label=f"After GRPO   (mean={np.mean(all_after):.4f})", edgecolor="none")
+    ax_hist.axvline(np.mean(all_before), color=COLORS["before"], linestyle="--", linewidth=1.5)
+    ax_hist.axvline(np.mean(all_after),  color=COLORS["after"],  linestyle="--", linewidth=1.5)
+    ax_hist.set_title("Score Distribution Across All Workflows", color="white", fontsize=10)
+    ax_hist.set_xlabel("Final Score", color="#94a3b8", fontsize=9)
+    ax_hist.set_ylabel("Count", color="#94a3b8", fontsize=9)
+    ax_hist.tick_params(colors="#64748b", labelsize=8)
+    ax_hist.legend(fontsize=9, facecolor="#1e293b", labelcolor="white",
+                   edgecolor="#475569", framealpha=0.9)
+    for spine in ax_hist.spines.values():
+        spine.set_edgecolor("#334155")
+
+    plt.savefig("before_after_curves.png", dpi=150, bbox_inches="tight",
+                facecolor="#0f172a", edgecolor="none")
+    plt.close()
+    tlog("[ARTIFACT] file=before_after_curves.png")
+
+
+# ------------------------------------------------------------------
+# Training callback
+# ------------------------------------------------------------------
+
+class OrgOSLogCallback(TrainerCallback):
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if not logs:
+            return
+        step        = state.global_step
+        loss        = logs.get("loss", logs.get("train_loss", "?"))
+        mean_reward = logs.get("reward", logs.get("mean_reward", "?"))
+        kl          = logs.get("kl", logs.get("approx_kl", "?"))
+        lr_now      = logs.get("learning_rate", "?")
+
+        loss_str   = f"{loss:.6f}"        if isinstance(loss, float)        else str(loss)
+        reward_str = f"{mean_reward:.4f}" if isinstance(mean_reward, float) else str(mean_reward)
+        kl_str     = f"{kl:.6f}"          if isinstance(kl, float)          else str(kl)
+        lr_str     = f"{lr_now:.2e}"      if isinstance(lr_now, float)      else str(lr_now)
+
+        tlog(f"[TRAIN_STEP] step={step} loss={loss_str} "
+             f"mean_reward={reward_str} kl={kl_str} lr={lr_str}")
+
+
+# ------------------------------------------------------------------
+# Main
+# ------------------------------------------------------------------
+
+def main():
+    server_proc = start_env_server()
+
+    try:
+        model, tokenizer = load_model()
+
+        prompt_dataset = build_prompt_dataset(tokenizer)
+
+        # Sanity-check reward function
+        test_r = orgos_reward_fn(
+            completions = ['{"app": "zendesk", "operation": "list_tickets", "args": {"state": "new"}}',
+                           "not json"],
+            prompts     = ["", ""],
+            workflow_id = ["A", "A"],
+        )
+        tlog(f"[REWARD_FN_CHECK] valid_action={test_r[0]:.4f} invalid_action={test_r[1]:.4f}")
+
+        # Baseline evaluation
+        FastLanguageModel.for_inference(model)
+        baseline_scores = evaluate(model, tokenizer, phase="baseline")
+        baseline_mean   = np.mean([s for v in baseline_scores.values() for s in v])
+
+        # GRPO training
+        model.train()
+        tlog(f"[TRAIN_CONFIG] epochs={NUM_EPOCHS} batch_size={BATCH_SIZE} "
+             f"grad_accum={GRAD_ACCUM} lr={LR} num_generations={NUM_GEN} "
+             f"temperature={TEMPERATURE} beta_kl={BETA}")
+
+        grpo_config = GRPOConfig(
+            output_dir                  = "./orgos_grpo_ckpt",
+            num_train_epochs            = NUM_EPOCHS,
+            per_device_train_batch_size = BATCH_SIZE,
+            gradient_accumulation_steps = GRAD_ACCUM,
+            learning_rate               = LR,
+            warmup_steps                = 10,
+            logging_steps               = 5,
+            save_steps                  = 100,
+            bf16                        = torch.cuda.is_bf16_supported(),
+            fp16                        = not torch.cuda.is_bf16_supported(),
+            max_grad_norm               = 1.0,
+            num_generations             = NUM_GEN,
+            max_new_tokens              = 256,
+            temperature                 = TEMPERATURE,
+            beta                        = BETA,
+            report_to                   = "none",
+            seed                        = 42,
+        )
+
+        trainer = GRPOTrainer(
+            model            = model,
+            args             = grpo_config,
+            reward_funcs     = orgos_reward_fn,
+            train_dataset    = prompt_dataset,
+            processing_class = tokenizer,
+            callbacks        = [OrgOSLogCallback()],
+        )
+
+        tlog("[TRAIN_START]")
+        train_result = trainer.train()
+        tlog(f"[TRAIN_END] total_steps={train_result.global_step} "
+             f"train_loss={train_result.training_loss:.6f} "
+             f"train_runtime_s={train_result.metrics.get('train_runtime', 0):.1f}")
+
+        # Post-training evaluation
+        FastLanguageModel.for_inference(model)
+        post_scores = evaluate(model, tokenizer, phase="post_training")
+        post_mean   = np.mean([s for v in post_scores.values() for s in v])
+        improvement = post_mean - baseline_mean
+
+        tlog(
+            f"[TRAIN_SUMMARY] "
+            f"model={MODEL_NAME} algorithm=GRPO "
+            f"baseline_mean={baseline_mean:.4f} "
+            f"post_training_mean={post_mean:.4f} "
+            f"improvement={improvement:+.4f} "
+            f"workflow_A_before={np.mean(baseline_scores['A']):.4f} "
+            f"workflow_A_after={np.mean(post_scores['A']):.4f} "
+            f"workflow_B_before={np.mean(baseline_scores['B']):.4f} "
+            f"workflow_B_after={np.mean(post_scores['B']):.4f} "
+            f"workflow_C_before={np.mean(baseline_scores['C']):.4f} "
+            f"workflow_C_after={np.mean(post_scores['C']):.4f}"
+        )
+
+        # Save artifacts
+        plot_results(baseline_scores, post_scores)
+        model.save_pretrained("orgos_lora_adapter")
+        tokenizer.save_pretrained("orgos_lora_adapter")
+        tlog("[ARTIFACT] file=orgos_lora_adapter/")
+        tlog("[ARTIFACT] file=training_log.txt")
+
+        print(f"\nDone. Improvement: {baseline_mean:.4f} → {post_mean:.4f} ({improvement:+.4f})")
+
+    finally:
+        server_proc.terminate()
+
+
+if __name__ == "__main__":
+    main()