Fix GRPO all-zero rewards: parse JSON actions from completions

Remove environment_factory from GRPOTrainer — reward_func now parses
JSON action objects from model text completions via regex and executes
them in fresh HCM21Env instances. This bypasses TRL's tool-calling
infrastructure which fails with a 0.6B model + Unsloth patches.

Changes:
- SYSTEM_PROMPT: instruct model to output JSON action lines
- reward_func: parse actions, run in env, return scores
- GRPOTrainer: remove environment_factory arg
- Evaluation: single-shot generation matching training format
- Fix BatchEncoding .shape bug in eval cell
- Add CLAUDE.md with GRPO training architecture docs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

CLAUDE.md

@@ -0,0 +1,177 @@
+# HR Productivity Environment — AgentX-AgentBeats
+
+## Overview
+A fully-implemented HR simulation environment (OpenEnv-compliant) where an LLM agent serves as Chief Human Capital Officer (CHCO), managing a simulated company over 6 quarters. The agent uses Fitz-enz HR Analytics frameworks to optimize human capital metrics through strategic hiring, training, compensation, and retention decisions.
+
+## Repository Structure
+```
+/tmp/autoresearch/
+├── hr_env/
+│   └── server/
+│       ├── environment.py    # Main simulation loop (OpenEnv Environment)
+│       ├── company.py        # Company + Department models (Cobb-Douglas production)
+│       ├── employee.py       # Employee dataclass with lifecycle methods
+│       ├── metrics.py        # Fitz-enz metrics: HCVA, HCROI, QIPS, Five Indexes
+│       ├── scoring.py        # Reward computation + final score (normalized [0,1])
+│       ├── phases.py         # HCM:21 phase management (Scanning→Planning→Producing→Controlling)
+│       ├── events.py         # Stochastic quarterly events (market, competitor, etc.)
+│       ├── data_gen.py       # Synthetic company generation (Faker + numpy)
+│       └── app.py            # OpenEnv server entrypoint
+├── agent/
+│   ├── hr_agent.py           # Claude-based CHCO agent with phase-aware prompting
+│   ├── prompts.py            # System prompt, phase prompts, summary prompts
+│   └── run.py                # Agent runner script
+├── tests/                    # Test suite
+├── openenv.yaml              # OpenEnv configuration
+├── pyproject.toml             # Python project config (uv)
+└── Dockerfile                 # Container build
+```
+
+## Development Commands
+```bash
+# Install dependencies
+uv sync
+
+# Run tests
+uv run pytest tests/ -v
+
+# Start OpenEnv server
+uv run python -m hr_env.server.app
+
+# Run agent against server
+uv run python -m agent.run
+
+# Docker build
+docker build -t hr-productivity-env .
+```
+
+## Architecture
+
+### Simulation Engine
+- **Company Model**: 5 departments (Engineering, Sales, Operations, HR, Finance) with Cobb-Douglas production functions
+- **Employee Model**: 17 tracked attributes including performance, engagement, flight_risk, promotability
+- **Revenue**: `A * headcount^alpha * avg_skill^beta * engagement_mult * market_modifier` per department
+- **Turnover**: Stochastic based on flight_risk (logistic function of pay, engagement, tenure, market)
+
+### Fitz-enz Metrics
+- **HCVA**: (Revenue - Non-Employment Costs) / FTE
+- **HCROI**: Revenue / Employment Cost
+- **QIPS**: Quality(0.25) + Innovation(0.25) + Productivity(0.30) + Service(0.20)
+- **Five Indexes**: Cost, Time, Quantity, Quality, Human Reactions (% change vs prior quarter)
+- **Employee Value**: avg(Productivity + Promotability + Transferability + Retainability)
+
+### HCM:21 Phase System (per quarter)
+1. **Scanning** (min 2 actions): query_department, query_employees, calculate_metric, review_financials
+2. **Planning** (min 1): set_hiring_target, set_training_budget, set_compensation_policy, set_retention_program
+3. **Producing** (min 1): execute_hiring, execute_promotion, execute_transfer, execute_training, execute_termination
+4. **Controlling** (min 1): submit_report, calculate_metric → then advance_quarter
+
+### Scoring
+- **Quarterly reward**: HCVA improvement (40%) + HCROI (20%) + QIPS (20%) + Five Indexes (20%)
+- **Final score** [0,1]: HCVA trajectory (25%) + HCROI improvement (20%) + Employee Value (20%) + QIPS consistency (15%) + Five Indexes cumulative (10%) + Financial health (10%)
+
+### Stochastic Events
+MarketDownturn, CompetitorPoaching, ProductLaunch, ExecDeparture, MinWageIncrease, EmployerAward, BudgetCut, TrainingBreakthrough — each with independent probability per quarter, max 2 per quarter.
+
+## Databricks Integration
+The simulation is adapted into a Databricks notebook (`/tmp/AimpointDigital/01_agentic_wikipedia_aimpoint_interview.ipynb`) with:
+- **Wikipedia knowledge layer**: HR theory articles (Fitz-enz, HCM, talent management) indexed into FAISS via BGE-Large embeddings
+- **LangGraph orchestration**: StateGraph with 5 nodes (initialize → agent_decide → execute_tools → check_done → summarize)
+- **RAG bridge**: `consult_hr_knowledge` tool lets the agent ground decisions in Wikipedia-sourced HR theory
+- **Databricks LLM**: `databricks-meta-llama-3-1-8b-instruct` for agent reasoning
+
+## Key Conventions
+- Pydantic models for OpenEnv; plain dataclasses/dicts for Databricks adaptation
+- Phase-gated action space prevents random action sequences
+- Company generation uses Faker (names) + numpy (distributions) with deterministic seeds
+- Salary distributions: lognormal; tenure: exponential; engagement: beta(7,3)
+- Level distribution: pyramid [0.35, 0.30, 0.20, 0.10, 0.05] for levels 1-5
+- Recruiting cost: $8,000/hire; Benefits: 30% of salary; Training: $50/hr/employee
+
+## Databricks Deployment
+
+### Workspace
+- **URL**: `https://dbc-c005fe3a-5584.cloud.databricks.com`
+- **Org ID**: `7474655330631867`
+- **Target Folder ID**: `1876640772048528`
+- **Folder URL**: `https://dbc-c005fe3a-5584.cloud.databricks.com/browse/folders/1876640772048528?o=7474655330631867`
+- **PAT**: (stored in env / .databrickscfg — not committed)
+- **User**: `n8.mauer@gmail.com`
+- **Edition**: Community Edition (Free)
+- **LLM endpoint**: `databricks-meta-llama-3-1-8b-instruct`
+- **Embedding endpoint**: `databricks-bge-large-en`
+
+### .env file (for local development)
+```
+DATABRICKS_HOST=https://dbc-c005fe3a-5584.cloud.databricks.com
+DATABRICKS_TOKEN=<your-token>
+DATABRICKS_ORG_ID=7474655330631867
+```
+
+### Cold-Start Setup
+1. `pip install databricks-cli`
+2. Write `~/.databrickscfg`:
+   ```ini
+   [DEFAULT]
+   host = https://dbc-c005fe3a-5584.cloud.databricks.com
+   token = <your-token>
+   ```
+3. Deploy notebook:
+   ```bash
+   databricks workspace import \
+     "/tmp/AimpointDigital/01_agentic_wikipedia_aimpoint_interview.ipynb" \
+     /Users/n8.mauer@gmail.com/01_agentic_wikipedia_aimpoint_interview \
+     --format JUPYTER --language PYTHON --overwrite
+   ```
+4. Open workspace URL → attach cluster → Run All
+
+### Deploy Script
+```bash
+./deploy.sh   # from /tmp/AimpointDigital/
+```
+
+### Runtime Fixes Applied
+- **Rate limiting**: `agent_decide` LLM calls wrapped with `backoff.on_exception` (exponential, max 8 retries, 120s) + 1.5s proactive delay between calls. Community Edition has low QPS for foundation model endpoints.
+- **MLflow tracing**: `mlflow.langchain.autolog()` added at top of Cell A2 to enable tracing and suppress Databricks "Enhance GenAI observability" hint. `mlflow-skinny[databricks]` already in pip install cell.
+
+### GitHub Repo
+- URL: `https://github.com/n8mauer/AimpointDigital`
+- Contains: notebook, build script, deploy script, README (runbook + talk track), assignment PDF
+
+## TRL/GRPO Training (`notebooks/hcm21_trl_training.ipynb`)
+
+### Architecture
+- **Model**: `unsloth/Qwen3-0.6B` with LoRA (r=16, 4-bit quantization)
+- **Method**: GRPO (Group Relative Policy Optimization) via TRL's `GRPOTrainer`
+- **Environment**: HCM:21 OpenEnv on HF Spaces (`https://paretooptimal-hcm21.hf.space`)
+- **Colab URL**: `https://colab.research.google.com/github/n8mauer/autoresearch/blob/main/notebooks/hcm21_trl_training.ipynb`
+
+### GRPO Reward Strategy (JSON action parsing — no environment_factory)
+The model generates text completions containing JSON action objects (one per line).
+`reward_func` parses these with regex (`\{[^{}]*"action_type"[^{}]*\}`), executes
+them in a fresh `HCM21Env` instance, and returns the environment's final score.
+
+**Why not `environment_factory`?**
+- TRL's `environment_factory` exposes env methods as tools for tool-calling, but a 0.6B model can't reliably generate structured tool calls
+- `environment_factory` requires `transformers>=5.2.0` which conflicts with Unsloth patches
+- Disabling Qwen3 thinking mode via `chat_template.replace("enable_thinking", "")` may break tool-calling templates
+- Direct JSON parsing in `reward_func` is simpler, more robust, and works with Unsloth
+
+### Key Implementation Details
+- `SYSTEM_PROMPT` instructs the model to output 20-40 JSON action objects covering 6 quarters
+- `reward_func(completions, **kwargs)` — no `environments` param (environments created internally)
+- `GRPOTrainer` has NO `environment_factory` argument
+- Evaluation (cell 20) generates one completion per seed, parses all actions, executes them
+- `tokenizer.apply_chat_template()` may return `BatchEncoding` — extract `.input_ids` before using `.shape`
+
+### Training Config
+- `per_device_train_batch_size=2`, `gradient_accumulation_steps=4`, `num_generations=4`
+- `max_completion_length=2048`, `learning_rate=5e-6`, `fp16=True`
+- Dataset: 20 prompts with different episode seeds
+
+### Verification Checklist
+After running in Colab, confirm:
+- `reward` column shows **non-zero values**
+- `reward_std` is **non-zero** (learning signal exists)
+- `Training Loss` is **non-zero** (GRPO has gradients)
+- Rewards trend upward over training steps

notebooks/hcm21_trl_training.ipynb

@@ -301,39 +301,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from datasets import Dataset\n",
-    "\n",
-    "SYSTEM_PROMPT = \"\"\"You are the Chief Human Capital Officer (CHCO) of a simulated company.\n",
-    "You manage 300 employees across 5 departments over 6 quarters.\n",
-    "Each quarter has 4 phases: Scanning, Planning, Producing, Controlling.\n",
-    "Maximize HCVA, HCROI, and QIPS metrics by making strategic HR decisions.\n",
-    "\n",
-    "When calling the step tool, provide a JSON action object like:\n",
-    "{\"action_type\": \"query_department\", \"department\": \"Engineering\"}\n",
-    "\n",
-    "Available action types vary by phase:\n",
-    "- Scanning: query_department, query_employees, calculate_metric, review_financials\n",
-    "- Planning: set_hiring_target, set_training_budget, set_compensation_policy, set_retention_program\n",
-    "- Producing: execute_hiring, execute_promotion, execute_transfer, execute_training, execute_termination\n",
-    "- Controlling: submit_report, calculate_metric, advance_quarter\n",
-    "- advance_phase: move to next phase (after minimum actions taken)\n",
-    "\n",
-    "Think step by step. Start by scanning each department.\"\"\"\n",
-    "\n",
-    "# Create prompts with different seeds for diversity\n",
-    "prompts = []\n",
-    "for seed in range(20):\n",
-    "    prompts.append({\n",
-    "        \"prompt\": [\n",
-    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "            {\"role\": \"user\", \"content\": f\"Begin managing the company. Episode seed: {seed}. Start by scanning your organization to understand its current state.\"},\n",
-    "        ]\n",
-    "    })\n",
-    "\n",
-    "dataset = Dataset.from_list(prompts)\n",
-    "print(f\"Training dataset: {len(dataset)} episodes\")"
-   ]
+   "source": "from datasets import Dataset\n\nSYSTEM_PROMPT = \"\"\"You are the Chief Human Capital Officer (CHCO) of a simulated company.\nYou manage 300 employees across 5 departments over 6 quarters.\nEach quarter has 4 phases: Scanning, Planning, Producing, Controlling.\nMaximize HCVA, HCROI, and QIPS metrics by making strategic HR decisions.\n\nOutput a sequence of JSON action objects, one per line. Example:\n{\"action_type\": \"query_department\", \"department\": \"Engineering\"}\n{\"action_type\": \"query_department\", \"department\": \"Sales\"}\n{\"action_type\": \"advance_phase\"}\n{\"action_type\": \"set_hiring_target\", \"department\": \"Engineering\", \"count\": 5}\n{\"action_type\": \"advance_phase\"}\n\nAvailable action types by phase:\n- Scanning: query_department, query_employees, calculate_metric, review_financials\n- Planning: set_hiring_target, set_training_budget, set_compensation_policy, set_retention_program\n- Producing: execute_hiring, execute_promotion, execute_transfer, execute_training, execute_termination\n- Controlling: submit_report, calculate_metric\n- advance_phase / advance_quarter: transition between phases/quarters\n\nOutput 20-40 actions covering all 6 quarters. Include advance_phase and advance_quarter actions to progress through the simulation.\"\"\"\n\n# Create prompts with different seeds for diversity\nprompts = []\nfor seed in range(20):\n    prompts.append({\n        \"prompt\": [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": f\"Begin managing the company. Episode seed: {seed}. Start by scanning your organization to understand its current state.\"},\n        ]\n    })\n\ndataset = Dataset.from_list(prompts)\nprint(f\"Training dataset: {len(dataset)} episodes\")"
   },
   {
    "cell_type": "markdown",
@@ -401,40 +369,17 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "def reward_func(completions, environments, **kwargs):\n",
-    "    \"\"\"Extract rewards from the HCM:21 environment after each rollout.\"\"\"\n",
-    "    rewards = []\n",
-    "    for env in environments:\n",
-    "        r = env.get_reward()\n",
-    "        rewards.append(r)\n",
-    "    return rewards\n",
-    "\n",
-    "print(\"Reward function defined.\")"
-   ]
+   "source": "import re\n\ndef reward_func(completions, **kwargs):\n    \"\"\"Parse actions from model completions, execute in environment, return rewards.\"\"\"\n    rewards = []\n    for completion in completions:\n        # Extract text content from completion\n        if isinstance(completion, list):\n            text = completion[-1].get(\"content\", \"\") if completion else \"\"\n        else:\n            text = str(completion)\n\n        # Parse all JSON objects that have action_type\n        actions = []\n        for match in re.finditer(r'\\{[^{}]*\"action_type\"[^{}]*\\}', text):\n            try:\n                action = json.loads(match.group())\n                actions.append(action)\n            except json.JSONDecodeError:\n                continue\n\n        if not actions:\n            rewards.append(0.0)\n            continue\n\n        # Run episode with parsed actions\n        try:\n            env = HCM21Env()\n            env.reset(seed=42)\n            for action in actions:\n                env.step(json.dumps(action))\n                if env._done:\n                    break\n            reward = env.get_reward()\n            env.close()\n            rewards.append(float(reward))\n        except Exception as e:\n            rewards.append(0.0)\n\n    return rewards\n\nprint(\"Reward function defined.\")"
   },
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "## 6. Configure and Run GRPO Training\n",
-    "\n",
-    "We use TRL's `GRPOTrainer` with `environment_factory=HCM21Env`.\n",
-    "The model generates tool calls to interact with the environment,\n",
-    "and GRPO optimizes the policy based on the environment rewards.\n",
-    "\n",
-    "On a free T4 GPU, we use:\n",
-    "- Small batch size (2) with gradient accumulation\n",
-    "- 4-bit quantization via Unsloth\n",
-    "- Short max_completion_length to fit in memory\n",
-    "\n",
-    "After training, we plot the reward trajectory to visualize learning progress."
-   ]
+   "source": "## 6. Configure and Run GRPO Training\n\nWe use TRL's `GRPOTrainer` with a custom `reward_func` that parses JSON\naction sequences from model completions and executes them in the HCM:21\nenvironment. GRPO generates G completions per prompt, scores them via\nthe reward function, and reinforces higher-scoring completions.\n\nOn a free T4 GPU, we use:\n- Small batch size (2) with gradient accumulation\n- 4-bit quantization via Unsloth\n- Short max_completion_length to fit in memory\n\nAfter training, we plot the reward trajectory to visualize learning progress."
   },
   {
    "cell_type": "code",
    "metadata": {},
-   "source": "from trl import GRPOTrainer, GRPOConfig\nimport os\n\n# Disable Qwen3 thinking mode at the tokenizer level (Unsloth-compatible)\ntokenizer.chat_template = tokenizer.chat_template.replace(\"enable_thinking\", \"\")\n\nos.environ[\"TENSORBOARD_LOGGING_DIR\"] = \"./hcm21-grpo-logs\"\n\ntraining_args = GRPOConfig(\n    output_dir=\"./hcm21-grpo-output\",\n    num_train_epochs=1,\n    per_device_train_batch_size=2,\n    gradient_accumulation_steps=4,\n    num_generations=4,\n    max_completion_length=2048,\n    learning_rate=5e-6,\n    logging_steps=1,\n    save_steps=10,\n    fp16=True,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    train_dataset=dataset,\n    reward_funcs=reward_func,\n    environment_factory=HCM21Env,\n    args=training_args,\n)\n\nprint(\"Trainer initialized. Starting training...\")\ntrainer.train()\nprint(\"Training complete!\")\n\n# Extract and plot reward curve from training logs\nimport matplotlib.pyplot as plt\n\nlog_history = trainer.state.log_history\nrewards = [entry[\"reward\"] for entry in log_history if \"reward\" in entry]\nlosses = [entry[\"loss\"] for entry in log_history if \"loss\" in entry]\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# Reward curve\nif rewards:\n    axes[0].plot(rewards, alpha=0.4, label=\"Per-step reward\", color=\"steelblue\")\n    if len(rewards) >= 3:\n        window = min(5, len(rewards))\n        rolling = np.convolve(rewards, np.ones(window)/window, mode=\"valid\")\n        axes[0].plot(range(window-1, len(rewards)), rolling, linewidth=2,\n                     label=f\"{window}-step rolling avg\", color=\"darkblue\")\n    axes[0].axhline(y=np.mean(random_scores), color=\"red\", linestyle=\"--\",\n                    label=f\"Random baseline ({np.mean(random_scores):.3f})\")\n    axes[0].axhline(y=np.mean(heuristic_scores), color=\"green\", linestyle=\"--\",\n                    label=f\"Heuristic baseline ({np.mean(heuristic_scores):.3f})\")\n    axes[0].set_xlabel(\"Training Step\")\n    axes[0].set_ylabel(\"Episode Reward\")\n    axes[0].set_title(\"GRPO Training Reward Curve\")\n    axes[0].legend()\n    axes[0].grid(True, alpha=0.3)\n\n# Loss curve\nif losses:\n    axes[1].plot(losses, color=\"orange\", alpha=0.7)\n    axes[1].set_xlabel(\"Training Step\")\n    axes[1].set_ylabel(\"Loss\")\n    axes[1].set_title(\"Training Loss\")\n    axes[1].grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.savefig(\"./hcm21-grpo-output/training_curves.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\nprint(f\"\\nTraining stats: {len(rewards)} reward entries, mean={np.mean(rewards):.4f}\" if rewards else \"No reward data logged.\")",
+   "source": "from trl import GRPOTrainer, GRPOConfig\nimport os\n\n# Disable Qwen3 thinking mode at the tokenizer level (Unsloth-compatible)\ntokenizer.chat_template = tokenizer.chat_template.replace(\"enable_thinking\", \"\")\n\nos.environ[\"TENSORBOARD_LOGGING_DIR\"] = \"./hcm21-grpo-logs\"\n\ntraining_args = GRPOConfig(\n    output_dir=\"./hcm21-grpo-output\",\n    num_train_epochs=1,\n    per_device_train_batch_size=2,\n    gradient_accumulation_steps=4,\n    num_generations=4,\n    max_completion_length=2048,\n    learning_rate=5e-6,\n    logging_steps=1,\n    save_steps=10,\n    fp16=True,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    train_dataset=dataset,\n    reward_funcs=reward_func,\n    args=training_args,\n)\n\nprint(\"Trainer initialized. Starting training...\")\ntrainer.train()\nprint(\"Training complete!\")\n\n# Extract and plot reward curve from training logs\nimport matplotlib.pyplot as plt\n\nlog_history = trainer.state.log_history\nrewards = [entry[\"reward\"] for entry in log_history if \"reward\" in entry]\nlosses = [entry[\"loss\"] for entry in log_history if \"loss\" in entry]\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# Reward curve\nif rewards:\n    axes[0].plot(rewards, alpha=0.4, label=\"Per-step reward\", color=\"steelblue\")\n    if len(rewards) >= 3:\n        window = min(5, len(rewards))\n        rolling = np.convolve(rewards, np.ones(window)/window, mode=\"valid\")\n        axes[0].plot(range(window-1, len(rewards)), rolling, linewidth=2,\n                     label=f\"{window}-step rolling avg\", color=\"darkblue\")\n    axes[0].axhline(y=np.mean(random_scores), color=\"red\", linestyle=\"--\",\n                    label=f\"Random baseline ({np.mean(random_scores):.3f})\")\n    axes[0].axhline(y=np.mean(heuristic_scores), color=\"green\", linestyle=\"--\",\n                    label=f\"Heuristic baseline ({np.mean(heuristic_scores):.3f})\")\n    axes[0].set_xlabel(\"Training Step\")\n    axes[0].set_ylabel(\"Episode Reward\")\n    axes[0].set_title(\"GRPO Training Reward Curve\")\n    axes[0].legend()\n    axes[0].grid(True, alpha=0.3)\n\n# Loss curve\nif losses:\n    axes[1].plot(losses, color=\"orange\", alpha=0.7)\n    axes[1].set_xlabel(\"Training Step\")\n    axes[1].set_ylabel(\"Loss\")\n    axes[1].set_title(\"Training Loss\")\n    axes[1].grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.savefig(\"./hcm21-grpo-output/training_curves.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\nprint(f\"\\nTraining stats: {len(rewards)} reward entries, mean={np.mean(rewards):.4f}\" if rewards else \"No reward data logged.\")",
    "outputs": [],
    "execution_count": null
   },
@@ -476,72 +421,7 @@
   {
    "cell_type": "code",
    "metadata": {},
-   "source": [
-    "import json\n",
-    "\n",
-    "FastLanguageModel.for_inference(model)\n",
-    "\n",
-    "def run_trained_episode(seed, max_steps=200):\n",
-    "    \"\"\"Run one episode with the trained model, return final score.\"\"\"\n",
-    "    env = HCM21Env()\n",
-    "    try:\n",
-    "        obs_text = env.reset(seed=seed)\n",
-    "        messages = [\n",
-    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "            {\"role\": \"user\", \"content\": f\"Begin. Observation:\\n{obs_text}\"},\n",
-    "        ]\n",
-    "        step_num = 0\n",
-    "        while not env._done and step_num < max_steps:\n",
-    "            inputs = tokenizer.apply_chat_template(\n",
-    "                messages, return_tensors=\"pt\", add_generation_prompt=True\n",
-    "            ).to(model.device)\n",
-    "            outputs = model.generate(\n",
-    "                input_ids=inputs, max_new_tokens=256,\n",
-    "                temperature=0.7, do_sample=True,\n",
-    "            )\n",
-    "            response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
-    "            try:\n",
-    "                start = response.index(\"{\")\n",
-    "                depth = 0\n",
-    "                for i, c in enumerate(response[start:], start):\n",
-    "                    if c == \"{\": depth += 1\n",
-    "                    elif c == \"}\": depth -= 1\n",
-    "                    if depth == 0:\n",
-    "                        action_json = response[start:i+1]\n",
-    "                        break\n",
-    "                obs_text = env.step(action_json)\n",
-    "            except (ValueError, json.JSONDecodeError):\n",
-    "                obs_text = env.step(json.dumps({\"action_type\": \"advance_phase\"}))\n",
-    "            messages.append({\"role\": \"assistant\", \"content\": response})\n",
-    "            messages.append({\"role\": \"user\", \"content\": obs_text})\n",
-    "            if len(messages) > 12:\n",
-    "                messages = messages[:2] + messages[-10:]\n",
-    "            step_num += 1\n",
-    "        return env.get_reward(), step_num\n",
-    "    finally:\n",
-    "        env.close()\n",
-    "\n",
-    "print(\"Evaluating trained agent on baseline seeds...\")\n",
-    "trained_scores = []\n",
-    "for seed in BASELINE_SEEDS:\n",
-    "    score, steps = run_trained_episode(seed)\n",
-    "    trained_scores.append(score)\n",
-    "    print(f\"  Seed {seed:>3d}: {score:.4f} ({steps} steps)\")\n",
-    "\n",
-    "sep = \"=\" * 50\n",
-    "print(f\"\\n{sep}\")\n",
-    "print(f\"         RESULTS COMPARISON\")\n",
-    "print(f\"{sep}\")\n",
-    "print(f\"  Random Agent:    {np.mean(random_scores):.4f} +/- {np.std(random_scores):.4f}\")\n",
-    "print(f\"  Heuristic Agent: {np.mean(heuristic_scores):.4f} +/- {np.std(heuristic_scores):.4f}\")\n",
-    "print(f\"  Trained Agent:   {np.mean(trained_scores):.4f} +/- {np.std(trained_scores):.4f}\")\n",
-    "print(f\"{sep}\")\n",
-    "\n",
-    "improvement = np.mean(trained_scores) - np.mean(random_scores)\n",
-    "print(f\"\\n  Improvement over random: +{improvement:.4f} ({improvement/max(np.mean(random_scores), 0.01)*100:.1f}%)\")\n",
-    "if np.mean(trained_scores) > np.mean(random_scores):\n",
-    "    print(\"  >>> Training shows positive improvement! <<<\")"
-   ],
+   "source": "import json\nimport re\n\nFastLanguageModel.for_inference(model)\n\ndef run_trained_episode(seed):\n    \"\"\"Run one episode with the trained model, return final score.\"\"\"\n    env = HCM21Env()\n    try:\n        env.reset(seed=seed)\n        messages = [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": f\"Begin managing the company. Episode seed: {seed}. Start by scanning your organization to understand its current state.\"},\n        ]\n\n        # Generate a single completion with all actions (matches training format)\n        input_ids = tokenizer.apply_chat_template(\n            messages, return_tensors=\"pt\", add_generation_prompt=True\n        )\n        # Handle both tensor and BatchEncoding returns\n        if hasattr(input_ids, \"input_ids\"):\n            input_ids = input_ids.input_ids\n        input_ids = input_ids.to(model.device)\n\n        outputs = model.generate(\n            input_ids=input_ids, max_new_tokens=2048,\n            temperature=0.7, do_sample=True,\n        )\n        response = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)\n\n        # Parse JSON actions using the same regex as reward_func\n        actions = []\n        for match in re.finditer(r'\\{[^{}]*\"action_type\"[^{}]*\\}', response):\n            try:\n                action = json.loads(match.group())\n                actions.append(action)\n            except json.JSONDecodeError:\n                continue\n\n        # Execute all parsed actions\n        for action in actions:\n            env.step(json.dumps(action))\n            if env._done:\n                break\n\n        return env.get_reward(), len(actions)\n    finally:\n        env.close()\n\nprint(\"Evaluating trained agent on baseline seeds...\")\ntrained_scores = []\nfor seed in BASELINE_SEEDS:\n    score, num_actions = run_trained_episode(seed)\n    trained_scores.append(score)\n    print(f\"  Seed {seed:>3d}: {score:.4f} ({num_actions} actions parsed)\")\n\nsep = \"=\" * 50\nprint(f\"\\n{sep}\")\nprint(f\"         RESULTS COMPARISON\")\nprint(f\"{sep}\")\nprint(f\"  Random Agent:    {np.mean(random_scores):.4f} +/- {np.std(random_scores):.4f}\")\nprint(f\"  Heuristic Agent: {np.mean(heuristic_scores):.4f} +/- {np.std(heuristic_scores):.4f}\")\nprint(f\"  Trained Agent:   {np.mean(trained_scores):.4f} +/- {np.std(trained_scores):.4f}\")\nprint(f\"{sep}\")\n\nimprovement = np.mean(trained_scores) - np.mean(random_scores)\nprint(f\"\\n  Improvement over random: +{improvement:.4f} ({improvement/max(np.mean(random_scores), 0.01)*100:.1f}%)\")\nif np.mean(trained_scores) > np.mean(random_scores):\n    print(\"  >>> Training shows positive improvement! <<<\")",
    "outputs": [],
    "execution_count": null
   },