{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}, "accelerator": "GPU", "colab": {"gpuType": "T4"} }, "cells": [ { "cell_type": "markdown", "id": "cell-0", "metadata": {}, "source": [ "# OrgOS — GRPO Training on a Multi-App Enterprise RL Environment\n", "\n", "**Project:** OrgOS — an OpenEnv environment that simulates enterprise workflows across **Jira, Zendesk, Salesforce, and Workday** with realistic challenges: schema drift, RBAC, SLA constraints, and policy drift.\n", "\n", "**Goal of this notebook:** Fine-tune `Qwen2.5-3B-Instruct` with **GRPO** (Group Relative Policy Optimization) using **live environment rewards**, then compare the trained agent against the untrained baseline.\n", "\n", "**Hardware:** Colab T4 (free tier, 16 GB VRAM). End-to-end runtime ≈ 45–60 min.\n", "\n", "**Outputs (committed to the repo):**\n", "- `training/training_log.txt` — structured logs (`[TRAIN_CONFIG]`, `[EVAL]`, `[TRAIN_STEP]`, …)\n", "- `training/plots/training_curve.png` — mean reward vs GRPO step\n", "- `training/plots/baseline_vs_trained.png` — per-workflow comparison\n", "- `training/plots/score_distribution.png` — per-episode score distribution\n", "- `training/orgos_lora_adapter/` — trained LoRA weights\n", "\n", "Reviewers can open this notebook on Colab → Runtime → *Run all* and reproduce every artifact end-to-end." ] }, { "cell_type": "markdown", "id": "cell-1", "metadata": {}, "source": ["## 1. Setup — install dependencies and clone the repo"] }, { "cell_type": "code", "id": "cell-2", "metadata": {}, "outputs": [], "source": [ "# Pin TRL to the version Unsloth requires BEFORE installing unsloth.\n", "# trl 1.x breaks Unsloth's GRPOTrainer patches — keep it <=0.24.\n", "%pip install -q \"trl>=0.18.2,<=0.24.0\" peft accelerate bitsandbytes datasets\n", "# Install Unsloth after TRL so its patches apply to the right TRL version.\n", "%pip install -q --upgrade unsloth\n", "%pip install -q fastapi 'uvicorn[standard]' pydantic httpx faker openai aiofiles" ] }, { "cell_type": "code", "id": "cell-3", "metadata": {}, "outputs": [], "source": [ "# Clone the OrgOS dev repo (env server, models, business rules)\n", "import os\n", "REPO_URL = 'https://github.com/Tanvi51204/OpenEnv-Round-2.git'\n", "if not os.path.isdir('/content/OpenEnv-Round-2'):\n", " !git clone {REPO_URL} /content/OpenEnv-Round-2\n", "%cd /content/OpenEnv-Round-2" ] }, { "cell_type": "code", "id": "cell-4", "metadata": {}, "outputs": [], "source": [ "# Imports — keep `import unsloth` first to register its patches.\n", "import unsloth\n", "\n", "import json, os, re, sys, time, subprocess\n", "from pathlib import Path\n", "from typing import List\n", "\n", "import httpx\n", "import numpy as np\n", "import torch\n", "import matplotlib.pyplot as plt\n", "from datasets import Dataset\n", "from transformers import TrainerCallback\n", "from trl import GRPOConfig, GRPOTrainer\n", "from unsloth import FastLanguageModel\n", "\n", "torch.set_float32_matmul_precision('high')" ] }, { "cell_type": "markdown", "id": "cell-5", "metadata": {}, "source": ["## 2. Configuration"] }, { "cell_type": "code", "id": "cell-6", "metadata": {}, "outputs": [], "source": [ "# ---- Model ----\n", "MODEL_NAME = 'unsloth/Qwen2.5-3B-Instruct-bnb-4bit'\n", "MAX_SEQ_LEN = 4096\n", "LORA_R = 16\n", "LORA_ALPHA = 16\n", "\n", "# ---- Environment ----\n", "ENV_URL = 'http://localhost:8000'\n", "WORKFLOWS = ['A', 'B', 'C']\n", "\n", "# ---- Data / eval ----\n", "N_PROMPTS_PER_WORKFLOW = 20 # 20 × 3 = 60 prompts\n", "N_EVAL_EPISODES = 5 # episodes per workflow at eval time\n", "MAX_EPISODE_STEPS = 15\n", "\n", "# ---- GRPO ----\n", "MAX_TRAIN_STEPS = 150 # more steps for better convergence\n", "NUM_GENERATIONS = 2 # G = candidates per prompt (keep low for T4 VRAM)\n", "PER_DEVICE_BATCH = 1\n", "GRAD_ACCUM = 2 # effective batch = 2 with grad accum\n", "LEARNING_RATE = 8e-6\n", "MAX_COMPLETION_LENGTH = 256\n", "REWARD_STEPS = 2 # multi-step rollout depth in reward fn\n", "\n", "# ---- Drive checkpoint (saves every N steps so Colab disconnects don't lose progress) ----\n", "CKPT_EVERY_STEPS = 30\n", "\n", "# ---- Output paths ----\n", "TRAIN_DIR = Path('/content/OpenEnv-Round-2/training')\n", "PLOTS_DIR = TRAIN_DIR / 'plots'\n", "ADAPTER_DIR = TRAIN_DIR / 'orgos_lora_adapter'\n", "LOG_PATH = TRAIN_DIR / 'training_log.txt'\n", "PLOTS_DIR.mkdir(parents=True, exist_ok=True)" ] }, { "cell_type": "code", "id": "cell-7", "metadata": {}, "outputs": [], "source": [ "# Structured logger — every important event goes through this so submission has a clean log.\n", "LOG_PATH.write_text('') # truncate\n", "\n", "def tlog(line: str):\n", " print(line, flush=True)\n", " with open(LOG_PATH, 'a') as f:\n", " f.write(line + '\\n')" ] }, { "cell_type": "markdown", "id": "cell-8", "metadata": {}, "source": [ "## 3. Start the OrgOS environment server\n", "\n", "We launch the FastAPI env server (`server/app.py`) as a background subprocess. The reward function and eval loop call it over HTTP at `localhost:8000`." ] }, { "cell_type": "code", "id": "cell-9", "metadata": {}, "outputs": [], "source": [ "ENV_PROC = subprocess.Popen(\n", " [sys.executable, '-m', 'uvicorn', 'server.app:app', '--host', '0.0.0.0', '--port', '8000'],\n", " cwd='/content/OpenEnv-Round-2',\n", " stdout=subprocess.DEVNULL,\n", " stderr=subprocess.DEVNULL,\n", ")\n", "for _ in range(30):\n", " try:\n", " r = httpx.get(f'{ENV_URL}/health', timeout=2)\n", " if r.status_code == 200:\n", " tlog(f\"[ENV] status={r.json().get('status')} version={r.json().get('version','?')}\")\n", " break\n", " except Exception:\n", " time.sleep(1)\n", "else:\n", " raise RuntimeError('Env server failed to start')" ] }, { "cell_type": "markdown", "id": "cell-10", "metadata": {}, "source": ["## 4. Load model — Qwen2.5-3B-Instruct, 4-bit, with LoRA adapters"] }, { "cell_type": "code", "id": "cell-11", "metadata": {}, "outputs": [], "source": [ "model, tokenizer = FastLanguageModel.from_pretrained(\n", " model_name = MODEL_NAME,\n", " max_seq_length = MAX_SEQ_LEN,\n", " dtype = None,\n", " load_in_4bit = True,\n", ")\n", "\n", "model = FastLanguageModel.get_peft_model(\n", " model,\n", " r = LORA_R,\n", " lora_alpha = LORA_ALPHA,\n", " target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n", " use_gradient_checkpointing = 'unsloth',\n", ")\n", "\n", "# Clear max_length so generate() doesn't warn about max_new_tokens vs max_length conflict.\n", "model.config.max_length = None\n", "\n", "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "tlog(f'[TRAIN_CONFIG] model={MODEL_NAME} lora_r={LORA_R} max_seq_len={MAX_SEQ_LEN} '\n", " f'trainable_params={trainable:,} quantization=4bit')" ] }, { "cell_type": "markdown", "id": "cell-12", "metadata": {}, "source": [ "## 5. Helpers — system prompt, observation formatting, action parsing\n", "\n", "The agent gets a **stateless single-turn prompt**: `[SYSTEM_PROMPT] + [observation]` → `[action JSON]`. This matches what GRPO trains on, which is critical for eval/train alignment, and prevents context accumulation over a multi-step episode." ] }, { "cell_type": "code", "id": "cell-13", "metadata": {}, "outputs": [], "source": [ "SYSTEM_PROMPT = '''You are OrgOS Agent — an enterprise workflow automation agent.\n", "You operate across four SaaS apps: Jira, Zendesk, Salesforce, and Workday.\n", "\n", "Each turn you receive a JSON observation with workflow_goal, pending_steps, app_states,\n", "schema_hints (field renames in effect this episode, e.g. {\"jira.priority\": \"severity\"}),\n", "active_rules, message (feedback from last action), and current_score.\n", "\n", "Respond ONLY with a valid JSON object — no markdown, no explanation:\n", " {\"app\": \"\", \"operation\": \"\", \"args\": {...}}\n", "\n", "Available apps and key operations:\n", " jira: get_issue, create_issue, update_status, set_priority, assign_owner,\n", " add_label, link_zendesk_ticket, close_issue, list_issues\n", " zendesk: get_ticket, acknowledge_ticket, set_urgency, assign_agent,\n", " escalate_to_jira, resolve_ticket, add_note, list_tickets, create_agent_profile\n", " salesforce: get_account, list_accounts, update_deal_stage, flag_churn_risk,\n", " assign_account_owner, log_interaction, get_opportunity\n", " workday: get_employee, list_employees, provision_access, log_sla_event,\n", " request_budget_approval, create_onboarding_task, complete_task\n", "\n", "CRITICAL RULES:\n", "1. Read schema_hints FIRST. If \"salesforce.owner\" -> \"rep_email\", use \"rep_email\" not \"owner\".\n", "2. Complete pending_steps in order.\n", "3. Never repeat a failed action unchanged — read the message and adapt.\n", "4. Use list_* operations to discover record IDs.\n", "5. Stop when pending_steps is empty or done=true.'''" ] }, { "cell_type": "code", "id": "cell-14", "metadata": {}, "outputs": [], "source": [ "WORKFLOW_APPS = {\n", " 'A': {'jira', 'zendesk', 'salesforce', 'workday'},\n", " 'B': {'zendesk', 'salesforce', 'workday'},\n", " 'C': {'jira', 'zendesk', 'salesforce'},\n", "}\n", "\n", "def obs_to_text(obs: dict) -> str:\n", " hints = obs.get('schema_hints', {})\n", " pending = obs.get('pending_steps', [])\n", " lines = [\n", " f\"current_score: {obs['current_score']}\",\n", " f\"step_count: {obs['step_count']}\",\n", " f\"workflow_id: {obs['workflow_id']}\",\n", " '',\n", " '=== WORKFLOW GOAL ===' , obs['workflow_goal'], '',\n", " '=== PENDING STEPS ===',\n", " '\\n'.join(f' - {s}' for s in pending) or ' (all steps complete!)',\n", " '',\n", " '=== SCHEMA HINTS (use these field names) ===',\n", " json.dumps(hints, indent=2) if hints else ' (no drift — use canonical names)',\n", " '',\n", " '=== ACTIVE RULES ===',\n", " json.dumps(obs.get('active_rules', {}), indent=2),\n", " '',\n", " '=== LAST MESSAGE ===', obs['message'], '',\n", " '=== APP STATES ===',\n", " ]\n", " relevant = WORKFLOW_APPS.get(obs.get('workflow_id', 'A'),\n", " {'jira','zendesk','salesforce','workday'})\n", " for app_name, view in obs.get('app_states', {}).items():\n", " if app_name not in relevant:\n", " continue\n", " view_str = str(view)\n", " if len(view_str) > 600:\n", " view_str = view_str[:600] + '...[truncated]'\n", " lines += [f' [{app_name.upper()}]', f' {view_str}', '']\n", " return '\\n'.join(lines)\n", "\n", "def parse_action(text: str):\n", " text = re.sub(r'```(?:json)?\\s*', '', text.strip()).strip()\n", " try:\n", " return json.loads(text)\n", " except json.JSONDecodeError:\n", " m = re.search(r'\\{.*\\}', text, re.DOTALL)\n", " if m:\n", " try: return json.loads(m.group())\n", " except Exception: return None\n", " return None\n", "\n", "def build_prompt(obs_text: str) -> str:\n", " msgs = [{'role': 'user', 'content': SYSTEM_PROMPT + '\\n\\n---\\n\\n' + obs_text}]\n", " return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)" ] }, { "cell_type": "markdown", "id": "cell-15", "metadata": {}, "source": [ "## 6. Episode runner & evaluator\n", "\n", "`run_episode_with_model` is **stateless** — each step sends just `[system + current obs]`, no chat history. This (a) keeps prompts under `MAX_SEQ_LEN`, (b) matches the GRPO training format exactly, and (c) avoids context accumulation across multi-step episodes." ] }, { "cell_type": "code", "id": "cell-16", "metadata": {}, "outputs": [], "source": [ "def run_episode_with_model(workflow_id: str, max_steps: int = MAX_EPISODE_STEPS) -> float:\n", " obs = httpx.post(f'{ENV_URL}/reset', json={'workflow_id': workflow_id}).json()['observation']\n", " for _ in range(max_steps):\n", " if obs['done']:\n", " break\n", " prompt = build_prompt(obs_to_text(obs))\n", " inputs = tokenizer(prompt, return_tensors='pt').to(model.device)\n", " with torch.no_grad():\n", " out = model.generate(\n", " **inputs,\n", " max_new_tokens = 256,\n", " do_sample = False,\n", " pad_token_id = tokenizer.eos_token_id,\n", " )\n", " action_str = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:],\n", " skip_special_tokens=True).strip()\n", " action = parse_action(action_str)\n", " if action is None:\n", " break\n", " result = httpx.post(f'{ENV_URL}/step', json=action).json()\n", " obs = result['observation']\n", " if obs['done']:\n", " break\n", " return float(obs.get('current_score', 0.001))\n", "\n", "def evaluate(phase: str, n_eval: int = N_EVAL_EPISODES) -> dict:\n", " scores = {wf: [] for wf in WORKFLOWS}\n", " tlog(f'[EVAL_START] phase={phase}')\n", " for wf in WORKFLOWS:\n", " for ep in range(n_eval):\n", " s = run_episode_with_model(wf)\n", " scores[wf].append(s)\n", " tlog(f'[EVAL] phase={phase} workflow={wf} episode={ep+1} score={s:.4f}')\n", " m = float(np.mean(scores[wf]))\n", " tlog(f'[EVAL_WORKFLOW] phase={phase} workflow={wf} '\n", " f'mean={m:.4f} min={min(scores[wf]):.4f} max={max(scores[wf]):.4f}')\n", " overall = float(np.mean([s for v in scores.values() for s in v]))\n", " tlog(f'[EVAL_END] phase={phase} overall_mean={overall:.4f}')\n", " return scores" ] }, { "cell_type": "markdown", "id": "cell-17", "metadata": {}, "source": [ "## 7. Baseline evaluation — *before* training\n", "\n", "This is the untrained Qwen2.5-3B reference point. We will compare against this after GRPO training." ] }, { "cell_type": "code", "id": "cell-18", "metadata": {}, "outputs": [], "source": [ "FastLanguageModel.for_inference(model)\n", "baseline_scores = evaluate(phase='baseline')\n", "baseline_overall = float(np.mean([s for v in baseline_scores.values() for s in v]))" ] }, { "cell_type": "markdown", "id": "cell-19", "metadata": {}, "source": [ "## 8. Build the prompt dataset for GRPO\n", "\n", "We collect 60 fresh observations (20 per workflow) by resetting the env. GRPO will sample from this dataset, generate G=2 candidate actions per prompt, score each via the live env, and update the policy from the relative advantages." ] }, { "cell_type": "code", "id": "cell-20", "metadata": {}, "outputs": [], "source": [ "rows = []\n", "for wf in WORKFLOWS:\n", " for _ in range(N_PROMPTS_PER_WORKFLOW):\n", " obs = httpx.post(f'{ENV_URL}/reset', json={'workflow_id': wf}).json()['observation']\n", " rows.append({\n", " 'prompt': build_prompt(obs_to_text(obs)),\n", " 'workflow_id': wf,\n", " })\n", "prompt_dataset = Dataset.from_list(rows)\n", "tlog(f'[TRAIN_CONFIG] algorithm=GRPO prompts={len(prompt_dataset)} '\n", " f'workflows={\",\".join(WORKFLOWS)} prompts_per_workflow={N_PROMPTS_PER_WORKFLOW}')\n", "\n", "tok_len = len(tokenizer(prompt_dataset[0]['prompt']).input_ids)\n", "tlog(f'[PROMPT_DEBUG] first_prompt_tokens={tok_len}')" ] }, { "cell_type": "markdown", "id": "cell-21", "metadata": {}, "source": [ "## 9. Reward function — multi-step live environment rollout\n", "\n", "For each candidate completion we:\n", "1. Parse it as a JSON action.\n", "2. Reset the env and apply the action (step 1).\n", "3. Continue `REWARD_STEPS-1` more steps with the current model (greedy), accumulating env transitions.\n", "4. Return the **cumulative episode score** — not just single-step reward.\n", "\n", "This multi-step signal prevents the model from collapsing to always outputting `list_tickets` (which gives a small single-step reward but doesn't advance the workflow). Invalid JSON gets a −0.1 penalty." ] }, { "cell_type": "code", "id": "cell-22", "metadata": {}, "outputs": [], "source": [ "def orgos_reward_fn(completions: List[str], prompts: List[str] = None, **kwargs) -> List[float]:\n", " workflow_ids = kwargs.get('workflow_id', ['A'] * len(completions))\n", " rewards = []\n", " for completion, wf_id in zip(completions, workflow_ids):\n", " action = parse_action(completion)\n", " if action is None:\n", " rewards.append(-0.1)\n", " continue\n", " try:\n", " # Reset env and apply GRPO-generated action (step 1)\n", " obs = httpx.post(f'{ENV_URL}/reset', json={'workflow_id': wf_id}, timeout=10).json()['observation']\n", " result = httpx.post(f'{ENV_URL}/step', json=action, timeout=10).json()\n", " obs = result['observation']\n", "\n", " # Continue REWARD_STEPS-1 more steps with the current model (greedy)\n", " for _ in range(REWARD_STEPS - 1):\n", " if obs.get('done'):\n", " break\n", " prompt = build_prompt(obs_to_text(obs))\n", " inputs = tokenizer(prompt, return_tensors='pt').to(model.device)\n", " with torch.no_grad():\n", " out = model.generate(\n", " **inputs,\n", " max_new_tokens = 128,\n", " do_sample = False,\n", " pad_token_id = tokenizer.eos_token_id,\n", " )\n", " cont_str = tokenizer.decode(\n", " out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True\n", " ).strip()\n", " cont_action = parse_action(cont_str)\n", " if cont_action is None:\n", " break\n", " result = httpx.post(f'{ENV_URL}/step', json=cont_action, timeout=10).json()\n", " obs = result['observation']\n", "\n", " # Return cumulative score — rewards multi-step progress, not just single actions\n", " rewards.append(float(obs.get('current_score', 0.001)))\n", " except Exception as e:\n", " rewards.append(-0.1)\n", " return rewards\n", "\n", "# Sanity check\n", "_v = orgos_reward_fn(['{\\'app\\':\\'zendesk\\',\\'operation\\':\\'list_tickets\\',\\'args\\':{}}'], workflow_id=['A'])\n", "_i = orgos_reward_fn(['not json'], workflow_id=['A'])\n", "tlog(f'[REWARD_FN_CHECK] valid_action={_v[0]:.4f} invalid_action={_i[0]:.4f}')" ] }, { "cell_type": "markdown", "id": "cell-23", "metadata": {}, "source": ["## 10. GRPO training\n", "\n", "We log every training step's reward into `[TRAIN_STEP]` lines so we can plot a meaningful learning curve.\n", "A Drive checkpoint callback saves the adapter every 30 steps so a Colab disconnect doesn't lose progress."] }, { "cell_type": "code", "id": "cell-24", "metadata": {}, "outputs": [], "source": [ "training_step_rewards = [] # captured by callback for the plot\n", "\n", "class OrgOSLogCallback(TrainerCallback):\n", " def on_log(self, args, state, control, logs=None, **kwargs):\n", " if not logs:\n", " return\n", " step = state.global_step\n", " reward = logs.get('reward') or logs.get('rewards/orgos_reward_fn') or logs.get('reward/mean')\n", " loss = logs.get('loss')\n", " kl = logs.get('kl')\n", " if reward is not None:\n", " training_step_rewards.append((step, float(reward)))\n", " tlog(f'[TRAIN_STEP] step={step} reward={float(reward):.4f} '\n", " f\"loss={('%.4f'%loss) if loss is not None else 'NA'} \"\n", " f\"kl={('%.4f'%kl) if kl is not None else 'NA'}\")\n", "\n", " def on_step_end(self, args, state, control, **kwargs):\n", " \"\"\"Save checkpoint to Drive every CKPT_EVERY_STEPS steps.\"\"\"\n", " if state.global_step % CKPT_EVERY_STEPS == 0 and state.global_step > 0:\n", " try:\n", " from google.colab import drive\n", " drive.mount('/content/drive', force_remount=False)\n", " ckpt_dir = Path(f'/content/drive/MyDrive/orgos-training/ckpt_step{state.global_step}')\n", " ckpt_dir.mkdir(parents=True, exist_ok=True)\n", " model.save_pretrained(str(ckpt_dir))\n", " import shutil\n", " shutil.copy(LOG_PATH, ckpt_dir / 'training_log.txt')\n", " tlog(f'[CHECKPOINT] step={state.global_step} saved to {ckpt_dir}')\n", " except Exception:\n", " pass # not on Colab or Drive not mounted — skip silently" ] }, { "cell_type": "code", "id": "cell-25", "metadata": {}, "outputs": [], "source": [ "FastLanguageModel.for_training(model)\n", "\n", "# GRPOConfig — using TRL <=0.24 (pinned in cell 2) so max_new_tokens is accepted.\n", "# Unsloth patches this config; max_prompt_length / max_completion_length are NOT supported.\n", "grpo_config = GRPOConfig(\n", " output_dir = '/content/grpo_ckpt',\n", " num_train_epochs = 1,\n", " max_steps = MAX_TRAIN_STEPS,\n", " per_device_train_batch_size = PER_DEVICE_BATCH,\n", " gradient_accumulation_steps = GRAD_ACCUM,\n", " learning_rate = LEARNING_RATE,\n", " num_generations = NUM_GENERATIONS,\n", " max_new_tokens = MAX_COMPLETION_LENGTH,\n", " temperature = 0.9,\n", " logging_steps = 1,\n", " save_strategy = 'no',\n", " report_to = 'none',\n", " bf16 = False,\n", " fp16 = True,\n", " optim = 'adamw_8bit',\n", ")\n", "\n", "trainer = GRPOTrainer(\n", " model = model,\n", " processing_class = tokenizer,\n", " reward_funcs = [orgos_reward_fn],\n", " train_dataset = prompt_dataset,\n", " args = grpo_config,\n", " callbacks = [OrgOSLogCallback()],\n", ")\n", "\n", "tlog(f'[TRAIN_START] max_steps={MAX_TRAIN_STEPS} G={NUM_GENERATIONS} lr={LEARNING_RATE} reward_steps={REWARD_STEPS}')\n", "trainer.train()\n", "tlog(f'[TRAIN_END] steps_completed={trainer.state.global_step}')" ] }, { "cell_type": "markdown", "id": "cell-26", "metadata": {}, "source": [ "## 11. Post-training evaluation\n", "\n", "Same protocol as the baseline (3 workflows × 5 episodes), so the comparison is apples-to-apples." ] }, { "cell_type": "code", "id": "cell-27", "metadata": {}, "outputs": [], "source": [ "FastLanguageModel.for_inference(model)\n", "trained_scores = evaluate(phase='trained')\n", "trained_overall = float(np.mean([s for v in trained_scores.values() for s in v]))\n", "\n", "tlog('[TRAIN_SUMMARY] '\n", " f'baseline_overall={baseline_overall:.4f} trained_overall={trained_overall:.4f} '\n", " f'delta={trained_overall - baseline_overall:+.4f}')" ] }, { "cell_type": "markdown", "id": "cell-28", "metadata": {}, "source": [ "## 12. Plots\n", "\n", "All plots are saved to `training/plots/` and committed to the repo so reviewers can see them in the README." ] }, { "cell_type": "code", "id": "cell-29", "metadata": {}, "outputs": [], "source": [ "# 12a. Training curve — mean reward vs GRPO step\n", "if training_step_rewards:\n", " steps, rewards = zip(*training_step_rewards)\n", " plt.figure(figsize=(8,5))\n", " plt.plot(steps, rewards, marker='o', markersize=3, linewidth=1.5, color='tab:blue', label='per-step reward')\n", " if len(rewards) >= 5:\n", " win = max(3, len(rewards)//10)\n", " roll = np.convolve(rewards, np.ones(win)/win, mode='valid')\n", " plt.plot(steps[win-1:], roll, color='tab:orange', linewidth=2.5, label=f'rolling mean (w={win})')\n", " plt.xlabel('GRPO training step')\n", " plt.ylabel('mean reward (per batch)')\n", " plt.title('OrgOS GRPO training curve — Qwen2.5-3B-Instruct')\n", " plt.legend()\n", " plt.grid(alpha=0.3)\n", " plt.tight_layout()\n", " plt.savefig(PLOTS_DIR / 'training_curve.png', dpi=150)\n", " plt.show()\n", " tlog('[ARTIFACT] training_curve.png saved')" ] }, { "cell_type": "code", "id": "cell-30", "metadata": {}, "outputs": [], "source": [ "# 12b. Baseline vs trained — grouped bar per workflow\n", "x = np.arange(len(WORKFLOWS))\n", "width = 0.35\n", "baseline_means = [np.mean(baseline_scores[wf]) for wf in WORKFLOWS]\n", "trained_means = [np.mean(trained_scores[wf]) for wf in WORKFLOWS]\n", "\n", "fig, ax = plt.subplots(figsize=(8,5))\n", "ax.bar(x - width/2, baseline_means, width, label='baseline (untrained)', color='tab:gray')\n", "ax.bar(x + width/2, trained_means, width, label='GRPO-trained', color='tab:green')\n", "ax.set_xticks(x)\n", "ax.set_xticklabels([f'Workflow {wf}' for wf in WORKFLOWS])\n", "ax.set_ylabel('mean episode score (0–1)')\n", "ax.set_ylim(0, 1)\n", "ax.set_title(f'Baseline vs GRPO-trained — overall {baseline_overall:.3f} → {trained_overall:.3f}')\n", "ax.legend()\n", "ax.grid(axis='y', alpha=0.3)\n", "for i, (b, t) in enumerate(zip(baseline_means, trained_means)):\n", " ax.text(i - width/2, b + 0.01, f'{b:.2f}', ha='center', fontsize=9)\n", " ax.text(i + width/2, t + 0.01, f'{t:.2f}', ha='center', fontsize=9)\n", "plt.tight_layout()\n", "plt.savefig(PLOTS_DIR / 'baseline_vs_trained.png', dpi=150)\n", "plt.show()\n", "tlog('[ARTIFACT] baseline_vs_trained.png saved')" ] }, { "cell_type": "code", "id": "cell-31", "metadata": {}, "outputs": [], "source": [ "# 12c. Per-episode score distribution (box plot)\n", "fig, ax = plt.subplots(figsize=(9,5))\n", "data, labels, colors = [], [], []\n", "for wf in WORKFLOWS:\n", " data += [baseline_scores[wf], trained_scores[wf]]\n", " labels += [f'{wf} (base)', f'{wf} (trained)']\n", " colors += ['lightgray', 'lightgreen']\n", "bp = ax.boxplot(data, labels=labels, patch_artist=True)\n", "for patch, c in zip(bp['boxes'], colors):\n", " patch.set_facecolor(c)\n", "ax.set_ylabel('episode score (0–1)')\n", "ax.set_title('Per-episode score distribution — baseline vs GRPO-trained')\n", "ax.grid(axis='y', alpha=0.3)\n", "plt.tight_layout()\n", "plt.savefig(PLOTS_DIR / 'score_distribution.png', dpi=150)\n", "plt.show()\n", "tlog('[ARTIFACT] score_distribution.png saved')" ] }, { "cell_type": "markdown", "id": "cell-32", "metadata": {}, "source": [ "## 13. Save artifacts\n", "\n", "Saves the LoRA adapter and copies all artifacts to Google Drive so they survive a Colab disconnect." ] }, { "cell_type": "code", "id": "cell-33", "metadata": {}, "outputs": [], "source": [ "model.save_pretrained(str(ADAPTER_DIR))\n", "tokenizer.save_pretrained(str(ADAPTER_DIR))\n", "tlog(f'[ARTIFACT] lora_adapter saved to {ADAPTER_DIR}')\n", "\n", "try:\n", " from google.colab import drive\n", " drive.mount('/content/drive', force_remount=False)\n", " DRIVE_DIR = Path('/content/drive/MyDrive/orgos-training')\n", " DRIVE_DIR.mkdir(parents=True, exist_ok=True)\n", " !cp {LOG_PATH} {DRIVE_DIR}/\n", " !cp -r {PLOTS_DIR} {DRIVE_DIR}/\n", " !cp -r {ADAPTER_DIR} {DRIVE_DIR}/\n", " print(f'Artifacts copied to {DRIVE_DIR}')\n", "except ImportError:\n", " print('Not on Colab — skipping Drive copy')" ] }, { "cell_type": "code", "id": "cell-34", "metadata": {}, "outputs": [], "source": [ "# Stop the env server cleanly\n", "ENV_PROC.terminate()\n", "tlog('[RUN_END]')\n", "print('\\nDone. Commit these to the repo:')\n", "print(f' - {LOG_PATH}')\n", "print(f' - {PLOTS_DIR}/*.png')\n", "print(f' - {ADAPTER_DIR}/')" ] } ] }