fix(notebook): pin accelerate==1.0.1 to keep huggingface-hub at 0.26.x

Latest accelerate (1.13+) hard-requires huggingface-hub 1.x. Without
the accelerate pin, STEP 2b's only-if-needed install resolved a newer
accelerate that in turn upgraded hub past 0.26.5, breaking peft 0.14
and trl 0.21 imports. Pin accelerate to the last 0.x-hub-compatible
version (1.0.1) and pull it into STEP 2a alongside the rest of the
training stack.

notebooks/train_merchant_agent.ipynb

@@ -4,13 +4,29 @@
    "cell_type": "markdown",
    "id": "intro-md",
    "metadata": {},
-   "source": "# ChargebackOps Merchant Agent - outcome-based RL on Qwen2.5-3B (fp16 LoRA)\n\nEnd-to-end pipeline for a single Colab T4 (or Kaggle T4):\n\n1. **Phase A - JSON SFT** on heuristic rollouts. Teaches the model the env action schema.\n2. **Phase B - GRPO with outcome reward**. Reward = terminal `$` PnL after the model's action plus heuristic tail-rollout. The merchant policy is pushed toward actions that *win money* against the scripted Issuer + arbitration, not actions that match the heuristic. This is real RLVR (verifiable rewards) - the verifier is the dispute outcome.\n3. **Eval** every checkpoint against the heuristic baseline + naive baseline.\n\n**Why outcome reward over heuristic-match**: heuristic-match training is supervised distillation in disguise - the model can never beat the teacher and the reward is gameable by mimicry. Outcome reward is dollar-denominated, adversarially-verified by the Issuer, and ungameable: the only way to earn it is to actually win disputes.\n\n**Theme alignment**: Multi-Agent (merchant vs scripted Issuer) primary; Professional Tasks (B2B chargeback workflow) and Long-Horizon (multi-round arbitration) secondary.\n\nModel: `Qwen/Qwen2.5-3B-Instruct` fp16 + LoRA r=16. Fits T4 with no bitsandbytes."
+   "source": [
+    "# ChargebackOps Merchant Agent - outcome-based RL on Qwen2.5-3B (fp16 LoRA)\n",
+    "\n",
+    "End-to-end pipeline for a single Colab T4 (or Kaggle T4):\n",
+    "\n",
+    "1. **Phase A - JSON SFT** on heuristic rollouts. Teaches the model the env action schema.\n",
+    "2. **Phase B - GRPO with outcome reward**. Reward = terminal `$` PnL after the model's action plus heuristic tail-rollout. The merchant policy is pushed toward actions that *win money* against the scripted Issuer + arbitration, not actions that match the heuristic. This is real RLVR (verifiable rewards) - the verifier is the dispute outcome.\n",
+    "3. **Eval** every checkpoint against the heuristic baseline + naive baseline.\n",
+    "\n",
+    "**Why outcome reward over heuristic-match**: heuristic-match training is supervised distillation in disguise - the model can never beat the teacher and the reward is gameable by mimicry. Outcome reward is dollar-denominated, adversarially-verified by the Issuer, and ungameable: the only way to earn it is to actually win disputes.\n",
+    "\n",
+    "**Theme alignment**: Multi-Agent (merchant vs scripted Issuer) primary; Professional Tasks (B2B chargeback workflow) and Long-Horizon (multi-round arbitration) secondary.\n",
+    "\n",
+    "Model: `Qwen/Qwen2.5-3B-Instruct` fp16 + LoRA r=16. Fits T4 with no bitsandbytes."
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "setup-md",
    "metadata": {},
-   "source": "## 0. Setup - install deps + clone repo"
+   "source": [
+    "## 0. Setup - install deps + clone repo"
+   ]
   },
   {
    "cell_type": "code",
@@ -18,7 +34,7 @@
    "id": "setup-code",
    "metadata": {},
    "outputs": [],
-   "source": "# GPU + repo setup for Colab T4 (also runs on Kaggle T4).\n# Pin the training stack to the combo that supports both:\n#   - SFTTrainer.compute_loss without the shape-bug regression\n#   - vanilla GRPOTrainer with reward_funcs returning per-completion floats\n# Upper bounds matter: Colab/Kaggle pre-install transformers 5.x + tokenizers\n# 0.22+ + huggingface-hub 1.x which break peft 0.14 and trl 0.21.\nimport os\nimport shutil\nimport subprocess\n\nif os.path.isdir('/content'):\n    WORK_DIR = '/content'\nelif os.path.isdir('/kaggle/working'):\n    WORK_DIR = '/kaggle/working'\nelse:\n    raise RuntimeError('Notebook expects Colab or Kaggle. Set WORK_DIR manually otherwise.')\nos.chdir(WORK_DIR)\nDRIVE_ROOT = '/content/drive/MyDrive'\nPERSIST_ROOT = (\n    os.path.join(DRIVE_ROOT, 'chargebackops-artifacts')\n    if os.path.isdir(DRIVE_ROOT) else WORK_DIR\n)\nos.makedirs(PERSIST_ROOT, exist_ok=True)\nprint('work dir:', WORK_DIR)\nprint('artifact root:', PERSIST_ROOT)\nprint(subprocess.check_output(['nvidia-smi', '-L']).decode())\n\n# STEP 1 - torch trio for cu128. Turing T4 supports cu118+; cu128 wheels work.\nsubprocess.run(\n    ['pip', 'install', '-q', '--no-cache-dir',\n     'torch==2.10.0', 'torchvision==0.25.0', 'torchaudio==2.10.0',\n     '--index-url', 'https://download.pytorch.org/whl/cu128'],\n    check=True, cwd=WORK_DIR,\n)\n\n# STEP 2a - force exact pins for the training stack. `--no-deps` keeps torch\n# untouched. tokenizers + huggingface-hub MUST be pinned explicitly because\n# transformers 4.51 declares ranges that resolve up to incompatible majors\n# when Colab/Kaggle has 0.22+ / 1.x preinstalled.\nsubprocess.run(\n    ['pip', 'install', '-q', '--no-cache-dir', '--force-reinstall', '--no-deps',\n     'transformers==4.51.3',\n     'trl==0.21.0',\n     'peft==0.14.0',\n     'tokenizers==0.21.4',\n     'huggingface-hub==0.26.5'],\n    check=True, cwd=WORK_DIR,\n)\n\n# STEP 2b - supporting libs with only-if-needed so torch + the pinned stack\n# stay put. accelerate >=1.0 plays nice with trl 0.21.\nsubprocess.run(\n    ['pip', 'install', '-q', '--upgrade-strategy=only-if-needed',\n     'accelerate>=1.0,<2.0',\n     'datasets>=2.20,<4.0',\n     'matplotlib>=3.8',\n     'pydantic>=2.10',\n     'openenv-core>=0.2.2'],\n    check=True, cwd=WORK_DIR,\n)\n\n# Clone repo (always fresh) so the editable install matches main.\nREPO_DIR = os.path.join(WORK_DIR, 'chargebackops')\nif os.path.isdir(REPO_DIR):\n    shutil.rmtree(REPO_DIR)\nsubprocess.run(\n    ['git', 'clone', '--depth', '1',\n     'https://github.com/MitudruDutta/ChargeBackOps.git', REPO_DIR],\n    check=True, cwd=WORK_DIR,\n)\nos.chdir(REPO_DIR)\nprint('cwd:', os.getcwd())\n\n# Editable install with --no-deps so pyproject's app/server deps do not alter\n# the training stack we just pinned.\nsubprocess.run(['pip', 'install', '-q', '-e', '.', '--no-deps'],\n               check=True, cwd=REPO_DIR)\n\n# Verify all critical pins land at the exact requested versions.\nimport importlib.metadata as md_\nprint('torch          ', md_.version('torch'))\nprint('torchvision    ', md_.version('torchvision'))\nprint('transformers   ', md_.version('transformers'),    '(want 4.51.3)')\nprint('tokenizers     ', md_.version('tokenizers'),      '(want 0.21.4)')\nprint('huggingface-hub', md_.version('huggingface-hub'), '(want 0.26.5)')\nprint('trl            ', md_.version('trl'),             '(want 0.21.0)')\nprint('peft           ', md_.version('peft'),            '(want 0.14.0)')\nprint('accelerate     ', md_.version('accelerate'))\nprint('openenv-core   ', md_.version('openenv-core'))\nassert md_.version('transformers') == '4.51.3', 'transformers pin failed'\nassert md_.version('trl') == '0.21.0', 'trl pin failed'\nassert md_.version('peft') == '0.14.0', 'peft pin failed'\nassert md_.version('tokenizers').startswith('0.21'), 'tokenizers pin failed'\nassert md_.version('huggingface-hub').startswith('0.26'), 'huggingface-hub pin failed'"
+   "source": "# GPU + repo setup for Colab T4 (also runs on Kaggle T4).\n# Pin the training stack to the combo that supports both:\n#   - SFTTrainer.compute_loss without the shape-bug regression\n#   - vanilla GRPOTrainer with reward_funcs returning per-completion floats\n# Upper bounds matter: Colab/Kaggle pre-install transformers 5.x + tokenizers\n# 0.22+ + huggingface-hub 1.x which break peft 0.14 and trl 0.21.\nimport os\nimport shutil\nimport subprocess\n\nif os.path.isdir('/content'):\n    WORK_DIR = '/content'\nelif os.path.isdir('/kaggle/working'):\n    WORK_DIR = '/kaggle/working'\nelse:\n    raise RuntimeError('Notebook expects Colab or Kaggle. Set WORK_DIR manually otherwise.')\nos.chdir(WORK_DIR)\nDRIVE_ROOT = '/content/drive/MyDrive'\nPERSIST_ROOT = (\n    os.path.join(DRIVE_ROOT, 'chargebackops-artifacts')\n    if os.path.isdir(DRIVE_ROOT) else WORK_DIR\n)\nos.makedirs(PERSIST_ROOT, exist_ok=True)\nprint('work dir:', WORK_DIR)\nprint('artifact root:', PERSIST_ROOT)\nprint(subprocess.check_output(['nvidia-smi', '-L']).decode())\n\n# STEP 1 - torch trio for cu128. Turing T4 supports cu118+; cu128 wheels work.\nsubprocess.run(\n    ['pip', 'install', '-q', '--no-cache-dir',\n     'torch==2.10.0', 'torchvision==0.25.0', 'torchaudio==2.10.0',\n     '--index-url', 'https://download.pytorch.org/whl/cu128'],\n    check=True, cwd=WORK_DIR,\n)\n\n# STEP 2a - force exact pins for the training stack AND accelerate. `--no-deps`\n# keeps torch untouched. accelerate must be pinned here too because the latest\n# accelerate (1.13+) hard-requires huggingface-hub 1.x which breaks peft 0.14\n# and trl 0.21. accelerate 1.0.1 is the last version compatible with hub 0.26.x.\nsubprocess.run(\n    ['pip', 'install', '-q', '--no-cache-dir', '--force-reinstall', '--no-deps',\n     'transformers==4.51.3',\n     'trl==0.21.0',\n     'peft==0.14.0',\n     'accelerate==1.0.1',\n     'tokenizers==0.21.4',\n     'huggingface-hub==0.26.5'],\n    check=True, cwd=WORK_DIR,\n)\n\n# STEP 2b - supporting libs. `--upgrade-strategy=only-if-needed` keeps the\n# stack we just pinned in place (datasets/matplotlib/pydantic/openenv-core\n# don't pull hub or accelerate up).\nsubprocess.run(\n    ['pip', 'install', '-q', '--upgrade-strategy=only-if-needed',\n     'datasets>=2.20,<4.0',\n     'matplotlib>=3.8',\n     'pydantic>=2.10',\n     'openenv-core>=0.2.2'],\n    check=True, cwd=WORK_DIR,\n)\n\n# Clone repo (always fresh) so the editable install matches main.\nREPO_DIR = os.path.join(WORK_DIR, 'chargebackops')\nif os.path.isdir(REPO_DIR):\n    shutil.rmtree(REPO_DIR)\nsubprocess.run(\n    ['git', 'clone', '--depth', '1',\n     'https://github.com/MitudruDutta/ChargeBackOps.git', REPO_DIR],\n    check=True, cwd=WORK_DIR,\n)\nos.chdir(REPO_DIR)\nprint('cwd:', os.getcwd())\n\n# Editable install with --no-deps so pyproject's app/server deps do not alter\n# the training stack we just pinned.\nsubprocess.run(['pip', 'install', '-q', '-e', '.', '--no-deps'],\n               check=True, cwd=REPO_DIR)\n\n# Verify all critical pins land at the exact requested versions.\nimport importlib.metadata as md_\nprint('torch          ', md_.version('torch'))\nprint('torchvision    ', md_.version('torchvision'))\nprint('transformers   ', md_.version('transformers'),    '(want 4.51.3)')\nprint('tokenizers     ', md_.version('tokenizers'),      '(want 0.21.4)')\nprint('huggingface-hub', md_.version('huggingface-hub'), '(want 0.26.5)')\nprint('trl            ', md_.version('trl'),             '(want 0.21.0)')\nprint('peft           ', md_.version('peft'),            '(want 0.14.0)')\nprint('accelerate     ', md_.version('accelerate'),      '(want 1.0.1)')\nprint('openenv-core   ', md_.version('openenv-core'))\nassert md_.version('transformers') == '4.51.3', 'transformers pin failed'\nassert md_.version('trl') == '0.21.0', 'trl pin failed'\nassert md_.version('peft') == '0.14.0', 'peft pin failed'\nassert md_.version('accelerate') == '1.0.1', 'accelerate pin failed'\nassert md_.version('tokenizers').startswith('0.21'), 'tokenizers pin failed'\nassert md_.version('huggingface-hub').startswith('0.26'), 'huggingface-hub pin failed'\n"
   },
   {
    "cell_type": "code",
@@ -26,13 +42,51 @@
    "id": "patch-code",
    "metadata": {},
    "outputs": [],
-   "source": "# Path + module-cache flush so the editable install resolves before any other import.\nimport os, sys, importlib, logging, torch\nos.environ['TOKENIZERS_PARALLELISM'] = 'false'\n\n# Silence transformers per-layer \"Caching is incompatible...\" spam at scale\n# (every GRPO rollout fires it once per layer otherwise, hiding loss logs).\nimport transformers\ntransformers.logging.set_verbosity_error()\nlogging.getLogger('transformers').setLevel(logging.ERROR)\nlogging.getLogger('transformers.models.qwen2.modeling_qwen2').setLevel(logging.ERROR)\n\nREPO_DIR = globals().get('REPO_DIR') or (\n    '/content/chargebackops' if os.path.isdir('/content/chargebackops')\n    else '/kaggle/working/chargebackops'\n)\nPERSIST_ROOT = globals().get('PERSIST_ROOT') or (\n    os.path.join('/content/drive/MyDrive', 'chargebackops-artifacts')\n    if os.path.isdir('/content/drive/MyDrive') else os.path.dirname(REPO_DIR)\n)\nos.makedirs(PERSIST_ROOT, exist_ok=True)\nif not os.path.isdir(REPO_DIR):\n    raise FileNotFoundError(f'Repository checkout missing at {REPO_DIR}. Run setup first.')\nsys.path.insert(0, REPO_DIR)\nimportlib.invalidate_caches()\nfor mod in list(sys.modules):\n    if mod.startswith(('scenarios', 'training', 'evaluation', 'server', 'core', 'runners', 'connectors')):\n        del sys.modules[mod]\nprint('torch', torch.__version__, '| cuda', torch.cuda.is_available(),\n      torch.cuda.get_device_name(0) if torch.cuda.is_available() else '')\nprint('repo:', REPO_DIR)\nprint('artifact root:', PERSIST_ROOT)"
+   "source": [
+    "# Path + module-cache flush so the editable install resolves before any other import.\n",
+    "import os, sys, importlib, logging, torch\n",
+    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
+    "\n",
+    "# Silence transformers per-layer \"Caching is incompatible...\" spam at scale\n",
+    "# (every GRPO rollout fires it once per layer otherwise, hiding loss logs).\n",
+    "import transformers\n",
+    "transformers.logging.set_verbosity_error()\n",
+    "logging.getLogger('transformers').setLevel(logging.ERROR)\n",
+    "logging.getLogger('transformers.models.qwen2.modeling_qwen2').setLevel(logging.ERROR)\n",
+    "\n",
+    "REPO_DIR = globals().get('REPO_DIR') or (\n",
+    "    '/content/chargebackops' if os.path.isdir('/content/chargebackops')\n",
+    "    else '/kaggle/working/chargebackops'\n",
+    ")\n",
+    "PERSIST_ROOT = globals().get('PERSIST_ROOT') or (\n",
+    "    os.path.join('/content/drive/MyDrive', 'chargebackops-artifacts')\n",
+    "    if os.path.isdir('/content/drive/MyDrive') else os.path.dirname(REPO_DIR)\n",
+    ")\n",
+    "os.makedirs(PERSIST_ROOT, exist_ok=True)\n",
+    "if not os.path.isdir(REPO_DIR):\n",
+    "    raise FileNotFoundError(f'Repository checkout missing at {REPO_DIR}. Run setup first.')\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "importlib.invalidate_caches()\n",
+    "for mod in list(sys.modules):\n",
+    "    if mod.startswith(('scenarios', 'training', 'evaluation', 'server', 'core', 'runners', 'connectors')):\n",
+    "        del sys.modules[mod]\n",
+    "print('torch', torch.__version__, '| cuda', torch.cuda.is_available(),\n",
+    "      torch.cuda.get_device_name(0) if torch.cuda.is_available() else '')\n",
+    "print('repo:', REPO_DIR)\n",
+    "print('artifact root:', PERSIST_ROOT)"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "model-md",
    "metadata": {},
-   "source": "## 1. Load Qwen2.5-3B-Instruct fp16 + attach LoRA\n\n* `dtype=torch.float16` - T4 = Turing sm_75, no bf16 hardware.\n* LoRA r=16 on q/k/v/o + gate/up/down. ~9M trainable params.\n* `gradient_checkpointing` + `enable_input_require_grads` keeps activations off VRAM during backward."
+   "source": [
+    "## 1. Load Qwen2.5-3B-Instruct fp16 + attach LoRA\n",
+    "\n",
+    "* `dtype=torch.float16` - T4 = Turing sm_75, no bf16 hardware.\n",
+    "* LoRA r=16 on q/k/v/o + gate/up/down. ~9M trainable params.\n",
+    "* `gradient_checkpointing` + `enable_input_require_grads` keeps activations off VRAM during backward."
+   ]
   },
   {
    "cell_type": "code",
@@ -40,13 +94,55 @@
    "id": "model-code",
    "metadata": {},
    "outputs": [],
-   "source": "from transformers import AutoModelForCausalLM, AutoTokenizer\nfrom peft import LoraConfig, get_peft_model\n\nMODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen2.5-3B-Instruct')\nprint('MODEL_ID:', MODEL_ID)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\ntokenizer.padding_side = 'left'  # required for GRPO generation\n\nbase_model = AutoModelForCausalLM.from_pretrained(\n    MODEL_ID,\n    torch_dtype=torch.float16,\n    device_map='auto',\n    trust_remote_code=True,\n)\nbase_model.gradient_checkpointing_enable()\nbase_model.enable_input_require_grads()\n\nlora_target_modules = ['q_proj', 'k_proj', 'v_proj', 'o_proj',\n                       'gate_proj', 'up_proj', 'down_proj']\nlora_rank = 16\nlora_alpha = 32\n\nlora_config = LoraConfig(\n    r=lora_rank,\n    lora_alpha=lora_alpha,\n    target_modules=lora_target_modules,\n    lora_dropout=0.05,\n    bias='none',\n    task_type='CAUSAL_LM',\n)\nmodel = get_peft_model(base_model, lora_config)\nmodel.print_trainable_parameters()\nprint(f'VRAM allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB | '\n      f'free: {torch.cuda.mem_get_info()[0]/1e9:.2f} GB')"
+   "source": [
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from peft import LoraConfig, get_peft_model\n",
+    "\n",
+    "MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen2.5-3B-Instruct')\n",
+    "print('MODEL_ID:', MODEL_ID)\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)\n",
+    "if tokenizer.pad_token is None:\n",
+    "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "tokenizer.padding_side = 'left'  # required for GRPO generation\n",
+    "\n",
+    "base_model = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_ID,\n",
+    "    torch_dtype=torch.float16,\n",
+    "    device_map='auto',\n",
+    "    trust_remote_code=True,\n",
+    ")\n",
+    "base_model.gradient_checkpointing_enable()\n",
+    "base_model.enable_input_require_grads()\n",
+    "\n",
+    "lora_target_modules = ['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
+    "                       'gate_proj', 'up_proj', 'down_proj']\n",
+    "lora_rank = 16\n",
+    "lora_alpha = 32\n",
+    "\n",
+    "lora_config = LoraConfig(\n",
+    "    r=lora_rank,\n",
+    "    lora_alpha=lora_alpha,\n",
+    "    target_modules=lora_target_modules,\n",
+    "    lora_dropout=0.05,\n",
+    "    bias='none',\n",
+    "    task_type='CAUSAL_LM',\n",
+    ")\n",
+    "model = get_peft_model(base_model, lora_config)\n",
+    "model.print_trainable_parameters()\n",
+    "print(f'VRAM allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB | '\n",
+    "      f'free: {torch.cuda.mem_get_info()[0]/1e9:.2f} GB')"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "sft-md",
    "metadata": {},
-   "source": "## 2. Phase A - SFT on heuristic rollouts\n\nBuilds (prompt, oracle_completion) pairs by rolling the scripted heuristic on every headline + generated task. Wraps in the Qwen chat template so the model learns the same prompt format used at inference. After this phase the model emits valid JSON with the right `action_type` per state - solves the *always emit `select_case`* collapse before GRPO ever runs."
+   "source": [
+    "## 2. Phase A - SFT on heuristic rollouts\n",
+    "\n",
+    "Builds (prompt, oracle_completion) pairs by rolling the scripted heuristic on every headline + generated task. Wraps in the Qwen chat template so the model learns the same prompt format used at inference. After this phase the model emits valid JSON with the right `action_type` per state - solves the *always emit `select_case`* collapse before GRPO ever runs."
+   ]
   },
   {
    "cell_type": "code",
@@ -54,7 +150,79 @@
    "id": "sft-data-code",
    "metadata": {},
    "outputs": [],
-   "source": "from datasets import Dataset\nfrom scenarios.simulation import list_tasks, get_task\nfrom training.sft_dataset import build_sft_dataset\nfrom collections import Counter\n\n# Synthetic pool. Default 10k rows for T4 speed; override via env var.\nSFT_TARGET_ROWS = int(os.environ.get('SFT_TARGET_ROWS', '10000'))\nSFT_MAX_ROWS = int(os.environ.get('SFT_MAX_ROWS', str(SFT_TARGET_ROWS)))\nSFT_SEED_START = int(os.environ.get('SFT_SEED_START', '1000'))\nSFT_SEED_BATCH = int(os.environ.get('SFT_SEED_BATCH', '128'))\nSFT_MAX_STATES_PER_TASK = int(os.environ.get('SFT_MAX_STATES_PER_TASK', '24'))\nGRPO_SEED_COUNT = int(os.environ.get('GRPO_SEED_COUNT', '160'))\n\n# Holdout seeds excluded from training so eval is defensible.\nHOLDOUT_SEEDS_BY_DIFF = {\n    'easy': {42},\n    'medium': {17, 99},\n    'hard': {7, 53},\n    'nightmare': {31, 77},\n}\nDIFFICULTIES = ['easy', 'medium', 'hard', 'nightmare']\n\nheadline_task_ids = [t.task_id for t in list_tasks()]\ntask_ids = list(headline_task_ids)\nraw_sft = build_sft_dataset(headline_task_ids, max_states_per_task=SFT_MAX_STATES_PER_TASK)\ngenerated_train_task_ids = []\n\nseed_cursor = SFT_SEED_START\nwhile len(raw_sft) < SFT_TARGET_ROWS:\n    batch_task_ids = []\n    for diff in DIFFICULTIES:\n        blocked = HOLDOUT_SEEDS_BY_DIFF.get(diff, set())\n        for seed in range(seed_cursor, seed_cursor + SFT_SEED_BATCH):\n            if seed in blocked:\n                continue\n            tid = f'generated_{diff}_s{seed}'\n            get_task(tid)\n            batch_task_ids.append(tid)\n    raw_sft.extend(build_sft_dataset(batch_task_ids, max_states_per_task=SFT_MAX_STATES_PER_TASK))\n    generated_train_task_ids.extend(batch_task_ids)\n    task_ids.extend(batch_task_ids)\n    seed_cursor += SFT_SEED_BATCH\n    print(f'generated SFT rows: {len(raw_sft):,} / target {SFT_TARGET_ROWS:,}')\n\nif len(raw_sft) > SFT_MAX_ROWS:\n    raw_sft = raw_sft[:SFT_MAX_ROWS]\n\n# Seed list for GRPO state-action curriculum (smaller than SFT pool because\n# GRPO rollouts are slower than SFT forward passes).\nseeds = list(range(SFT_SEED_START, SFT_SEED_START + GRPO_SEED_COUNT))\n\ndef to_chat_text(prompt, completion):\n    return tokenizer.apply_chat_template(\n        [\n            {'role': 'user', 'content': prompt},\n            {'role': 'assistant', 'content': completion},\n        ],\n        tokenize=False,\n        add_generation_prompt=False,\n    )\n\nsft_rows = [{'text': to_chat_text(s['prompt'], s['completion'])} for s in raw_sft]\nsft_dataset = Dataset.from_list(sft_rows)\n\natype_counts = Counter(s['action_type'] for s in raw_sft)\nprint(f'SFT samples: {len(sft_dataset):,}, unique tasks: {len(set(s[\"task_id\"] for s in raw_sft)):,}')\nprint(f'headline tasks: {len(headline_task_ids)}, generated train tasks used: {len(generated_train_task_ids):,}')\nprint(f'excluded generated holdout seeds: {HOLDOUT_SEEDS_BY_DIFF}')\nprint(f'action_type distribution: {dict(atype_counts)}')\nprint('sample (first 500 chars):')\nprint(sft_rows[0]['text'][:500])"
+   "source": [
+    "from datasets import Dataset\n",
+    "from scenarios.simulation import list_tasks, get_task\n",
+    "from training.sft_dataset import build_sft_dataset\n",
+    "from collections import Counter\n",
+    "\n",
+    "# Synthetic pool. Default 10k rows for T4 speed; override via env var.\n",
+    "SFT_TARGET_ROWS = int(os.environ.get('SFT_TARGET_ROWS', '10000'))\n",
+    "SFT_MAX_ROWS = int(os.environ.get('SFT_MAX_ROWS', str(SFT_TARGET_ROWS)))\n",
+    "SFT_SEED_START = int(os.environ.get('SFT_SEED_START', '1000'))\n",
+    "SFT_SEED_BATCH = int(os.environ.get('SFT_SEED_BATCH', '128'))\n",
+    "SFT_MAX_STATES_PER_TASK = int(os.environ.get('SFT_MAX_STATES_PER_TASK', '24'))\n",
+    "GRPO_SEED_COUNT = int(os.environ.get('GRPO_SEED_COUNT', '160'))\n",
+    "\n",
+    "# Holdout seeds excluded from training so eval is defensible.\n",
+    "HOLDOUT_SEEDS_BY_DIFF = {\n",
+    "    'easy': {42},\n",
+    "    'medium': {17, 99},\n",
+    "    'hard': {7, 53},\n",
+    "    'nightmare': {31, 77},\n",
+    "}\n",
+    "DIFFICULTIES = ['easy', 'medium', 'hard', 'nightmare']\n",
+    "\n",
+    "headline_task_ids = [t.task_id for t in list_tasks()]\n",
+    "task_ids = list(headline_task_ids)\n",
+    "raw_sft = build_sft_dataset(headline_task_ids, max_states_per_task=SFT_MAX_STATES_PER_TASK)\n",
+    "generated_train_task_ids = []\n",
+    "\n",
+    "seed_cursor = SFT_SEED_START\n",
+    "while len(raw_sft) < SFT_TARGET_ROWS:\n",
+    "    batch_task_ids = []\n",
+    "    for diff in DIFFICULTIES:\n",
+    "        blocked = HOLDOUT_SEEDS_BY_DIFF.get(diff, set())\n",
+    "        for seed in range(seed_cursor, seed_cursor + SFT_SEED_BATCH):\n",
+    "            if seed in blocked:\n",
+    "                continue\n",
+    "            tid = f'generated_{diff}_s{seed}'\n",
+    "            get_task(tid)\n",
+    "            batch_task_ids.append(tid)\n",
+    "    raw_sft.extend(build_sft_dataset(batch_task_ids, max_states_per_task=SFT_MAX_STATES_PER_TASK))\n",
+    "    generated_train_task_ids.extend(batch_task_ids)\n",
+    "    task_ids.extend(batch_task_ids)\n",
+    "    seed_cursor += SFT_SEED_BATCH\n",
+    "    print(f'generated SFT rows: {len(raw_sft):,} / target {SFT_TARGET_ROWS:,}')\n",
+    "\n",
+    "if len(raw_sft) > SFT_MAX_ROWS:\n",
+    "    raw_sft = raw_sft[:SFT_MAX_ROWS]\n",
+    "\n",
+    "# Seed list for GRPO state-action curriculum (smaller than SFT pool because\n",
+    "# GRPO rollouts are slower than SFT forward passes).\n",
+    "seeds = list(range(SFT_SEED_START, SFT_SEED_START + GRPO_SEED_COUNT))\n",
+    "\n",
+    "def to_chat_text(prompt, completion):\n",
+    "    return tokenizer.apply_chat_template(\n",
+    "        [\n",
+    "            {'role': 'user', 'content': prompt},\n",
+    "            {'role': 'assistant', 'content': completion},\n",
+    "        ],\n",
+    "        tokenize=False,\n",
+    "        add_generation_prompt=False,\n",
+    "    )\n",
+    "\n",
+    "sft_rows = [{'text': to_chat_text(s['prompt'], s['completion'])} for s in raw_sft]\n",
+    "sft_dataset = Dataset.from_list(sft_rows)\n",
+    "\n",
+    "atype_counts = Counter(s['action_type'] for s in raw_sft)\n",
+    "print(f'SFT samples: {len(sft_dataset):,}, unique tasks: {len(set(s[\"task_id\"] for s in raw_sft)):,}')\n",
+    "print(f'headline tasks: {len(headline_task_ids)}, generated train tasks used: {len(generated_train_task_ids):,}')\n",
+    "print(f'excluded generated holdout seeds: {HOLDOUT_SEEDS_BY_DIFF}')\n",
+    "print(f'action_type distribution: {dict(atype_counts)}')\n",
+    "print('sample (first 500 chars):')\n",
+    "print(sft_rows[0]['text'][:500])"
+   ]
   },
   {
    "cell_type": "code",
@@ -62,13 +230,70 @@
    "id": "sft-train-code",
    "metadata": {},
    "outputs": [],
-   "source": "from trl import SFTConfig, SFTTrainer\n\nOUT_ROOT = PERSIST_ROOT\nSFT_DIR = os.path.join(OUT_ROOT, 'sft-merchant-agent')\nGRPO_DIR = os.path.join(OUT_ROOT, 'grpo-merchant-agent')\n\nSFT_FINAL_DIR = os.path.join(SFT_DIR, 'final')\nRUN_SFT_TRAIN = os.environ.get('RUN_SFT_TRAIN', 'auto').strip().lower()\nTRAIN_SFT = RUN_SFT_TRAIN in {'1', 'true', 'yes', 'y', 'on'} or (\n    RUN_SFT_TRAIN == 'auto' and not os.path.isdir(SFT_FINAL_DIR)\n)\nSFT_EPOCHS = float(os.environ.get('SFT_EPOCHS', '1'))\nSFT_LR = float(os.environ.get('SFT_LR', '1e-4'))\nSFT_MAX_STEPS = int(os.environ.get('SFT_MAX_STEPS', '800'))\n\nif not TRAIN_SFT:\n    print(f'Skipping SFT train; using existing adapter at {SFT_FINAL_DIR}')\nelse:\n    sft_config = SFTConfig(\n        output_dir=SFT_DIR,\n        per_device_train_batch_size=1,\n        gradient_accumulation_steps=8,\n        num_train_epochs=SFT_EPOCHS,\n        max_steps=SFT_MAX_STEPS,\n        learning_rate=SFT_LR,\n        logging_steps=10,\n        save_steps=500,\n        save_total_limit=2,\n        bf16=False,\n        fp16=True,\n        gradient_checkpointing=True,\n        gradient_checkpointing_kwargs={'use_reentrant': False},\n        max_length=1024,\n        dataset_text_field='text',\n        report_to='none',\n        optim='adamw_torch',\n        warmup_ratio=0.03,\n    )\n    print(f'SFT config: rows={len(sft_dataset):,}, epochs={SFT_EPOCHS}, lr={SFT_LR}, max_steps={SFT_MAX_STEPS}')\n    if hasattr(model, 'config'):\n        model.config.use_cache = False\n    sft_trainer = SFTTrainer(\n        model=model,\n        args=sft_config,\n        train_dataset=sft_dataset,\n        processing_class=tokenizer,\n    )\n    sft_trainer.train()\n    sft_trainer.save_model(SFT_FINAL_DIR)\n    del sft_trainer\n    torch.cuda.empty_cache()\n    print(f'PEAK VRAM (SFT): {torch.cuda.max_memory_allocated()/1e9:.2f} GB')"
+   "source": [
+    "from trl import SFTConfig, SFTTrainer\n",
+    "\n",
+    "OUT_ROOT = PERSIST_ROOT\n",
+    "SFT_DIR = os.path.join(OUT_ROOT, 'sft-merchant-agent')\n",
+    "GRPO_DIR = os.path.join(OUT_ROOT, 'grpo-merchant-agent')\n",
+    "\n",
+    "SFT_FINAL_DIR = os.path.join(SFT_DIR, 'final')\n",
+    "RUN_SFT_TRAIN = os.environ.get('RUN_SFT_TRAIN', 'auto').strip().lower()\n",
+    "TRAIN_SFT = RUN_SFT_TRAIN in {'1', 'true', 'yes', 'y', 'on'} or (\n",
+    "    RUN_SFT_TRAIN == 'auto' and not os.path.isdir(SFT_FINAL_DIR)\n",
+    ")\n",
+    "SFT_EPOCHS = float(os.environ.get('SFT_EPOCHS', '1'))\n",
+    "SFT_LR = float(os.environ.get('SFT_LR', '1e-4'))\n",
+    "SFT_MAX_STEPS = int(os.environ.get('SFT_MAX_STEPS', '800'))\n",
+    "\n",
+    "if not TRAIN_SFT:\n",
+    "    print(f'Skipping SFT train; using existing adapter at {SFT_FINAL_DIR}')\n",
+    "else:\n",
+    "    sft_config = SFTConfig(\n",
+    "        output_dir=SFT_DIR,\n",
+    "        per_device_train_batch_size=1,\n",
+    "        gradient_accumulation_steps=8,\n",
+    "        num_train_epochs=SFT_EPOCHS,\n",
+    "        max_steps=SFT_MAX_STEPS,\n",
+    "        learning_rate=SFT_LR,\n",
+    "        logging_steps=10,\n",
+    "        save_steps=500,\n",
+    "        save_total_limit=2,\n",
+    "        bf16=False,\n",
+    "        fp16=True,\n",
+    "        gradient_checkpointing=True,\n",
+    "        gradient_checkpointing_kwargs={'use_reentrant': False},\n",
+    "        max_length=1024,\n",
+    "        dataset_text_field='text',\n",
+    "        report_to='none',\n",
+    "        optim='adamw_torch',\n",
+    "        warmup_ratio=0.03,\n",
+    "    )\n",
+    "    print(f'SFT config: rows={len(sft_dataset):,}, epochs={SFT_EPOCHS}, lr={SFT_LR}, max_steps={SFT_MAX_STEPS}')\n",
+    "    if hasattr(model, 'config'):\n",
+    "        model.config.use_cache = False\n",
+    "    sft_trainer = SFTTrainer(\n",
+    "        model=model,\n",
+    "        args=sft_config,\n",
+    "        train_dataset=sft_dataset,\n",
+    "        processing_class=tokenizer,\n",
+    "    )\n",
+    "    sft_trainer.train()\n",
+    "    sft_trainer.save_model(SFT_FINAL_DIR)\n",
+    "    del sft_trainer\n",
+    "    torch.cuda.empty_cache()\n",
+    "    print(f'PEAK VRAM (SFT): {torch.cuda.max_memory_allocated()/1e9:.2f} GB')"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "merge-md",
    "metadata": {},
-   "source": "## 2.5. Merge SFT LoRA into base, attach fresh LoRA for GRPO\n\n`accelerate.unwrap_model_for_generation` calls `merge_adapter()` + `unmerge_adapter()` around generation. With fp16 LoRA the round-trip can lose enough precision that completions degrade. Fix: bake SFT into base via `merge_and_unload()`, then attach a fresh zero-initialized LoRA. The fresh adapter starts as identity, so generation emits SFT-quality output regardless of TRL adapter toggling. GRPO then trains the fresh adapter on top."
+   "source": [
+    "## 2.5. Merge SFT LoRA into base, attach fresh LoRA for GRPO\n",
+    "\n",
+    "`accelerate.unwrap_model_for_generation` calls `merge_adapter()` + `unmerge_adapter()` around generation. With fp16 LoRA the round-trip can lose enough precision that completions degrade. Fix: bake SFT into base via `merge_and_unload()`, then attach a fresh zero-initialized LoRA. The fresh adapter starts as identity, so generation emits SFT-quality output regardless of TRL adapter toggling. GRPO then trains the fresh adapter on top."
+   ]
   },
   {
    "cell_type": "code",
@@ -76,13 +301,84 @@
    "id": "merge-code",
    "metadata": {},
    "outputs": [],
-   "source": "# Reload saved SFT LoRA into a fresh base, merge it, then attach a fresh Phase B LoRA.\nfrom peft import PeftModel\nimport gc\n\nif not os.path.isdir(SFT_FINAL_DIR):\n    raise FileNotFoundError(f'Missing SFT adapter: {SFT_FINAL_DIR}. Run Phase A or upload the adapter first.')\n\nfor name in ['model', 'base_model', 'merged_base']:\n    if name in globals():\n        del globals()[name]\ngc.collect()\ntorch.cuda.empty_cache()\nprint(f'before SFT reload: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')\n\nfresh_base = AutoModelForCausalLM.from_pretrained(\n    MODEL_ID,\n    torch_dtype=torch.float16,\n    device_map='auto',\n    trust_remote_code=True,\n)\nsft_model = PeftModel.from_pretrained(fresh_base, SFT_FINAL_DIR)\nmerged_base = sft_model.merge_and_unload()\ndel sft_model, fresh_base\ngc.collect()\ntorch.cuda.empty_cache()\nprint(f'after SFT merge: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')\nmerged_base.enable_input_require_grads()\n\n# Sanity: SFT-baked base should emit clean JSON deterministically.\nfrom training.env_adapter import build_prompt\nfrom server.chargeback_ops_environment import ChargebackOpsEnvironment\nenv = ChargebackOpsEnvironment()\nobs = env.reset(task_id='goods_not_received_easy')\nchat = tokenizer.apply_chat_template(\n    [{'role': 'user', 'content': build_prompt(obs.model_dump())}],\n    tokenize=False, add_generation_prompt=True,\n)\ninp = tokenizer(chat, return_tensors='pt').to(merged_base.device)\nmerged_base.eval()\nwith torch.no_grad():\n    out = merged_base.generate(\n        **inp,\n        max_new_tokens=160,\n        do_sample=False,\n        pad_token_id=tokenizer.eos_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    )\nprint('merged-base gen:', repr(tokenizer.decode(out[0][inp.input_ids.shape[1]:], skip_special_tokens=False)))\nmerged_base.train()\n\n# Attach fresh Phase B LoRA. lora_dropout=0 so generation can't be polluted by\n# random adapter contributions while train() mode is active.\nlora_phase_b = LoraConfig(\n    r=lora_rank,\n    lora_alpha=lora_alpha,\n    target_modules=lora_target_modules,\n    lora_dropout=0.0,\n    bias='none',\n    task_type='CAUSAL_LM',\n)\nmodel = get_peft_model(merged_base, lora_phase_b)\nmodel.enable_input_require_grads()\nmodel.print_trainable_parameters()\nprint(f'after fresh Phase B LoRA: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')"
+   "source": [
+    "# Reload saved SFT LoRA into a fresh base, merge it, then attach a fresh Phase B LoRA.\n",
+    "from peft import PeftModel\n",
+    "import gc\n",
+    "\n",
+    "if not os.path.isdir(SFT_FINAL_DIR):\n",
+    "    raise FileNotFoundError(f'Missing SFT adapter: {SFT_FINAL_DIR}. Run Phase A or upload the adapter first.')\n",
+    "\n",
+    "for name in ['model', 'base_model', 'merged_base']:\n",
+    "    if name in globals():\n",
+    "        del globals()[name]\n",
+    "gc.collect()\n",
+    "torch.cuda.empty_cache()\n",
+    "print(f'before SFT reload: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')\n",
+    "\n",
+    "fresh_base = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_ID,\n",
+    "    torch_dtype=torch.float16,\n",
+    "    device_map='auto',\n",
+    "    trust_remote_code=True,\n",
+    ")\n",
+    "sft_model = PeftModel.from_pretrained(fresh_base, SFT_FINAL_DIR)\n",
+    "merged_base = sft_model.merge_and_unload()\n",
+    "del sft_model, fresh_base\n",
+    "gc.collect()\n",
+    "torch.cuda.empty_cache()\n",
+    "print(f'after SFT merge: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')\n",
+    "merged_base.enable_input_require_grads()\n",
+    "\n",
+    "# Sanity: SFT-baked base should emit clean JSON deterministically.\n",
+    "from training.env_adapter import build_prompt\n",
+    "from server.chargeback_ops_environment import ChargebackOpsEnvironment\n",
+    "env = ChargebackOpsEnvironment()\n",
+    "obs = env.reset(task_id='goods_not_received_easy')\n",
+    "chat = tokenizer.apply_chat_template(\n",
+    "    [{'role': 'user', 'content': build_prompt(obs.model_dump())}],\n",
+    "    tokenize=False, add_generation_prompt=True,\n",
+    ")\n",
+    "inp = tokenizer(chat, return_tensors='pt').to(merged_base.device)\n",
+    "merged_base.eval()\n",
+    "with torch.no_grad():\n",
+    "    out = merged_base.generate(\n",
+    "        **inp,\n",
+    "        max_new_tokens=160,\n",
+    "        do_sample=False,\n",
+    "        pad_token_id=tokenizer.eos_token_id,\n",
+    "        eos_token_id=tokenizer.eos_token_id,\n",
+    "    )\n",
+    "print('merged-base gen:', repr(tokenizer.decode(out[0][inp.input_ids.shape[1]:], skip_special_tokens=False)))\n",
+    "merged_base.train()\n",
+    "\n",
+    "# Attach fresh Phase B LoRA. lora_dropout=0 so generation can't be polluted by\n",
+    "# random adapter contributions while train() mode is active.\n",
+    "lora_phase_b = LoraConfig(\n",
+    "    r=lora_rank,\n",
+    "    lora_alpha=lora_alpha,\n",
+    "    target_modules=lora_target_modules,\n",
+    "    lora_dropout=0.0,\n",
+    "    bias='none',\n",
+    "    task_type='CAUSAL_LM',\n",
+    ")\n",
+    "model = get_peft_model(merged_base, lora_phase_b)\n",
+    "model.enable_input_require_grads()\n",
+    "model.print_trainable_parameters()\n",
+    "print(f'after fresh Phase B LoRA: VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB')"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "grpo-md",
    "metadata": {},
-   "source": "## 3. Phase B - GRPO with outcome reward (RLVR, not distillation)\n\nReward source: terminal `$` PnL after the model's action plus heuristic tail-rollout. The merchant earns positive reward only when its action leads to a winning packet against the scripted Issuer or arbitration. A second format-shaping reward provides dense early-training signal so GRPO has gradient before the policy can produce winning packets.\n\nThis is the core of the v2 pivot: the trainer no longer optimizes \"match the heuristic.\" It optimizes \"win the dispute.\" The Issuer + arbitration are the verifier - they cannot be tricked because reward is dollar-denominated outcome."
+   "source": [
+    "## 3. Phase B - GRPO with outcome reward (RLVR, not distillation)\n",
+    "\n",
+    "Reward source: terminal `$` PnL after the model's action plus heuristic tail-rollout. The merchant earns positive reward only when its action leads to a winning packet against the scripted Issuer or arbitration. A second format-shaping reward provides dense early-training signal so GRPO has gradient before the policy can produce winning packets.\n",
+    "\n",
+    "It optimizes \"win the dispute.\" The Issuer + arbitration are the verifier - they cannot be tricked because reward is dollar-denominated outcome."
+   ]
   },
   {
    "cell_type": "code",
@@ -90,7 +386,62 @@
    "id": "grpo-data-code",
    "metadata": {},
    "outputs": [],
-   "source": "from training.reward_adapter import build_state_action_dataset\n\n# State-action samples: (task_id, state_step, prompt) tuples captured by\n# rolling the heuristic forward on each task. The model will be asked to\n# pick the next action at each captured state.\nPHASE_B_MAX_STATES_PER_TASK = int(os.environ.get('PHASE_B_MAX_STATES_PER_TASK', '10'))\nGRPO_DIFFICULTIES = tuple(\n    d.strip()\n    for d in os.environ.get('GRPO_DIFFICULTIES', 'easy,medium,hard,nightmare').split(',')\n    if d.strip()\n)\ncurriculum_task_ids = [\n    t.task_id for t in list_tasks()\n    if t.difficulty in GRPO_DIFFICULTIES\n]\nfor diff in GRPO_DIFFICULTIES:\n    blocked = HOLDOUT_SEEDS_BY_DIFF.get(diff, set())\n    for s in seeds:\n        if s in blocked:\n            continue\n        tid = f'generated_{diff}_s{s}'\n        try:\n            get_task(tid)\n            if tid not in curriculum_task_ids:\n                curriculum_task_ids.append(tid)\n        except Exception:\n            pass\n\nraw_grpo = build_state_action_dataset(\n    curriculum_task_ids, max_states_per_task=PHASE_B_MAX_STATES_PER_TASK,\n)\n\ndef to_chat_prompt(prompt):\n    return tokenizer.apply_chat_template(\n        [{'role': 'user', 'content': prompt}],\n        tokenize=False, add_generation_prompt=True,\n    )\n\ngrpo_rows = []\nfor sample in raw_grpo:\n    chat_prompt = to_chat_prompt(sample['prompt'])\n    n_tokens = len(tokenizer(chat_prompt, add_special_tokens=False)['input_ids'])\n    if n_tokens <= 1000:\n        grpo_rows.append({\n            'prompt': chat_prompt,\n            'task_id': sample['task_id'],\n            'state_step': int(sample['state_step']),\n        })\n\ngrpo_dataset = Dataset.from_list(grpo_rows)\nunique_tasks = len({row['task_id'] for row in grpo_rows})\nprint(f'GRPO state-action samples: {len(grpo_dataset)}, '\n      f'unique tasks: {unique_tasks}, difficulties={GRPO_DIFFICULTIES}, '\n      f'max_states_per_task={PHASE_B_MAX_STATES_PER_TASK}')"
+   "source": [
+    "from training.reward_adapter import build_state_action_dataset\n",
+    "\n",
+    "# State-action samples: (task_id, state_step, prompt) tuples captured by\n",
+    "# rolling the heuristic forward on each task. The model will be asked to\n",
+    "# pick the next action at each captured state.\n",
+    "PHASE_B_MAX_STATES_PER_TASK = int(os.environ.get('PHASE_B_MAX_STATES_PER_TASK', '10'))\n",
+    "GRPO_DIFFICULTIES = tuple(\n",
+    "    d.strip()\n",
+    "    for d in os.environ.get('GRPO_DIFFICULTIES', 'easy,medium,hard,nightmare').split(',')\n",
+    "    if d.strip()\n",
+    ")\n",
+    "curriculum_task_ids = [\n",
+    "    t.task_id for t in list_tasks()\n",
+    "    if t.difficulty in GRPO_DIFFICULTIES\n",
+    "]\n",
+    "for diff in GRPO_DIFFICULTIES:\n",
+    "    blocked = HOLDOUT_SEEDS_BY_DIFF.get(diff, set())\n",
+    "    for s in seeds:\n",
+    "        if s in blocked:\n",
+    "            continue\n",
+    "        tid = f'generated_{diff}_s{s}'\n",
+    "        try:\n",
+    "            get_task(tid)\n",
+    "            if tid not in curriculum_task_ids:\n",
+    "                curriculum_task_ids.append(tid)\n",
+    "        except Exception:\n",
+    "            pass\n",
+    "\n",
+    "raw_grpo = build_state_action_dataset(\n",
+    "    curriculum_task_ids, max_states_per_task=PHASE_B_MAX_STATES_PER_TASK,\n",
+    ")\n",
+    "\n",
+    "def to_chat_prompt(prompt):\n",
+    "    return tokenizer.apply_chat_template(\n",
+    "        [{'role': 'user', 'content': prompt}],\n",
+    "        tokenize=False, add_generation_prompt=True,\n",
+    "    )\n",
+    "\n",
+    "grpo_rows = []\n",
+    "for sample in raw_grpo:\n",
+    "    chat_prompt = to_chat_prompt(sample['prompt'])\n",
+    "    n_tokens = len(tokenizer(chat_prompt, add_special_tokens=False)['input_ids'])\n",
+    "    if n_tokens <= 1000:\n",
+    "        grpo_rows.append({\n",
+    "            'prompt': chat_prompt,\n",
+    "            'task_id': sample['task_id'],\n",
+    "            'state_step': int(sample['state_step']),\n",
+    "        })\n",
+    "\n",
+    "grpo_dataset = Dataset.from_list(grpo_rows)\n",
+    "unique_tasks = len({row['task_id'] for row in grpo_rows})\n",
+    "print(f'GRPO state-action samples: {len(grpo_dataset)}, '\n",
+    "      f'unique tasks: {unique_tasks}, difficulties={GRPO_DIFFICULTIES}, '\n",
+    "      f'max_states_per_task={PHASE_B_MAX_STATES_PER_TASK}')"
+   ]
   },
   {
    "cell_type": "code",
@@ -98,13 +449,102 @@
    "id": "grpo-train-code",
    "metadata": {},
    "outputs": [],
-   "source": "from trl import GRPOConfig, GRPOTrainer\nfrom training.outcome_reward import compute_outcome_reward, compute_format_reward\nimport torch.nn as nn\n\n# Re-arm hooks and disable cache/noisy stochastic modules.\nmodel.enable_input_require_grads()\nif hasattr(model, 'config'):\n    model.config.use_cache = False\nmodel.config.eos_token_id = tokenizer.eos_token_id\nmodel.config.pad_token_id = tokenizer.pad_token_id\nif hasattr(model, 'generation_config'):\n    model.generation_config.eos_token_id = tokenizer.eos_token_id\n    model.generation_config.pad_token_id = tokenizer.pad_token_id\n\n# Belt-and-suspenders dropout zeroing in case any new Dropout slipped in.\nfor m in model.modules():\n    if isinstance(m, nn.Dropout):\n        m.p = 0.0\n\ndef outcome_reward_fn(prompts, completions, **kwargs):\n    task_ids = kwargs.get('task_id') or kwargs.get('task_ids')\n    state_steps = kwargs.get('state_step') or kwargs.get('state_steps')\n    return compute_outcome_reward(\n        prompts, completions,\n        task_ids=task_ids, state_steps=state_steps,\n    )\n\ndef format_reward_fn(prompts, completions, **kwargs):\n    return compute_format_reward(prompts, completions)\n\n# 3B fp16 + LoRA + batch 1 x 8 grad-accum + num_generations=4 fits T4. Disabling\n# grad-checkpoint roughly halves per-step time at the cost of more activation\n# VRAM; SFT peaked at <8GB so there is headroom.\ngrpo_config = GRPOConfig(\n    output_dir=GRPO_DIR,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=8,\n    num_generations=4,\n    max_prompt_length=1024,\n    max_completion_length=192,\n    learning_rate=float(os.environ.get('GRPO_LR', '5e-6')),\n    max_steps=int(os.environ.get('GRPO_MAX_STEPS', '120')),\n    logging_steps=5,\n    save_steps=40,\n    save_total_limit=3,\n    bf16=False,\n    fp16=True,\n    max_grad_norm=0.5,\n    gradient_checkpointing=False,\n    report_to='none',\n    beta=0.0,\n    # Tightened sampling defaults: TRL 0.21 GRPOConfig defaults to\n    # top_p=1.0, top_k=0, repetition_penalty=1.0 - i.e. pure temperature\n    # sampling on Qwen's 152K vocab. Even when the SFT policy puts 99% mass\n    # on <|im_end|>, the 1% tail across 150K tokens is fat enough that\n    # temp=0.7 sampling can pick a junk token. These match Qwen2.5's own\n    # generation_config defaults.\n    temperature=0.7,\n    top_p=0.9,\n    top_k=50,\n    repetition_penalty=1.05,\n    use_vllm=False,\n    log_completions=True,\n    num_completions_to_print=2,\n    optim='adamw_torch',\n    lr_scheduler_type='constant',\n)\n\nRUN_GRPO = os.environ.get('RUN_GRPO', '1').strip().lower() not in {'0', 'false', 'no'}\nif RUN_GRPO:\n    grpo_trainer = GRPOTrainer(\n        model=model,\n        processing_class=tokenizer,\n        reward_funcs=[outcome_reward_fn, format_reward_fn],\n        args=grpo_config,\n        train_dataset=grpo_dataset,\n    )\n    grpo_trainer.train()\n    grpo_trainer.save_model(os.path.join(GRPO_DIR, 'final'))\n    del grpo_trainer\n    torch.cuda.empty_cache()\nelse:\n    print('RUN_GRPO=0: skipped GRPO training')\nprint(f'PEAK VRAM (GRPO): {torch.cuda.max_memory_allocated()/1e9:.2f} GB')"
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "from training.outcome_reward import compute_outcome_reward, compute_format_reward\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "# Re-arm hooks and disable cache/noisy stochastic modules.\n",
+    "model.enable_input_require_grads()\n",
+    "if hasattr(model, 'config'):\n",
+    "    model.config.use_cache = False\n",
+    "model.config.eos_token_id = tokenizer.eos_token_id\n",
+    "model.config.pad_token_id = tokenizer.pad_token_id\n",
+    "if hasattr(model, 'generation_config'):\n",
+    "    model.generation_config.eos_token_id = tokenizer.eos_token_id\n",
+    "    model.generation_config.pad_token_id = tokenizer.pad_token_id\n",
+    "\n",
+    "# Belt-and-suspenders dropout zeroing in case any new Dropout slipped in.\n",
+    "for m in model.modules():\n",
+    "    if isinstance(m, nn.Dropout):\n",
+    "        m.p = 0.0\n",
+    "\n",
+    "def outcome_reward_fn(prompts, completions, **kwargs):\n",
+    "    task_ids = kwargs.get('task_id') or kwargs.get('task_ids')\n",
+    "    state_steps = kwargs.get('state_step') or kwargs.get('state_steps')\n",
+    "    return compute_outcome_reward(\n",
+    "        prompts, completions,\n",
+    "        task_ids=task_ids, state_steps=state_steps,\n",
+    "    )\n",
+    "\n",
+    "def format_reward_fn(prompts, completions, **kwargs):\n",
+    "    return compute_format_reward(prompts, completions)\n",
+    "\n",
+    "# 3B fp16 + LoRA + batch 1 x 8 grad-accum + num_generations=4 fits T4. Disabling\n",
+    "# grad-checkpoint roughly halves per-step time at the cost of more activation\n",
+    "# VRAM; SFT peaked at <8GB so there is headroom.\n",
+    "grpo_config = GRPOConfig(\n",
+    "    output_dir=GRPO_DIR,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=8,\n",
+    "    num_generations=4,\n",
+    "    max_prompt_length=1024,\n",
+    "    max_completion_length=192,\n",
+    "    learning_rate=float(os.environ.get('GRPO_LR', '5e-6')),\n",
+    "    max_steps=int(os.environ.get('GRPO_MAX_STEPS', '120')),\n",
+    "    logging_steps=5,\n",
+    "    save_steps=40,\n",
+    "    save_total_limit=3,\n",
+    "    bf16=False,\n",
+    "    fp16=True,\n",
+    "    max_grad_norm=0.5,\n",
+    "    gradient_checkpointing=False,\n",
+    "    report_to='none',\n",
+    "    beta=0.0,\n",
+    "    # Tightened sampling defaults: TRL 0.21 GRPOConfig defaults to\n",
+    "    # top_p=1.0, top_k=0, repetition_penalty=1.0 - i.e. pure temperature\n",
+    "    # sampling on Qwen's 152K vocab. Even when the SFT policy puts 99% mass\n",
+    "    # on <|im_end|>, the 1% tail across 150K tokens is fat enough that\n",
+    "    # temp=0.7 sampling can pick a junk token. These match Qwen2.5's own\n",
+    "    # generation_config defaults.\n",
+    "    temperature=0.7,\n",
+    "    top_p=0.9,\n",
+    "    top_k=50,\n",
+    "    repetition_penalty=1.05,\n",
+    "    use_vllm=False,\n",
+    "    log_completions=True,\n",
+    "    num_completions_to_print=2,\n",
+    "    optim='adamw_torch',\n",
+    "    lr_scheduler_type='constant',\n",
+    ")\n",
+    "\n",
+    "RUN_GRPO = os.environ.get('RUN_GRPO', '1').strip().lower() not in {'0', 'false', 'no'}\n",
+    "if RUN_GRPO:\n",
+    "    grpo_trainer = GRPOTrainer(\n",
+    "        model=model,\n",
+    "        processing_class=tokenizer,\n",
+    "        reward_funcs=[outcome_reward_fn, format_reward_fn],\n",
+    "        args=grpo_config,\n",
+    "        train_dataset=grpo_dataset,\n",
+    "    )\n",
+    "    grpo_trainer.train()\n",
+    "    grpo_trainer.save_model(os.path.join(GRPO_DIR, 'final'))\n",
+    "    del grpo_trainer\n",
+    "    torch.cuda.empty_cache()\n",
+    "else:\n",
+    "    print('RUN_GRPO=0: skipped GRPO training')\n",
+    "print(f'PEAK VRAM (GRPO): {torch.cuda.max_memory_allocated()/1e9:.2f} GB')"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "eval-md",
    "metadata": {},
-   "source": "## 4. Per-checkpoint eval - overall + per-difficulty\n\nLoads each saved adapter, plays full episodes across the headline catalog, and plots the curve. Stall detection in `run_episode_with_text_policy` ensures degenerate checkpoints reach grading instead of returning 0."
+   "source": [
+    "## 4. Per-checkpoint eval - overall + per-difficulty\n",
+    "\n",
+    "Loads each saved adapter, plays full episodes across the headline catalog, and plots the curve. Stall detection in `run_episode_with_text_policy` ensures degenerate checkpoints reach grading instead of returning 0."
+   ]
   },
   {
    "cell_type": "code",
@@ -112,13 +552,107 @@
    "id": "eval-code",
    "metadata": {},
    "outputs": [],
-   "source": "import glob, re\nfrom peft import PeftModel\nfrom training.curve import (\n    evaluate_checkpoint, evaluate_checkpoint_by_family,\n    plot_training_curve, plot_training_curve_by_family,\n)\n\ndef make_text_policy(adapter_path, adapter_kind='base'):\n    base = AutoModelForCausalLM.from_pretrained(\n        MODEL_ID, torch_dtype=torch.float16, device_map='auto', trust_remote_code=True,\n    )\n    if adapter_kind == 'grpo':\n        sft = PeftModel.from_pretrained(base, SFT_FINAL_DIR)\n        base = sft.merge_and_unload()\n        m = PeftModel.from_pretrained(base, adapter_path)\n    elif adapter_path is not None:\n        m = PeftModel.from_pretrained(base, adapter_path)\n    else:\n        m = base\n    m.eval()\n    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)\n    if tok.pad_token is None:\n        tok.pad_token = tok.eos_token\n    tok.padding_side = 'left'\n\n    def policy(prompt):\n        chat = tok.apply_chat_template(\n            [{'role': 'user', 'content': prompt}],\n            tokenize=False, add_generation_prompt=True,\n        )\n        inputs = tok(chat, return_tensors='pt', truncation=True, max_length=1024).to(m.device)\n        with torch.no_grad():\n            out = m.generate(\n                **inputs, max_new_tokens=192, do_sample=False,\n                pad_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id,\n            )\n        return tok.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n    return policy, m\n\n# Catalog of checkpoints: untrained base, SFT final, every saved GRPO checkpoint.\nckpt_specs = [('base', None, 0, 'base'),\n              ('sft', SFT_FINAL_DIR, 1, 'sft')]\ngrpo_dirs = sorted(\n    glob.glob(os.path.join(GRPO_DIR, 'checkpoint-*')),\n    key=lambda p: int(re.search(r'checkpoint-(\\d+)', p).group(1)),\n)\nfor d in grpo_dirs:\n    step = int(re.search(r'checkpoint-(\\d+)', d).group(1))\n    ckpt_specs.append((f'grpo-{step}', d, 1 + step, 'grpo'))\nfinal_grpo = os.path.join(GRPO_DIR, 'final')\nif os.path.isdir(final_grpo):\n    ckpt_specs.append(('grpo-final', final_grpo, 1 + (max(int(re.search(r'checkpoint-(\\d+)', d).group(1)) for d in grpo_dirs) if grpo_dirs else 0) + 1, 'grpo'))\n\noverall = []\ngrouped = []\nfor label, path, step, kind in ckpt_specs:\n    print(f'eval {label} from {path}')\n    pol, m_ckpt = make_text_policy(path, kind)\n    overall.append(evaluate_checkpoint(step=step, policy=pol))\n    grouped.append(evaluate_checkpoint_by_family(step=step, policy=pol))\n    del m_ckpt\n    torch.cuda.empty_cache()\n\nprint('\\nOVERALL CURVE:')\nfor c in overall:\n    print(f'  step={c.step:4d} mean={c.mean_score:.4f}')\n\nprint('\\nPER-FAMILY CURVE:')\nfor g in grouped:\n    line = f'  step={g.step:4d}'\n    for fam in sorted(g.by_family.keys()):\n        line += f'  {fam}={g.by_family[fam].mean_score:.3f}'\n    print(line)\n\nfrom runners.benchmark_runner import run_policy_sweep\nsweep = run_policy_sweep()\nheur_overall = next(s.mean_score for s in sweep.policies if s.policy == 'heuristic')\n\nFIG_DIR = os.path.join(REPO_DIR, 'docs', 'figures')\nos.makedirs(FIG_DIR, exist_ok=True)\nplot_training_curve(\n    overall, os.path.join(FIG_DIR, 'training_curve.png'),\n    baseline_scores={'heuristic': heur_overall, 'naive': 0.0},\n)\nplot_training_curve_by_family(\n    grouped, os.path.join(FIG_DIR, 'training_curve_by_family.png'),\n    family_order=['easy', 'medium', 'hard', 'nightmare'],\n)\nprint(f'\\nfigures saved to {FIG_DIR}/')"
+   "source": [
+    "import glob, re\n",
+    "from peft import PeftModel\n",
+    "from training.curve import (\n",
+    "    evaluate_checkpoint, evaluate_checkpoint_by_family,\n",
+    "    plot_training_curve, plot_training_curve_by_family,\n",
+    ")\n",
+    "\n",
+    "def make_text_policy(adapter_path, adapter_kind='base'):\n",
+    "    base = AutoModelForCausalLM.from_pretrained(\n",
+    "        MODEL_ID, torch_dtype=torch.float16, device_map='auto', trust_remote_code=True,\n",
+    "    )\n",
+    "    if adapter_kind == 'grpo':\n",
+    "        sft = PeftModel.from_pretrained(base, SFT_FINAL_DIR)\n",
+    "        base = sft.merge_and_unload()\n",
+    "        m = PeftModel.from_pretrained(base, adapter_path)\n",
+    "    elif adapter_path is not None:\n",
+    "        m = PeftModel.from_pretrained(base, adapter_path)\n",
+    "    else:\n",
+    "        m = base\n",
+    "    m.eval()\n",
+    "    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)\n",
+    "    if tok.pad_token is None:\n",
+    "        tok.pad_token = tok.eos_token\n",
+    "    tok.padding_side = 'left'\n",
+    "\n",
+    "    def policy(prompt):\n",
+    "        chat = tok.apply_chat_template(\n",
+    "            [{'role': 'user', 'content': prompt}],\n",
+    "            tokenize=False, add_generation_prompt=True,\n",
+    "        )\n",
+    "        inputs = tok(chat, return_tensors='pt', truncation=True, max_length=1024).to(m.device)\n",
+    "        with torch.no_grad():\n",
+    "            out = m.generate(\n",
+    "                **inputs, max_new_tokens=192, do_sample=False,\n",
+    "                pad_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id,\n",
+    "            )\n",
+    "        return tok.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
+    "    return policy, m\n",
+    "\n",
+    "# Catalog of checkpoints: untrained base, SFT final, every saved GRPO checkpoint.\n",
+    "ckpt_specs = [('base', None, 0, 'base'),\n",
+    "              ('sft', SFT_FINAL_DIR, 1, 'sft')]\n",
+    "grpo_dirs = sorted(\n",
+    "    glob.glob(os.path.join(GRPO_DIR, 'checkpoint-*')),\n",
+    "    key=lambda p: int(re.search(r'checkpoint-(\\d+)', p).group(1)),\n",
+    ")\n",
+    "for d in grpo_dirs:\n",
+    "    step = int(re.search(r'checkpoint-(\\d+)', d).group(1))\n",
+    "    ckpt_specs.append((f'grpo-{step}', d, 1 + step, 'grpo'))\n",
+    "final_grpo = os.path.join(GRPO_DIR, 'final')\n",
+    "if os.path.isdir(final_grpo):\n",
+    "    ckpt_specs.append(('grpo-final', final_grpo, 1 + (max(int(re.search(r'checkpoint-(\\d+)', d).group(1)) for d in grpo_dirs) if grpo_dirs else 0) + 1, 'grpo'))\n",
+    "\n",
+    "overall = []\n",
+    "grouped = []\n",
+    "for label, path, step, kind in ckpt_specs:\n",
+    "    print(f'eval {label} from {path}')\n",
+    "    pol, m_ckpt = make_text_policy(path, kind)\n",
+    "    overall.append(evaluate_checkpoint(step=step, policy=pol))\n",
+    "    grouped.append(evaluate_checkpoint_by_family(step=step, policy=pol))\n",
+    "    del m_ckpt\n",
+    "    torch.cuda.empty_cache()\n",
+    "\n",
+    "print('\\nOVERALL CURVE:')\n",
+    "for c in overall:\n",
+    "    print(f'  step={c.step:4d} mean={c.mean_score:.4f}')\n",
+    "\n",
+    "print('\\nPER-FAMILY CURVE:')\n",
+    "for g in grouped:\n",
+    "    line = f'  step={g.step:4d}'\n",
+    "    for fam in sorted(g.by_family.keys()):\n",
+    "        line += f'  {fam}={g.by_family[fam].mean_score:.3f}'\n",
+    "    print(line)\n",
+    "\n",
+    "from runners.benchmark_runner import run_policy_sweep\n",
+    "sweep = run_policy_sweep()\n",
+    "heur_overall = next(s.mean_score for s in sweep.policies if s.policy == 'heuristic')\n",
+    "\n",
+    "FIG_DIR = os.path.join(REPO_DIR, 'docs', 'figures')\n",
+    "os.makedirs(FIG_DIR, exist_ok=True)\n",
+    "plot_training_curve(\n",
+    "    overall, os.path.join(FIG_DIR, 'training_curve.png'),\n",
+    "    baseline_scores={'heuristic': heur_overall, 'naive': 0.0},\n",
+    ")\n",
+    "plot_training_curve_by_family(\n",
+    "    grouped, os.path.join(FIG_DIR, 'training_curve_by_family.png'),\n",
+    "    family_order=['easy', 'medium', 'hard', 'nightmare'],\n",
+    ")\n",
+    "print(f'\\nfigures saved to {FIG_DIR}/')"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "diag-md",
    "metadata": {},
-   "source": "## 5. Diagnose final checkpoint\n\nPrint the trained checkpoint's completion vs. the heuristic oracle on three representative tasks (easy, hard, nightmare). Verifies the model emits valid JSON and shows what the trained policy actually does."
+   "source": [
+    "## 5. Diagnose final checkpoint\n",
+    "\n",
+    "Print the trained checkpoint's completion vs. the heuristic oracle on three representative tasks (easy, hard, nightmare). Verifies the model emits valid JSON and shows what the trained policy actually does."
+   ]
   },
   {
    "cell_type": "code",
@@ -126,13 +660,61 @@
    "id": "diag-code",
    "metadata": {},
    "outputs": [],
-   "source": "from training.env_adapter import build_prompt, parse_completion\nfrom training.outcome_reward import compute_outcome_reward\nfrom server.chargeback_ops_environment import ChargebackOpsEnvironment\nfrom runners.benchmark_runner import heuristic_policy\n\nfinal_adapter = (\n    os.path.join(GRPO_DIR, 'final')\n    if os.path.isdir(os.path.join(GRPO_DIR, 'final'))\n    else (grpo_dirs[-1] if grpo_dirs else SFT_FINAL_DIR)\n)\nfinal_kind = 'grpo' if 'grpo-merchant-agent' in final_adapter else 'sft'\nprint(f'diagnose adapter: {final_adapter}')\npolicy, m = make_text_policy(final_adapter, final_kind)\n\nfor tid in ['goods_not_received_easy', 'queue_optimization_hard', 'generated_nightmare_s31']:\n    env = ChargebackOpsEnvironment()\n    obs = env.reset(task_id=tid)\n    raw = build_prompt(obs.model_dump())\n    completion = policy(raw)\n    parsed = parse_completion(completion)\n    oracle = heuristic_policy(obs.model_dump())\n    pnl = compute_outcome_reward(['x'], [completion], task_ids=[tid], state_steps=[0])[0]\n    print(f'\\n=== {tid} ===')\n    print(f'oracle: {oracle.action_type} case={oracle.case_id}')\n    print(f'completion (first 200): {repr(completion[:200])}')\n    print(f'parsed: {parsed}')\n    print(f'outcome PnL (normalized): {pnl:+.3f}')\ndel m\ntorch.cuda.empty_cache()"
+   "source": [
+    "from training.env_adapter import build_prompt, parse_completion\n",
+    "from training.outcome_reward import compute_outcome_reward\n",
+    "from server.chargeback_ops_environment import ChargebackOpsEnvironment\n",
+    "from runners.benchmark_runner import heuristic_policy\n",
+    "\n",
+    "final_adapter = (\n",
+    "    os.path.join(GRPO_DIR, 'final')\n",
+    "    if os.path.isdir(os.path.join(GRPO_DIR, 'final'))\n",
+    "    else (grpo_dirs[-1] if grpo_dirs else SFT_FINAL_DIR)\n",
+    ")\n",
+    "final_kind = 'grpo' if 'grpo-merchant-agent' in final_adapter else 'sft'\n",
+    "print(f'diagnose adapter: {final_adapter}')\n",
+    "policy, m = make_text_policy(final_adapter, final_kind)\n",
+    "\n",
+    "for tid in ['goods_not_received_easy', 'queue_optimization_hard', 'generated_nightmare_s31']:\n",
+    "    env = ChargebackOpsEnvironment()\n",
+    "    obs = env.reset(task_id=tid)\n",
+    "    raw = build_prompt(obs.model_dump())\n",
+    "    completion = policy(raw)\n",
+    "    parsed = parse_completion(completion)\n",
+    "    oracle = heuristic_policy(obs.model_dump())\n",
+    "    pnl = compute_outcome_reward(['x'], [completion], task_ids=[tid], state_steps=[0])[0]\n",
+    "    print(f'\\n=== {tid} ===')\n",
+    "    print(f'oracle: {oracle.action_type} case={oracle.case_id}')\n",
+    "    print(f'completion (first 200): {repr(completion[:200])}')\n",
+    "    print(f'parsed: {parsed}')\n",
+    "    print(f'outcome PnL (normalized): {pnl:+.3f}')\n",
+    "del m\n",
+    "torch.cuda.empty_cache()"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "done-md",
    "metadata": {},
-   "source": "## Done\n\nArtifacts written to `docs/figures/`:\n* `training_curve.png` - overall mean rubric score across SFT + GRPO checkpoints, with heuristic baseline.\n* `training_curve_by_family.png` - per-difficulty curves (easy / medium / hard / nightmare).\n\nAdapter weights (under `PERSIST_ROOT`):\n* `sft-merchant-agent/final/` - Phase A output.\n* `grpo-merchant-agent/final/` - Phase B output (GRPO with outcome reward).\n\nTo use the trained model:\n```python\nfrom peft import PeftModel\nsft_model = PeftModel.from_pretrained(base, 'sft-merchant-agent/final')\nmerged = sft_model.merge_and_unload()\ntrained = PeftModel.from_pretrained(merged, 'grpo-merchant-agent/final')\n```"
+   "source": [
+    "## Done\n",
+    "\n",
+    "Artifacts written to `docs/figures/`:\n",
+    "* `training_curve.png` - overall mean rubric score across SFT + GRPO checkpoints, with heuristic baseline.\n",
+    "* `training_curve_by_family.png` - per-difficulty curves (easy / medium / hard / nightmare).\n",
+    "\n",
+    "Adapter weights (under `PERSIST_ROOT`):\n",
+    "* `sft-merchant-agent/final/` - Phase A output.\n",
+    "* `grpo-merchant-agent/final/` - Phase B output (GRPO with outcome reward).\n",
+    "\n",
+    "To use the trained model:\n",
+    "```python\n",
+    "from peft import PeftModel\n",
+    "sft_model = PeftModel.from_pretrained(base, 'sft-merchant-agent/final')\n",
+    "merged = sft_model.merge_and_unload()\n",
+    "trained = PeftModel.from_pretrained(merged, 'grpo-merchant-agent/final')\n",
+    "```"
+   ]
   }
  ],
  "metadata": {