v3_training: escape walk_away collapse with constructive fallback + few-shot

The training was stuck at reward=0.172 with 4/4 rollouts producing
turns=1 deal=N -- the 1B model was hitting the JSON parser fallback,
which was set to walk_away. Every rollout taking the same fallback
path means zero advantage variance, so GRPO can never escape.

Five changes to break the loop:

1) Few-shot prompt: add two example JSON actions inside the system
prompt. The 1B Llama needs a worked example to produce valid JSON
reliably -- a description alone is not enough. Also adds an
anti-walk_away nudge ("DO NOT walk away unless this is genuinely
the best option").

2) Constructive fallback: when JSON parsing fails, default to a
make_offer with the current draft terms instead of walk_away.
This generates reward variance when the Client responds, giving
GRPO something to learn from.

3) Action-type normalization: lowercase, strip, expand common
abbreviations ("offer" -> "make_offer", "counter" -> "counter_offer",
etc.) before validating against the allowed set.

4) Temperature schedule: ramp from 1.0 down to 0.7 over training.
Higher temperature early forces the 4 rollouts to diverge and
produce variance; lower temperature later locks in the policy.

5) Diagnostic dump on step 1: prints raw model output for first 2
turns of episode 1 so we can actually see what the model is
producing instead of inferring it from turns=1 / tom=0.00.

notebooks/v3_training.ipynb

@@ -174,7 +174,7 @@
    "id": "fccce752",
    "metadata": {},
    "outputs": [],
-   "source": "# Vendor prompt: SINGLE JSON action with optional inline tom_prediction.\n# (Original V3 design used two blocks separated by '---', but a 1B model\n# at temperature 0.8 fumbles the two-block format ~30% of the time, which\n# trapped training in a constant walk_away fallback. The single-block\n# format is what V2 used and what evaluation/llm_vendor's act_simple()\n# expects. The env still grades ToM whenever tom_prediction is present\n# inside the action -- see server/environment.py.)\n\nVENDOR_SYSTEM = '''You are an expert negotiation lawyer representing the VENDOR.\n\nYour private brief:\n{brief}\n\nContract context:\n{ctx}\n\nOpen issues:\n{issues}\n\nNegotiation history:\n{history}\n\nCurrent turn: {turn} of {max_turns}\n\nRespond with ONE JSON action. The \"tom_prediction\" field is OPTIONAL but\nrecommended -- it lets the environment grade your model of the Client\nand earns a per-turn bonus when accurate.\n\n{{\n  \"action_type\": \"make_offer|accept_offer|counter_offer|concede|demand|walk_away|ask_clarification\",\n  \"proposed_terms\": {{...}},\n  \"issue_name\": \"<for concede/demand>\",\n  \"new_value\": <value>,\n  \"reasoning\": \"<one sentence>\",\n  \"tom_prediction\": {{\n    \"predicted_client_priority\": \"<which issue you think Client cares most about>\",\n    \"predicted_client_walkaway\": {{...}},\n    \"predicted_next_action\": \"make_offer|accept_offer|counter_offer|concede|demand|walk_away\",\n    \"confidence\": <0.0-1.0>,\n    \"reasoning\": \"<one sentence>\"\n  }}\n}}\n\nOutput ONLY the JSON. No fences, no other text.'''\n\n\ndef fmt_brief(b):\n    if not b: return \"(no brief)\"\n    parts = [f\"Role: {b.get('role')}\"]\n    if b.get('strategy_hints'): parts.append(f\"Strategy: {b['strategy_hints']}\")\n    parts.append(\"Walkaways: \" + json.dumps(b.get('walkaway_thresholds', {})))\n    parts.append(\"Ideals: \" + json.dumps(b.get('ideal_outcomes', {})))\n    parts.append(\"Priorities: \" + json.dumps(b.get('private_priorities', {})))\n    parts.append(\"Dealbreakers: \" + json.dumps(b.get('dealbreakers', [])))\n    return \"\\n\".join(parts)\n\ndef fmt_issues(issues):\n    return \"\\n\".join(f\"- {i.get('name')}: current={i.get('current_value')} ({i.get('type')})\"\n                      for i in issues or []) or \"(none)\"\n\ndef fmt_history(h):\n    if not h: return \"(no history)\"\n    return \"\\n\".join(\n        f\"[{e.get('agent_role')}] {e.get('action_type')} terms={json.dumps(e.get('proposed_terms') or {})}\"\n        for e in h[-12:]\n    )\n\ndef build_vendor_prompt(obs, turn):\n    return VENDOR_SYSTEM.format(\n        brief=fmt_brief(obs.get('vendor_private_brief')),\n        ctx=obs.get('contract_draft', ''),\n        issues=fmt_issues(obs.get('open_issues', [])),\n        history=fmt_history(obs.get('offer_history', [])),\n        turn=turn,\n        max_turns=int(obs.get('max_turns', 30)),\n    )\n\n\n# ---- Single-block JSON parser ----\n# Parses one JSON action object. If the model included `tom_prediction`\n# inside it, we lift it out so the env's step handler grades it.\n\ndef _strip_fences(s):\n    s = s.strip()\n    m = re.search(r\"```(?:json)?\\s*([\\s\\S]*?)```\", s)\n    return m.group(1).strip() if m else s\n\ndef _balanced_objects(s, n=2):\n    out, i = [], 0\n    while i < len(s) and len(out) < n:\n        if s[i] != '{': i += 1; continue\n        depth = 0; in_str = False; esc = False; end = None\n        for j in range(i, len(s)):\n            ch = s[j]\n            if in_str:\n                if esc: esc = False\n                elif ch == '\\\\': esc = True\n                elif ch == '\"': in_str = False\n                continue\n            if ch == '\"': in_str = True\n            elif ch == '{': depth += 1\n            elif ch == '}':\n                depth -= 1\n                if depth == 0: end = j; break\n        if end is None: break\n        cand = s[i:end+1]\n        try: obj = json.loads(cand)\n        except json.JSONDecodeError:\n            try: obj = json.loads(re.sub(r\",\\s*([}\\]])\", r\"\\1\", cand))\n            except json.JSONDecodeError: i = end + 1; continue\n        if isinstance(obj, dict): out.append(obj)\n        i = end + 1\n    return out\n\n\nFALLBACK_ACTION = {\"action_type\": \"walk_away\", \"agent_role\": \"vendor\",\n                   \"reasoning\": \"fallback: parse failed\"}\n\n\n# Server-side ActionRequest (server/app.py) uses pydantic with extra='forbid'\n# -- any unknown key triggers a 422. We strip down to just the allowed keys.\n_VALID_ACTION_KEYS = {\n    \"action_type\", \"episode_id\", \"turn_number\", \"agent_role\",\n    \"proposed_terms\", \"issue_name\", \"new_value\", \"reasoning\",\n    \"metadata\", \"tom_prediction\",\n}\n_VALID_ACTION_TYPES = {\n    \"make_offer\", \"accept_offer\", \"counter_offer\",\n    \"concede\", \"demand\", \"walk_away\", \"ask_clarification\",\n}\n\n\ndef sanitize_action(raw):\n    \"\"\"Drop unknown keys, coerce types, fall back to walk_away on garbage.\n\n    Mirrors the V2 notebook's sanitize_action -- the server uses pydantic\n    with extra='forbid', so any unexpected key triggers a 422 even if the\n    rest of the action is valid. Sanitizing here turns parse-noise into\n    a clean walk_away instead of a destructive 422 in the training loop.\n    \"\"\"\n    if not isinstance(raw, dict):\n        return {\"action_type\": \"walk_away\", \"reasoning\": \"non-dict action\"}\n\n    cleaned = {k: v for k, v in raw.items() if k in _VALID_ACTION_KEYS}\n\n    atype = cleaned.get(\"action_type\")\n    if not isinstance(atype, str) or atype not in _VALID_ACTION_TYPES:\n        return {\"action_type\": \"walk_away\",\n                \"reasoning\": f\"invalid action_type: {atype!r}\"}\n\n    pt = cleaned.get(\"proposed_terms\")\n    if pt is not None:\n        if not isinstance(pt, dict):\n            cleaned.pop(\"proposed_terms\", None)\n        else:\n            cleaned[\"proposed_terms\"] = {\n                str(k): v for k, v in pt.items()\n                if isinstance(v, (int, float, str, bool))\n            } or None\n            if cleaned[\"proposed_terms\"] is None:\n                cleaned.pop(\"proposed_terms\", None)\n\n    if cleaned.get(\"issue_name\") is not None and not isinstance(cleaned[\"issue_name\"], str):\n        cleaned.pop(\"issue_name\", None)\n\n    nv = cleaned.get(\"new_value\")\n    if nv is not None and not isinstance(nv, (int, float, str, bool)):\n        cleaned.pop(\"new_value\", None)\n\n    if cleaned.get(\"reasoning\") is not None:\n        cleaned[\"reasoning\"] = str(cleaned[\"reasoning\"])[:400]\n\n    # tom_prediction must be a dict if present.\n    tp = cleaned.get(\"tom_prediction\")\n    if tp is not None and not isinstance(tp, dict):\n        cleaned.pop(\"tom_prediction\", None)\n\n    return cleaned\n\n\ndef parse_action(text):\n    \"\"\"Extract a single JSON action object from the model's response.\n\n    Returns a sanitized dict with action_type guaranteed valid; falls back\n    to walk_away on any parse failure rather than raising.\n    \"\"\"\n    if not text:\n        return dict(FALLBACK_ACTION)\n    s = _strip_fences(text)\n    objs = _balanced_objects(s, 1)\n    if not objs:\n        return dict(FALLBACK_ACTION)\n    return sanitize_action(objs[0])\n"
+   "source": "# Vendor prompt: SINGLE JSON action with optional inline tom_prediction.\n# Includes a few-shot example because the 1B Llama needs one to produce\n# valid JSON consistently. Without the example, every response is\n# malformed and the parser falls back to walk_away, trapping training.\n\nVENDOR_SYSTEM = '''You are an expert negotiation lawyer representing the VENDOR.\n\nYour private brief:\n{brief}\n\nContract context:\n{ctx}\n\nOpen issues:\n{issues}\n\nNegotiation history:\n{history}\n\nCurrent turn: {turn} of {max_turns}\n\nRespond with ONE JSON action. Valid action_types are: make_offer,\naccept_offer, counter_offer, concede, demand, walk_away, ask_clarification.\n\nExample valid response (this is what your output should look like):\n{{\"action_type\": \"counter_offer\", \"proposed_terms\": {{\"annual_price\": 95000, \"payment_net_days\": 30}}, \"reasoning\": \"Pushing back on price; client opened too low.\"}}\n\nAnother example with theory-of-mind prediction:\n{{\"action_type\": \"make_offer\", \"proposed_terms\": {{\"annual_price\": 110000}}, \"reasoning\": \"Anchor high.\", \"tom_prediction\": {{\"predicted_client_priority\": \"annual_price\", \"predicted_client_walkaway\": {{\"annual_price\": 90000}}, \"predicted_next_action\": \"counter_offer\", \"confidence\": 0.7, \"reasoning\": \"Client emphasized budget.\"}}}}\n\nNow respond with YOUR JSON action for the current turn. Output ONLY the\nJSON object, nothing else. No fences, no preamble, no explanation.\n\nDO NOT walk away unless this is genuinely the best option. Try to make\nor counter offers to advance the negotiation.\n\nYour JSON action:'''\n\n\ndef fmt_brief(b):\n    if not b: return \"(no brief)\"\n    parts = [f\"Role: {b.get('role')}\"]\n    if b.get('strategy_hints'): parts.append(f\"Strategy: {b['strategy_hints']}\")\n    parts.append(\"Walkaways: \" + json.dumps(b.get('walkaway_thresholds', {})))\n    parts.append(\"Ideals: \" + json.dumps(b.get('ideal_outcomes', {})))\n    parts.append(\"Priorities: \" + json.dumps(b.get('private_priorities', {})))\n    parts.append(\"Dealbreakers: \" + json.dumps(b.get('dealbreakers', [])))\n    return \"\\n\".join(parts)\n\ndef fmt_issues(issues):\n    return \"\\n\".join(f\"- {i.get('name')}: current={i.get('current_value')} ({i.get('type')})\"\n                      for i in issues or []) or \"(none)\"\n\ndef fmt_history(h):\n    if not h: return \"(no history)\"\n    return \"\\n\".join(\n        f\"[{e.get('agent_role')}] {e.get('action_type')} terms={json.dumps(e.get('proposed_terms') or {})}\"\n        for e in h[-12:]\n    )\n\ndef build_vendor_prompt(obs, turn):\n    return VENDOR_SYSTEM.format(\n        brief=fmt_brief(obs.get('vendor_private_brief')),\n        ctx=obs.get('contract_draft', ''),\n        issues=fmt_issues(obs.get('open_issues', [])),\n        history=fmt_history(obs.get('offer_history', [])),\n        turn=turn,\n        max_turns=int(obs.get('max_turns', 30)),\n    )\n\n\n# ---- Single-block JSON parser, V2-grade forgiving ----\n\ndef _strip_fences(s):\n    s = s.strip()\n    m = re.search(r\"```(?:json)?\\s*([\\s\\S]*?)```\", s)\n    return m.group(1).strip() if m else s\n\ndef _balanced_objects(s, n=2):\n    out, i = [], 0\n    while i < len(s) and len(out) < n:\n        if s[i] != '{': i += 1; continue\n        depth = 0; in_str = False; esc = False; end = None\n        for j in range(i, len(s)):\n            ch = s[j]\n            if in_str:\n                if esc: esc = False\n                elif ch == '\\\\': esc = True\n                elif ch == '\"': in_str = False\n                continue\n            if ch == '\"': in_str = True\n            elif ch == '{': depth += 1\n            elif ch == '}':\n                depth -= 1\n                if depth == 0: end = j; break\n        if end is None: break\n        cand = s[i:end+1]\n        try: obj = json.loads(cand)\n        except json.JSONDecodeError:\n            try: obj = json.loads(re.sub(r\",\\s*([}\\]])\", r\"\\1\", cand))\n            except json.JSONDecodeError: i = end + 1; continue\n        if isinstance(obj, dict): out.append(obj)\n        i = end + 1\n    return out\n\n\n_VALID_ACTION_KEYS = {\n    \"action_type\", \"episode_id\", \"turn_number\", \"agent_role\",\n    \"proposed_terms\", \"issue_name\", \"new_value\", \"reasoning\",\n    \"metadata\", \"tom_prediction\",\n}\n_VALID_ACTION_TYPES = {\n    \"make_offer\", \"accept_offer\", \"counter_offer\",\n    \"concede\", \"demand\", \"walk_away\", \"ask_clarification\",\n}\n\n\ndef _build_default_make_offer(obs):\n    \"\"\"Constructive fallback: a make_offer using the current draft terms.\n\n    Why this is much better than walking away: every rollout that walks away\n    on turn 1 gets identical reward, killing the GRPO advantage signal. A\n    make_offer with current terms keeps the negotiation alive long enough\n    for the Client to respond, generating reward variance the trainer can use.\n    \"\"\"\n    open_issues = obs.get(\"open_issues\") or []\n    proposed = {}\n    for i in open_issues:\n        name = i.get(\"name\")\n        cur = i.get(\"current_value\")\n        if name and cur is not None:\n            proposed[str(name)] = cur\n    if not proposed:\n        # Fall back to an ask_clarification (also produces variance vs walk_away).\n        return {\"action_type\": \"ask_clarification\",\n                \"reasoning\": \"constructive fallback: no terms available\"}\n    return {\"action_type\": \"make_offer\",\n            \"proposed_terms\": proposed,\n            \"reasoning\": \"constructive fallback: anchor on draft terms\"}\n\n\ndef sanitize_action(raw, obs=None):\n    \"\"\"Drop unknown keys, coerce types, fall back to make_offer (NOT walk_away).\n\n    The fallback is deliberately a make_offer or ask_clarification because\n    the constant walk_away fallback was the root cause of the training-time\n    GRPO collapse: 4 rollouts -> 4 walk_aways -> identical reward -> zero\n    advantage variance -> no gradient -> no learning.\n    \"\"\"\n    if not isinstance(raw, dict):\n        return _build_default_make_offer(obs or {}) if obs is not None else \\\n               {\"action_type\": \"ask_clarification\", \"reasoning\": \"non-dict action\"}\n\n    cleaned = {k: v for k, v in raw.items() if k in _VALID_ACTION_KEYS}\n\n    # Normalize action_type: lowercase, strip, common-typo fixes.\n    atype_raw = cleaned.get(\"action_type\")\n    if isinstance(atype_raw, str):\n        atype = atype_raw.strip().lower().replace(\"-\", \"_\").replace(\" \", \"_\")\n        # Common abbreviations / near-misses.\n        atype = {\n            \"offer\": \"make_offer\",\n            \"counter\": \"counter_offer\",\n            \"accept\": \"accept_offer\",\n            \"reject\": \"walk_away\",\n            \"walkaway\": \"walk_away\",\n            \"clarify\": \"ask_clarification\",\n            \"ask\": \"ask_clarification\",\n        }.get(atype, atype)\n        cleaned[\"action_type\"] = atype\n    else:\n        atype = None\n\n    if atype not in _VALID_ACTION_TYPES:\n        # Don't fall back to walk_away. Build a constructive default.\n        if obs is not None:\n            return _build_default_make_offer(obs)\n        return {\"action_type\": \"ask_clarification\",\n                \"reasoning\": f\"invalid action_type: {atype_raw!r}\"}\n\n    pt = cleaned.get(\"proposed_terms\")\n    if pt is not None:\n        if not isinstance(pt, dict):\n            cleaned.pop(\"proposed_terms\", None)\n        else:\n            cleaned[\"proposed_terms\"] = {\n                str(k): v for k, v in pt.items()\n                if isinstance(v, (int, float, str, bool))\n            } or None\n            if cleaned[\"proposed_terms\"] is None:\n                cleaned.pop(\"proposed_terms\", None)\n\n    if cleaned.get(\"issue_name\") is not None and not isinstance(cleaned[\"issue_name\"], str):\n        cleaned.pop(\"issue_name\", None)\n\n    nv = cleaned.get(\"new_value\")\n    if nv is not None and not isinstance(nv, (int, float, str, bool)):\n        cleaned.pop(\"new_value\", None)\n\n    if cleaned.get(\"reasoning\") is not None:\n        cleaned[\"reasoning\"] = str(cleaned[\"reasoning\"])[:400]\n\n    tp = cleaned.get(\"tom_prediction\")\n    if tp is not None and not isinstance(tp, dict):\n        cleaned.pop(\"tom_prediction\", None)\n\n    return cleaned\n\n\ndef parse_action(text, obs=None):\n    \"\"\"Extract a JSON action from the model output. Constructive fallback.\n\n    NOTE: the obs argument is required for the constructive fallback to\n    produce a make_offer with the current draft terms when JSON parsing\n    fails. Without obs, falls back to ask_clarification.\n    \"\"\"\n    if not text:\n        return _build_default_make_offer(obs) if obs is not None else \\\n               {\"action_type\": \"ask_clarification\", \"reasoning\": \"empty text\"}\n    s = _strip_fences(text)\n    objs = _balanced_objects(s, 1)\n    if not objs:\n        return _build_default_make_offer(obs) if obs is not None else \\\n               {\"action_type\": \"ask_clarification\", \"reasoning\": \"no JSON found\"}\n    return sanitize_action(objs[0], obs=obs)\n"
   },
   {
    "cell_type": "markdown",
@@ -190,7 +190,7 @@
    "id": "643d5625",
    "metadata": {},
    "outputs": [],
-   "source": "@dataclass\nclass TurnRecord:\n    prompt: str\n    completion: str\n    completion_ids: List[int]\n    prompt_ids: List[int]\n    reward: float\n    tom_acc: Optional[float] = None\n    done: bool = False\n\n\n@dataclass\nclass EpisodeRecord:\n    task_id: str\n    turns: List[TurnRecord] = field(default_factory=list)\n    deal_struck: bool = False\n    total_reward: float = 0.0\n    tribunal: Optional[Dict[str, float]] = None\n    error: Optional[str] = None\n\n    @property\n    def episode_reward(self):\n        if not self.turns: return 0.0\n        return float(self.turns[-1].reward)\n\n\n@torch.no_grad()\ndef generate_completion(model, tokenizer, prompt, max_new_tokens, temperature=0.7):\n    enc = tokenizer(prompt, return_tensors=\"pt\", truncation=True,\n                    max_length=CONFIG[\"max_seq_length\"] - max_new_tokens).to(model.device)\n    out = model.generate(\n        **enc,\n        max_new_tokens=max_new_tokens,\n        temperature=temperature,\n        do_sample=temperature > 0,\n        top_p=0.9,\n        pad_token_id=tokenizer.eos_token_id,\n    )\n    prompt_ids = enc.input_ids[0].tolist()\n    full_ids = out[0].tolist()\n    completion_ids = full_ids[len(prompt_ids):]\n    text = tokenizer.decode(completion_ids, skip_special_tokens=True)\n    return text, prompt_ids, completion_ids\n\n\ndef play_episode(model, tokenizer, env, task_id, seed=None, temperature=0.7):\n    rec = EpisodeRecord(task_id=task_id)\n    try:\n        obs = env.reset(task_id=task_id, seed=seed, client_strategy=\"balanced\")\n    except Exception as e:\n        rec.error = f\"reset_failed: {e}\"; return rec\n\n    episode_id = obs.get(\"episode_id\", \"\")\n    max_steps = int(obs.get(\"max_turns\", 30)) + CONFIG[\"max_turns_safety_margin\"]\n\n    for step in range(max_steps):\n        if obs.get(\"done\"): break\n        prompt = build_vendor_prompt(obs, turn=step + 1)\n        try:\n            text, p_ids, c_ids = generate_completion(\n                model, tokenizer, prompt,\n                max_new_tokens=CONFIG[\"max_completion_length\"], temperature=temperature)\n        except Exception as e:\n            rec.error = f\"generate_failed: {e}\"; break\n\n        # Single-block parse; sanitized action with action_type guaranteed valid.\n        action = parse_action(text)\n        # parse_action() may have included tom_prediction inside the action --\n        # the env's step handler accepts it as a top-level field too.\n        action.setdefault(\"episode_id\", episode_id)\n        action.setdefault(\"turn_number\", step + 1)\n        action.setdefault(\"agent_role\", \"vendor\")\n\n        # Step with 422-aware retry: pydantic rejection -> send a clean walk_away\n        # exactly once, then end the episode. Avoids the destructive loss spikes\n        # we saw when 422s were propagating to the trainer.\n        try:\n            obs = env.step(action)\n        except httpx.HTTPStatusError as he:\n            if he.response is not None and he.response.status_code == 422:\n                # Try a sanitized walk_away as a graceful end.\n                try:\n                    obs = env.step({\n                        \"action_type\": \"walk_away\",\n                        \"episode_id\": episode_id,\n                        \"turn_number\": step + 1,\n                        \"agent_role\": \"vendor\",\n                        \"reasoning\": \"fallback after 422\",\n                    })\n                except Exception as e2:\n                    rec.error = f\"step_422_then_failed: {e2}\"; break\n            else:\n                rec.error = f\"step_failed: {he}\"; break\n        except Exception as e:\n            rec.error = f\"step_failed: {e}\"; break\n\n        reward = float(obs.get(\"reward\") or 0.0)\n        feedback = obs.get(\"tom_grader_feedback\") or {}\n        tom_acc = float(feedback.get(\"composite\", 0.0)) if feedback else None\n        rec.turns.append(TurnRecord(\n            prompt=prompt, completion=text,\n            prompt_ids=p_ids, completion_ids=c_ids,\n            reward=reward, tom_acc=tom_acc, done=bool(obs.get(\"done\")),\n        ))\n        if obs.get(\"done\"):\n            rec.deal_struck = bool(obs.get(\"final_deal_struck\"))\n            rec.tribunal = obs.get(\"tribunal_breakdown\")\n            break\n\n    rec.total_reward = rec.episode_reward\n    return rec\n"
+   "source": "@dataclass\nclass TurnRecord:\n    prompt: str\n    completion: str\n    completion_ids: List[int]\n    prompt_ids: List[int]\n    reward: float\n    tom_acc: Optional[float] = None\n    done: bool = False\n\n\n@dataclass\nclass EpisodeRecord:\n    task_id: str\n    turns: List[TurnRecord] = field(default_factory=list)\n    deal_struck: bool = False\n    total_reward: float = 0.0\n    tribunal: Optional[Dict[str, float]] = None\n    error: Optional[str] = None\n\n    @property\n    def episode_reward(self):\n        if not self.turns: return 0.0\n        return float(self.turns[-1].reward)\n\n\n@torch.no_grad()\ndef generate_completion(model, tokenizer, prompt, max_new_tokens, temperature=0.7):\n    enc = tokenizer(prompt, return_tensors=\"pt\", truncation=True,\n                    max_length=CONFIG[\"max_seq_length\"] - max_new_tokens).to(model.device)\n    out = model.generate(\n        **enc,\n        max_new_tokens=max_new_tokens,\n        temperature=temperature,\n        do_sample=temperature > 0,\n        top_p=0.9,\n        pad_token_id=tokenizer.eos_token_id,\n    )\n    prompt_ids = enc.input_ids[0].tolist()\n    full_ids = out[0].tolist()\n    completion_ids = full_ids[len(prompt_ids):]\n    text = tokenizer.decode(completion_ids, skip_special_tokens=True)\n    return text, prompt_ids, completion_ids\n\n\ndef play_episode(model, tokenizer, env, task_id, seed=None, temperature=0.7, verbose=False):\n    \"\"\"Play one episode end-to-end.\n\n    verbose=True prints raw model completions for the first 2 turns and\n    422 details. Use this on step 1 of training to verify the model is\n    actually producing JSON.\n    \"\"\"\n    rec = EpisodeRecord(task_id=task_id)\n    try:\n        obs = env.reset(task_id=task_id, seed=seed, client_strategy=\"balanced\")\n    except Exception as e:\n        rec.error = f\"reset_failed: {e}\"; return rec\n\n    episode_id = obs.get(\"episode_id\", \"\")\n    max_steps = int(obs.get(\"max_turns\", 30)) + CONFIG[\"max_turns_safety_margin\"]\n\n    for step in range(max_steps):\n        if obs.get(\"done\"): break\n        prompt = build_vendor_prompt(obs, turn=step + 1)\n        try:\n            text, p_ids, c_ids = generate_completion(\n                model, tokenizer, prompt,\n                max_new_tokens=CONFIG[\"max_completion_length\"], temperature=temperature)\n        except Exception as e:\n            rec.error = f\"generate_failed: {e}\"; break\n\n        if verbose and step < 2:\n            print(f\"    [diag turn {step+1}] raw model output (first 300 chars):\")\n            print(f\"      {text[:300]!r}\")\n\n        # Parse with constructive fallback (NOT walk_away).\n        action = parse_action(text, obs=obs)\n        action.setdefault(\"episode_id\", episode_id)\n        action.setdefault(\"turn_number\", step + 1)\n        action.setdefault(\"agent_role\", \"vendor\")\n\n        if verbose and step < 2:\n            print(f\"    [diag turn {step+1}] parsed action_type={action.get('action_type')!r} \"\n                  f\"proposed_terms={action.get('proposed_terms')}\")\n\n        # Step with 422-aware retry.\n        try:\n            obs = env.step(action)\n        except httpx.HTTPStatusError as he:\n            if he.response is not None and he.response.status_code == 422:\n                if verbose:\n                    body = \"\"\n                    try: body = he.response.text[:200]\n                    except: pass\n                    print(f\"    [diag] 422 from server. action sent: {action}\")\n                    print(f\"    [diag] 422 response body: {body!r}\")\n                try:\n                    obs = env.step({\n                        \"action_type\": \"walk_away\",\n                        \"episode_id\": episode_id,\n                        \"turn_number\": step + 1,\n                        \"agent_role\": \"vendor\",\n                        \"reasoning\": \"fallback after 422\",\n                    })\n                except Exception as e2:\n                    rec.error = f\"step_422_then_failed: {e2}\"; break\n            else:\n                rec.error = f\"step_failed: {he}\"; break\n        except Exception as e:\n            rec.error = f\"step_failed: {e}\"; break\n\n        reward = float(obs.get(\"reward\") or 0.0)\n        feedback = obs.get(\"tom_grader_feedback\") or {}\n        tom_acc = float(feedback.get(\"composite\", 0.0)) if feedback else None\n        rec.turns.append(TurnRecord(\n            prompt=prompt, completion=text,\n            prompt_ids=p_ids, completion_ids=c_ids,\n            reward=reward, tom_acc=tom_acc, done=bool(obs.get(\"done\")),\n        ))\n        if obs.get(\"done\"):\n            rec.deal_struck = bool(obs.get(\"final_deal_struck\"))\n            rec.tribunal = obs.get(\"tribunal_breakdown\")\n            break\n\n    rec.total_reward = rec.episode_reward\n    return rec\n"
   },
   {
    "cell_type": "markdown",
@@ -268,7 +268,7 @@
    "id": "53dff448",
    "metadata": {},
    "outputs": [],
-   "source": "# GRPO core trick: at each step, play N rollouts with the current policy, then\n# z-score-normalize per-episode returns to get advantages. The policy gradient\n# is sum_t [advantage * grad log p(action_t | prompt_t)] with a small KL anchor\n# that pulls toward the LoRA-disabled (reference) distribution.\n\nFastLanguageModel.for_training(model)\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad],\n    lr=CONFIG[\"learning_rate\"],\n)\n\n\ndef compute_completion_logprobs(model, prompt_ids, completion_ids):\n    if not completion_ids:\n        return torch.tensor(0.0, device=model.device, requires_grad=True)\n    full_ids = torch.tensor([prompt_ids + completion_ids], device=model.device)\n    out = model(full_ids)\n    logits = out.logits[0, len(prompt_ids) - 1: -1]\n    target = torch.tensor(completion_ids, device=model.device)\n    log_probs = F.log_softmax(logits, dim=-1)\n    chosen = log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1)\n    return chosen.sum()\n\n\ndef kl_to_reference(model, prompt_ids, completion_ids):\n    if not completion_ids:\n        return torch.tensor(0.0, device=model.device)\n    full_ids = torch.tensor([prompt_ids + completion_ids], device=model.device)\n    pol_logits = model(full_ids).logits[0, len(prompt_ids) - 1: -1]\n    with torch.no_grad():\n        with model.disable_adapter():\n            ref_logits = model(full_ids).logits[0, len(prompt_ids) - 1: -1]\n    pol_logp = F.log_softmax(pol_logits, dim=-1)\n    ref_logp = F.log_softmax(ref_logits, dim=-1)\n    pol_p = pol_logp.exp()\n    return (pol_p * (pol_logp - ref_logp)).sum()\n\n\nclass _SimpleCurriculum:\n    \"\"\"Coarse mirror of server.curriculum.CurriculumScheduler for training-side task selection.\"\"\"\n    def __init__(self, tasks, threshold=0.65, window=20):\n        self.tasks = tasks; self.threshold = threshold; self.window = window\n        self.history = []; self.level = 0\n    def record(self, task, deal):\n        self.history.append((task, deal))\n        if self.level < len(self.tasks) - 1:\n            recent = [d for (t, d) in reversed(self.history)\n                      if t == self.tasks[self.level]][:self.window]\n            if len(recent) >= self.window and sum(recent) / len(recent) >= self.threshold:\n                self.level += 1\n                print(f\"  [curriculum] advancing to level {self.level}: {self.tasks[self.level]}\")\n    def select(self, rng):\n        if self.level == 0: return self.tasks[0]\n        if rng.random() < 0.20: return self.tasks[rng.randint(0, self.level - 1)]\n        return self.tasks[self.level]\n\n\ncurriculum = _SimpleCurriculum(CONFIG[\"tasks\"])\nrng = random.Random(CONFIG[\"seed\"])\n\n# TensorBoard writer -- streams scalars to the live dashboard above.\nwriter = SummaryWriter(log_dir=CONFIG[\"tb_logdir\"])\nwriter.add_text(\"config\", json.dumps({k: str(v) for k, v in CONFIG.items()}, indent=2))\n\n\ndef save_checkpoint(step):\n    path = Path(CONFIG[\"checkpoint_dir\"]) / f\"step_{step}\"\n    model.save_pretrained(str(path))\n    tokenizer.save_pretrained(str(path))\n    print(f\"  [checkpoint] saved {path}\")\n\n\nt_start = time.time()\nn_skipped_no_variance = 0\nn_skipped_all_errored = 0\n\nfor step in range(1, CONFIG[\"max_steps\"] + 1):\n    FastLanguageModel.for_inference(model)\n    task = curriculum.select(rng) if CONFIG[\"use_curriculum\"] else \\\n           CONFIG[\"tasks\"][rng.randrange(len(CONFIG[\"tasks\"]))]\n\n    print(f\"\\n--- step {step}/{CONFIG['max_steps']}  task={task} ---\")\n\n    episodes = []\n    for ep in range(CONFIG[\"num_episodes_per_step\"]):\n        seed = step * 1000 + ep\n        ep_rec = play_episode(model, tokenizer, env, task,\n                              seed=seed, temperature=0.8)\n        episodes.append(ep_rec)\n        curriculum.record(task, ep_rec.deal_struck)\n        ep_tom_accs = [t.tom_acc for t in ep_rec.turns if t.tom_acc is not None]\n        ep_tom = float(np.mean(ep_tom_accs)) if ep_tom_accs else 0.0\n        err_tag = f\" err={ep_rec.error[:60]}\" if ep_rec.error else \"\"\n        print(f\"  ep{ep+1}/{CONFIG['num_episodes_per_step']}  \"\n              f\"reward={ep_rec.episode_reward:.3f}  \"\n              f\"deal={'Y' if ep_rec.deal_struck else 'N'}  \"\n              f\"turns={len(ep_rec.turns)}  tom={ep_tom:.2f}{err_tag}\")\n\n    # ---- Skip the policy update when the batch is degenerate ----\n    # Two ways an update is destructive or wasted:\n    #   1) Every episode errored -> no usable trajectories.\n    #   2) All episodes got the same reward -> z-score advantages all zero,\n    #      gradient signal is exactly zero. Burning compute for nothing.\n    valid_eps = [e for e in episodes if not e.error and e.turns]\n    returns = np.array([e.episode_reward for e in valid_eps], dtype=np.float32)\n\n    skip_update_reason = None\n    if not valid_eps:\n        skip_update_reason = \"all_errored\"\n        n_skipped_all_errored += 1\n    elif len(returns) < 2 or float(returns.std()) < 1e-6:\n        skip_update_reason = \"no_variance\"\n        n_skipped_no_variance += 1\n\n    if skip_update_reason is None:\n        advantages = (returns - returns.mean()) / (returns.std() + 1e-8)\n\n        FastLanguageModel.for_training(model)\n        optimizer.zero_grad()\n        total_loss = torch.tensor(0.0, device=model.device)\n        n_turns_used = 0\n        for ep_idx, ep in enumerate(valid_eps):\n            adv = float(advantages[ep_idx])\n            for turn in ep.turns:\n                if not turn.completion_ids: continue\n                logp = compute_completion_logprobs(model, turn.prompt_ids, turn.completion_ids)\n                kl = kl_to_reference(model, turn.prompt_ids, turn.completion_ids)\n                total_loss = total_loss + (-adv * logp + CONFIG[\"kl_beta\"] * kl)\n                n_turns_used += 1\n\n        if n_turns_used > 0:\n            loss = total_loss / max(1, n_turns_used)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(\n                [p for p in model.parameters() if p.requires_grad], CONFIG[\"grad_clip\"])\n            optimizer.step()\n            loss_val = float(loss.detach())\n        else:\n            loss_val = 0.0\n    else:\n        loss_val = 0.0  # No update happened.\n\n    # Stats are computed across ALL episodes (errored or not) for the curve.\n    all_returns = np.array([e.episode_reward for e in episodes], dtype=np.float32)\n    mean_reward = float(all_returns.mean()) if len(all_returns) else 0.0\n    deal_rate = float(np.mean([1 if e.deal_struck else 0 for e in episodes]))\n    mean_turns = float(np.mean([len(e.turns) for e in episodes]))\n    tom_accs = [t.tom_acc for e in episodes for t in e.turns if t.tom_acc is not None]\n    mean_tom = float(np.mean(tom_accs)) if tom_accs else 0.0\n    tribunal_means = {k: 0.0 for k in (\"pro_vendor_score\", \"pro_client_score\", \"neutral_score\", \"trimmed_mean\")}\n    tribunal_count = 0\n    for e in episodes:\n        if e.tribunal:\n            tribunal_count += 1\n            for k in tribunal_means: tribunal_means[k] += float(e.tribunal.get(k, 0.0))\n    if tribunal_count:\n        for k in tribunal_means: tribunal_means[k] /= tribunal_count\n\n    entry = {\n        \"step\": step, \"task\": task,\n        \"mean_reward\": mean_reward, \"deal_rate\": deal_rate,\n        \"mean_turns\": mean_turns, \"tom_accuracy\": mean_tom,\n        \"tribunal\": tribunal_means, \"loss\": loss_val,\n        \"curriculum_level\": curriculum.level,\n        \"skipped\": skip_update_reason,\n    }\n    training_history[\"steps\"].append(entry)\n\n    # Stream scalars to TensorBoard.\n    writer.add_scalar(\"train/mean_reward\", mean_reward, step)\n    writer.add_scalar(\"train/deal_rate\", deal_rate, step)\n    writer.add_scalar(\"train/mean_turns\", mean_turns, step)\n    writer.add_scalar(\"train/tom_accuracy\", mean_tom, step)\n    writer.add_scalar(\"train/loss\", loss_val, step)\n    writer.add_scalar(\"train/curriculum_level\", curriculum.level, step)\n    if tribunal_count:\n        for k, v in tribunal_means.items():\n            writer.add_scalar(f\"tribunal/{k}\", v, step)\n    writer.flush()\n\n    elapsed = time.time() - t_start\n    tribunal_tag = \"\"\n    if tribunal_count:\n        tribunal_tag = f\"  tribunal={tribunal_means['trimmed_mean']:.2f}\"\n    skip_tag = f\"  [SKIPPED:{skip_update_reason}]\" if skip_update_reason else \"\"\n    print(f\"  step summary: reward={mean_reward:.3f}  deal={deal_rate:.0%}  \"\n          f\"turns={mean_turns:.1f}  tom={mean_tom:.2f}  loss={loss_val:.3f}\"\n          f\"{tribunal_tag}  elapsed={elapsed/60:.1f}m{skip_tag}\")\n\n    if step % CONFIG[\"checkpoint_every\"] == 0:\n        save_checkpoint(step)\n        with open(CONFIG[\"history_path\"], \"w\") as f:\n            json.dump(training_history, f, indent=2, default=str)\n\nwriter.close()\nprint(f\"\\nTraining done in {(time.time() - t_start)/60:.1f} minutes.\")\nprint(f\"Steps skipped: no_variance={n_skipped_no_variance}, all_errored={n_skipped_all_errored}\")\n"
+   "source": "# GRPO core trick: at each step, play N rollouts with the current policy, then\n# z-score-normalize per-episode returns to get advantages. The policy gradient\n# is sum_t [advantage * grad log p(action_t | prompt_t)] with a small KL anchor\n# that pulls toward the LoRA-disabled (reference) distribution.\n\nFastLanguageModel.for_training(model)\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad],\n    lr=CONFIG[\"learning_rate\"],\n)\n\n\ndef compute_completion_logprobs(model, prompt_ids, completion_ids):\n    if not completion_ids:\n        return torch.tensor(0.0, device=model.device, requires_grad=True)\n    full_ids = torch.tensor([prompt_ids + completion_ids], device=model.device)\n    out = model(full_ids)\n    logits = out.logits[0, len(prompt_ids) - 1: -1]\n    target = torch.tensor(completion_ids, device=model.device)\n    log_probs = F.log_softmax(logits, dim=-1)\n    chosen = log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1)\n    return chosen.sum()\n\n\ndef kl_to_reference(model, prompt_ids, completion_ids):\n    if not completion_ids:\n        return torch.tensor(0.0, device=model.device)\n    full_ids = torch.tensor([prompt_ids + completion_ids], device=model.device)\n    pol_logits = model(full_ids).logits[0, len(prompt_ids) - 1: -1]\n    with torch.no_grad():\n        with model.disable_adapter():\n            ref_logits = model(full_ids).logits[0, len(prompt_ids) - 1: -1]\n    pol_logp = F.log_softmax(pol_logits, dim=-1)\n    ref_logp = F.log_softmax(ref_logits, dim=-1)\n    pol_p = pol_logp.exp()\n    return (pol_p * (pol_logp - ref_logp)).sum()\n\n\nclass _SimpleCurriculum:\n    \"\"\"Coarse mirror of server.curriculum.CurriculumScheduler for training-side task selection.\"\"\"\n    def __init__(self, tasks, threshold=0.65, window=20):\n        self.tasks = tasks; self.threshold = threshold; self.window = window\n        self.history = []; self.level = 0\n    def record(self, task, deal):\n        self.history.append((task, deal))\n        if self.level < len(self.tasks) - 1:\n            recent = [d for (t, d) in reversed(self.history)\n                      if t == self.tasks[self.level]][:self.window]\n            if len(recent) >= self.window and sum(recent) / len(recent) >= self.threshold:\n                self.level += 1\n                print(f\"  [curriculum] advancing to level {self.level}: {self.tasks[self.level]}\")\n    def select(self, rng):\n        if self.level == 0: return self.tasks[0]\n        if rng.random() < 0.20: return self.tasks[rng.randint(0, self.level - 1)]\n        return self.tasks[self.level]\n\n\ncurriculum = _SimpleCurriculum(CONFIG[\"tasks\"])\nrng = random.Random(CONFIG[\"seed\"])\n\n# TensorBoard writer -- streams scalars to the live dashboard above.\nwriter = SummaryWriter(log_dir=CONFIG[\"tb_logdir\"])\nwriter.add_text(\"config\", json.dumps({k: str(v) for k, v in CONFIG.items()}, indent=2))\n\n\ndef save_checkpoint(step):\n    path = Path(CONFIG[\"checkpoint_dir\"]) / f\"step_{step}\"\n    model.save_pretrained(str(path))\n    tokenizer.save_pretrained(str(path))\n    print(f\"  [checkpoint] saved {path}\")\n\n\n# Per-step temperature: ramp DOWN from 1.0 -> 0.7 over training. Higher\n# temperature early forces exploration so the 4 rollouts diverge and\n# produce reward variance the trainer can use. Without this the 1B\n# model emits walk_away deterministically from the start.\ndef _temp_for_step(s, total):\n    progress = min(1.0, (s - 1) / max(1, total - 1))\n    return 1.0 - 0.3 * progress  # 1.0 at step 1, 0.7 at step max\n\n\nt_start = time.time()\nn_skipped_no_variance = 0\nn_skipped_all_errored = 0\n\nfor step in range(1, CONFIG[\"max_steps\"] + 1):\n    FastLanguageModel.for_inference(model)\n    task = curriculum.select(rng) if CONFIG[\"use_curriculum\"] else \\\n           CONFIG[\"tasks\"][rng.randrange(len(CONFIG[\"tasks\"]))]\n\n    cur_temp = _temp_for_step(step, CONFIG[\"max_steps\"])\n    print(f\"\\n--- step {step}/{CONFIG['max_steps']}  task={task}  temp={cur_temp:.2f} ---\")\n\n    episodes = []\n    for ep in range(CONFIG[\"num_episodes_per_step\"]):\n        seed = step * 1000 + ep\n        # On step 1 only, dump the first episode's raw model output so we\n        # can see exactly what the 1B is producing before any training.\n        ep_verbose = (step == 1 and ep == 0)\n        ep_rec = play_episode(model, tokenizer, env, task,\n                              seed=seed, temperature=cur_temp,\n                              verbose=ep_verbose)\n        episodes.append(ep_rec)\n        curriculum.record(task, ep_rec.deal_struck)\n        ep_tom_accs = [t.tom_acc for t in ep_rec.turns if t.tom_acc is not None]\n        ep_tom = float(np.mean(ep_tom_accs)) if ep_tom_accs else 0.0\n        err_tag = f\" err={ep_rec.error[:60]}\" if ep_rec.error else \"\"\n        print(f\"  ep{ep+1}/{CONFIG['num_episodes_per_step']}  \"\n              f\"reward={ep_rec.episode_reward:.3f}  \"\n              f\"deal={'Y' if ep_rec.deal_struck else 'N'}  \"\n              f\"turns={len(ep_rec.turns)}  tom={ep_tom:.2f}{err_tag}\")\n\n    # ---- Skip the policy update when the batch is degenerate ----\n    valid_eps = [e for e in episodes if not e.error and e.turns]\n    returns = np.array([e.episode_reward for e in valid_eps], dtype=np.float32)\n\n    skip_update_reason = None\n    if not valid_eps:\n        skip_update_reason = \"all_errored\"\n        n_skipped_all_errored += 1\n    elif len(returns) < 2 or float(returns.std()) < 1e-6:\n        skip_update_reason = \"no_variance\"\n        n_skipped_no_variance += 1\n\n    if skip_update_reason is None:\n        advantages = (returns - returns.mean()) / (returns.std() + 1e-8)\n\n        FastLanguageModel.for_training(model)\n        optimizer.zero_grad()\n        total_loss = torch.tensor(0.0, device=model.device)\n        n_turns_used = 0\n        for ep_idx, ep in enumerate(valid_eps):\n            adv = float(advantages[ep_idx])\n            for turn in ep.turns:\n                if not turn.completion_ids: continue\n                logp = compute_completion_logprobs(model, turn.prompt_ids, turn.completion_ids)\n                kl = kl_to_reference(model, turn.prompt_ids, turn.completion_ids)\n                total_loss = total_loss + (-adv * logp + CONFIG[\"kl_beta\"] * kl)\n                n_turns_used += 1\n\n        if n_turns_used > 0:\n            loss = total_loss / max(1, n_turns_used)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(\n                [p for p in model.parameters() if p.requires_grad], CONFIG[\"grad_clip\"])\n            optimizer.step()\n            loss_val = float(loss.detach())\n        else:\n            loss_val = 0.0\n    else:\n        loss_val = 0.0\n\n    all_returns = np.array([e.episode_reward for e in episodes], dtype=np.float32)\n    mean_reward = float(all_returns.mean()) if len(all_returns) else 0.0\n    deal_rate = float(np.mean([1 if e.deal_struck else 0 for e in episodes]))\n    mean_turns = float(np.mean([len(e.turns) for e in episodes]))\n    tom_accs = [t.tom_acc for e in episodes for t in e.turns if t.tom_acc is not None]\n    mean_tom = float(np.mean(tom_accs)) if tom_accs else 0.0\n    tribunal_means = {k: 0.0 for k in (\"pro_vendor_score\", \"pro_client_score\", \"neutral_score\", \"trimmed_mean\")}\n    tribunal_count = 0\n    for e in episodes:\n        if e.tribunal:\n            tribunal_count += 1\n            for k in tribunal_means: tribunal_means[k] += float(e.tribunal.get(k, 0.0))\n    if tribunal_count:\n        for k in tribunal_means: tribunal_means[k] /= tribunal_count\n\n    entry = {\n        \"step\": step, \"task\": task,\n        \"mean_reward\": mean_reward, \"deal_rate\": deal_rate,\n        \"mean_turns\": mean_turns, \"tom_accuracy\": mean_tom,\n        \"tribunal\": tribunal_means, \"loss\": loss_val,\n        \"curriculum_level\": curriculum.level,\n        \"skipped\": skip_update_reason,\n        \"temp\": cur_temp,\n    }\n    training_history[\"steps\"].append(entry)\n\n    writer.add_scalar(\"train/mean_reward\", mean_reward, step)\n    writer.add_scalar(\"train/deal_rate\", deal_rate, step)\n    writer.add_scalar(\"train/mean_turns\", mean_turns, step)\n    writer.add_scalar(\"train/tom_accuracy\", mean_tom, step)\n    writer.add_scalar(\"train/loss\", loss_val, step)\n    writer.add_scalar(\"train/curriculum_level\", curriculum.level, step)\n    writer.add_scalar(\"train/temperature\", cur_temp, step)\n    if tribunal_count:\n        for k, v in tribunal_means.items():\n            writer.add_scalar(f\"tribunal/{k}\", v, step)\n    writer.flush()\n\n    elapsed = time.time() - t_start\n    tribunal_tag = \"\"\n    if tribunal_count:\n        tribunal_tag = f\"  tribunal={tribunal_means['trimmed_mean']:.2f}\"\n    skip_tag = f\"  [SKIPPED:{skip_update_reason}]\" if skip_update_reason else \"\"\n    print(f\"  step summary: reward={mean_reward:.3f}  deal={deal_rate:.0%}  \"\n          f\"turns={mean_turns:.1f}  tom={mean_tom:.2f}  loss={loss_val:.3f}\"\n          f\"{tribunal_tag}  elapsed={elapsed/60:.1f}m{skip_tag}\")\n\n    if step % CONFIG[\"checkpoint_every\"] == 0:\n        save_checkpoint(step)\n        with open(CONFIG[\"history_path\"], \"w\") as f:\n            json.dump(training_history, f, indent=2, default=str)\n\nwriter.close()\nprint(f\"\\nTraining done in {(time.time() - t_start)/60:.1f} minutes.\")\nprint(f\"Steps skipped: no_variance={n_skipped_no_variance}, all_errored={n_skipped_all_errored}\")\n"
   },
   {
    "cell_type": "markdown",