Auto deploy from GitHub

space/Dockerfile

@@ -16,7 +16,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     NVIDIA_VISIBLE_DEVICES=all \
     NVIDIA_DRIVER_CAPABILITIES=compute,utility
 
-# Python + basic build deps (kept minimal)
+# Python + basic deps
 RUN apt-get update && apt-get install -y --no-install-recommends \
     python3.11 python3.11-venv python3-pip \
     git ca-certificates \
@@ -25,8 +25,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 COPY requirements.txt .
 RUN python3.11 -m pip install --upgrade pip && python3.11 -m pip install -r requirements.txt
 
+# CUDA-enabled PyTorch (so `torch.cuda.is_available()` is True on GPU Spaces)
+RUN python3.11 -m pip install --no-cache-dir \
+    torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+
 COPY . .
 
 EXPOSE 7860
-
 CMD ["python3.11", "app.py"]

space/space/space/Dockerfile

@@ -1,19 +1,32 @@
 # Optional: use Hugging Face Space SDK = docker (set `sdk: docker` in README.md).
-# Default README uses Gradio SDK and does not require this image.
+#
+# IMPORTANT:
+# - If your Space is on CPU hardware, `torch.cuda.is_available()` will be False no matter what.
+# - If your Space is on GPU hardware, you must use a CUDA-enabled base image (below),
+#   otherwise CUDA won't be visible inside the container.
 
-FROM python:3.11-slim
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
 
 WORKDIR /app
 
-ENV PYTHONUNBUFFERED=1 \
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
     PIP_NO_CACHE_DIR=1 \
-    GRADIO_SERVER_NAME=0.0.0.0
+    GRADIO_SERVER_NAME=0.0.0.0 \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+# Python + basic build deps (kept minimal)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3.11 python3.11-venv python3-pip \
+    git ca-certificates \
+  && rm -rf /var/lib/apt/lists/*
 
 COPY requirements.txt .
-RUN pip install --upgrade pip && pip install -r requirements.txt
+RUN python3.11 -m pip install --upgrade pip && python3.11 -m pip install -r requirements.txt
 
 COPY . .
 
 EXPOSE 7860
 
-CMD ["python", "app.py"]
+CMD ["python3.11", "app.py"]

space/space/space/space/space/space/program_generator.py

@@ -0,0 +1,116 @@
+"""
+Random Toy-IR list-of-dicts for GRPO / Deliverable2 / runtime_core.
+
+Schema matches `runtime_core.SAMPLE_PROGRAM` (op, args, dest, type) — the same
+shape `Deliverable2_Formatter` and `CompilerOptimizationEnv` expect.
+
+Note: `metahack1 (1).ipynb` uses a different TAC shape (CONST/src1/STORE).
+Use that notebook's generators only if you add a separate converter; this
+module is self-contained for training stack compatibility.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import random
+from typing import List
+
+# Hand-written seeds (same as Colab / train.py short list) — kept in sync for regression.
+_BUILTIN_EXTRA: List[List[dict]] = [
+    [
+        {"op": "const", "dest": "a", "args": ["1"], "type": "int"},
+        {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"},
+        {"op": "ret", "args": ["b"]},
+    ],
+    [
+        {"op": "const", "dest": "x", "args": ["2"], "type": "int"},
+        {"op": "const", "dest": "y", "args": ["4"], "type": "int"},
+        {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"},
+        {"op": "const", "dest": "k", "args": ["1"], "type": "int"},
+        {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"},
+        {"op": "ret", "args": ["w"]},
+    ],
+]
+
+
+def random_toy_ir_program(rng: random.Random) -> List[dict]:
+    """
+    One valid program: consts v0.., then a chain of add/mul on existing names, then ret.
+    All ops use the mock-engine-friendly list schema.
+    """
+    n_const = rng.randint(2, 5)
+    n_arith = rng.randint(1, 5)
+    progs: List[dict] = []
+    for i in range(n_const):
+        progs.append(
+            {
+                "op": "const",
+                "dest": f"v{i}",
+                "args": [str(rng.randint(0, 20))],
+                "type": "int",
+            }
+        )
+    available = [f"v{i}" for i in range(n_const)]
+    nxt = n_const
+    for _j in range(n_arith):
+        a = rng.choice(available)
+        b = rng.choice(available)
+        opn = rng.choice(["add", "mul"])
+        d = f"v{nxt}"
+        nxt += 1
+        progs.append({"op": opn, "dest": d, "args": [a, b], "type": "int"})
+        available.append(d)
+    progs.append({"op": "ret", "args": [available[-1]]})
+    return progs
+
+
+def build_training_program_corpus(
+    n_total: int = 120,
+    seed: int = 42,
+    *,
+    include_builtins: bool = True,
+) -> List[List[dict]]:
+    """
+    Return `n_total` programs for GRPO. Optionally prepend SAMPLE_PROGRAM + 2 hand-written IRs
+    (when include_builtins), then fill with random_toy_ir_program, deduplicating by JSON key.
+
+    Typical range: set `n_total` between 50 and 200 in the notebook.
+    """
+    if n_total < 1:
+        raise ValueError("n_total must be >= 1")
+
+    rng = random.Random(seed)
+    out: List[List[dict]] = []
+    seen: set[str] = set()
+
+    def _add(p: List[dict]) -> None:
+        k = json.dumps(p, sort_keys=True)
+        if k in seen:
+            return
+        seen.add(k)
+        out.append(copy.deepcopy(p))
+
+    if include_builtins:
+        from runtime_core import SAMPLE_PROGRAM
+
+        for p in (SAMPLE_PROGRAM, *_BUILTIN_EXTRA):
+            if len(out) >= n_total:
+                break
+            _add(p)
+
+    # Fill with random programs (dedupe by full JSON; allow dup if generator keeps colliding)
+    guard = 0
+    while len(out) < n_total:
+        guard += 1
+        if guard > 200_000:
+            out.append(random_toy_ir_program(rng))
+            continue
+        cand = random_toy_ir_program(rng)
+        k = json.dumps(cand, sort_keys=True)
+        if k in seen:
+            continue
+        seen.add(k)
+        out.append(cand)
+
+    return out

space/space/space/space/space/space/space/reverse_pass/README.md

@@ -0,0 +1,57 @@
+# Compiler Phase-Ordering RLVR (Toy-IR)
+
+Toy-IR compiler phase-ordering environment for RL with verifiable rewards (RLVR), including reverse passes to escape local optima.
+
+## What Is Implemented
+
+- Reverse passes are available in the pass library:
+  - `expand_constant`
+  - `duplicate_computation`
+- Action routing and validation include reverse passes.
+- Prompt guidance includes reverse-pass descriptions and usage intent.
+- Reward shaping is terminal-weighted with RLVR hard gate:
+  - non-equivalent -> `-1000.0`
+  - terminal -> `((original_cycles - current_cycles) / original_cycles) * 100.0`
+  - non-terminal -> `-0.1`
+- Episode termination:
+  - `STOP`/`done`, or
+  - hard cap at 5 steps.
+- Reverse-pass instrumentation logs every 50 episodes/completions:
+  - `reverse_pass_episodes`
+  - `expand_constant_count`
+  - `duplicate_computation_count`
+
+## Verified So Far
+
+- Smoke test runs end-to-end without runtime errors.
+- Rewards are scalar floats (not NaN) in tested rollouts.
+- Forced reverse-pass episodes show expected behavior:
+  - small negative intermediate rewards (`-0.1`)
+  - terminal reward depends on final outcome (positive if chain beats baseline, negative if not)
+
+## Pending (GPU Required)
+
+Comparative training runs are still pending and require a GPU environment with:
+
+- `torch`
+- `trl`
+- `unsloth`
+- `wandb`
+
+Required fair comparison:
+
+1. Baseline run: reverse passes disabled (`baseline_no_reverse`)
+2. Reverse run: reverse passes enabled (`with_reverse_passes`)
+3. Same episode budget for both runs
+4. Compare reward/cycle curves and reverse-pass usage metrics in WandB
+
+## Team Handoff Status
+
+- Reverse-pass feature: shipped
+- Training wiring + reward shaping + instrumentation: shipped
+- Comparative training: pending in GPU environment
+- README: this document
+
+## Suggested Final Pre-Training Sanity Check
+
+Run one smoke/equivalence check on a real Role 3 generated curriculum program (not just mock IR) to confirm verifier behavior in-pipeline before long training.

space/space/space/space/space/space/space/reverse_pass/compiler_optimization_grpo.ipynb

@@ -0,0 +1,850 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "title-cell",
+   "metadata": {},
+   "source": [
+    "# 🔧 Compiler Optimization RL Environment\n",
+    "### OpenEnv Hackathon 2026 — Theme #2: Long-Horizon Planning\n",
+    "\n",
+    "**What this notebook does:**\n",
+    "1. Defines a fully OpenEnv-compliant `CompilerOptimizationEnv`\n",
+    "2. Loads `Qwen2.5-3B-Instruct` via Unsloth (4-bit QLoRA)\n",
+    "3. Trains with GRPO (TRL) — LLM learns to pick compiler passes that reduce CPU cycles\n",
+    "4. Runs a smoke-test with a mock engine so you can verify reward logic without real hardware\n",
+    "5. Plots reward curves\n",
+    "\n",
+    "---\n",
+    "**Stack:** `unsloth` · `trl` · `openenv` · `wandb` · `matplotlib`\n",
+    "\n",
+    "> **Runtime:** Google Colab T4 GPU recommended. For the smoke-test only, CPU is fine."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "install-header",
+   "metadata": {},
+   "source": [
+    "## 📦 Cell 1 — Install Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "install-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install all required packages\n",
+    "# Unsloth must be installed before trl to get the right CUDA kernels\n",
+    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n",
+    "!pip install trl datasets transformers accelerate peft bitsandbytes --quiet\n",
+    "!pip install wandb matplotlib --quiet\n",
+    "\n",
+    "# Optional: install openenv for production use\n",
+    "# !pip install openenv --quiet\n",
+    "\n",
+    "print(\"✅ All packages installed\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "env-header",
+   "metadata": {},
+   "source": [
+    "## 🌍 Cell 2 — CompilerOptimizationEnv (OpenEnv-Compliant)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "env-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# CompilerOptimizationEnv, PASSES, reverse passes, terminal-weighted reward (toyir_rl_support)\n",
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "if os.getcwd() not in sys.path:\n",
+    "    sys.path.insert(0, os.getcwd())\n",
+    "\n",
+    "from toyir_rl_support import (\n",
+    "    MCPEnvironment,\n",
+    "    StepResult,\n",
+    "    EpisodeStats,\n",
+    "    CompilerOptimizationEnv,\n",
+    "    PASSES,\n",
+    "    MOCK_PASSES,\n",
+    "    compute_shaped_reward,\n",
+    "    rollout_shaped_return,\n",
+    "    log_reverse_pass_stats_for_completion,\n",
+    "    ensure_import_path,\n",
+    "    MockEngine,\n",
+    "    MOCK_ENGINE,\n",
+    "    SAMPLE_PROGRAM,\n",
+    ")\n",
+    "\n",
+    "ensure_import_path()\n",
+    "print(\"✅ Loaded CompilerOptimizationEnv, PASSES, and reward helpers from toyir_rl_support\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "smoke-header",
+   "metadata": {},
+   "source": [
+    "## 🧪 Cell 3 — Smoke Test (No GPU / Real Engine Needed)\n",
+    "Validates the entire reward pipeline with a mock engine. Run this before spending compute."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "smoke-test-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Smoke test (MockEngine; same PASSES as training)\n",
+    "engine = MockEngine()\n",
+    "env = CompilerOptimizationEnv(\n",
+    "    engine, MOCK_PASSES, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n",
+    ")\n",
+    "obs = env.reset(SAMPLE_PROGRAM)\n",
+    "\n",
+    "print(\"═\" * 55)\n",
+    "print(\"SMOKE TEST\")\n",
+    "print(\"═\" * 55)\n",
+    "print(f\"Initial state:\\n{obs}\")\n",
+    "print(f\"\\nBaseline cycles : {env.previous_cycles}\")\n",
+    "print(f\"Available actions: {env.available_actions()}\")\n",
+    "print()\n",
+    "\n",
+    "actions_to_try = [\n",
+    "    \"constant_folding\",\n",
+    "    \"dead_code_elimination\",\n",
+    "    \"peephole_optimization\",  # no-op\n",
+    "    \"hallucinated_pass\",    # invalid — but won't kill episode yet\n",
+    "    \"constant_folding\",\n",
+    "    \"dead_code_elimination\",\n",
+    "    \"expand_constant\",\n",
+    "    \"STOP\",\n",
+    "]\n",
+    "\n",
+    "for action in actions_to_try:\n",
+    "    result = env.step(action)\n",
+    "    tag = \"✗\" if result.reward < 0 else \"✓\"\n",
+    "    print(f\"{tag} '{action}'\")\n",
+    "    print(f\"   reward={result.reward:+.2f}  done={result.done}  is_terminal={result.info.get('is_terminal')}\")\n",
+    "    relevant = {k: v for k, v in result.info.items()\n",
+    "                if k in (\"delta_pct\", \"error\", \"no_op\", \"reason\",\n",
+    "                         \"terminal_bonus\", \"episode_stats\", \"is_terminal\")}\n",
+    "    if relevant:\n",
+    "        print(f\"   info: {relevant}\")\n",
+    "    print()\n",
+    "    if result.done:\n",
+    "        break\n",
+    "\n",
+    "print(\"✅ Smoke test passed\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "reward-plot-header",
+   "metadata": {},
+   "source": [
+    "## 📊 Cell 4 — Visualise Reward Across a Mock Episode"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "reward-plot-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.ticker as ticker\n",
+    "\n",
+    "# Run a full episode and collect data\n",
+    "env   = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n",
+    "env.reset(SAMPLE_PROGRAM)\n",
+    "\n",
+    "sequence = [\n",
+    "    \"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\",\n",
+    "    \"expand_constant\", \"peephole_optimization\",\n",
+    "]\n",
+    "\n",
+    "rewards, cycle_counts, actions_log = [], [], []\n",
+    "for act in sequence:\n",
+    "    r = env.step(act)\n",
+    "    rewards.append(r.reward)\n",
+    "    cycle_counts.append(env.previous_cycles)\n",
+    "    actions_log.append(act)\n",
+    "    if r.done:\n",
+    "        break\n",
+    "\n",
+    "steps = list(range(1, len(rewards) + 1))\n",
+    "\n",
+    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n",
+    "fig.suptitle(\"Compiler Optimization Episode — Mock Engine\", fontsize=14, fontweight=\"bold\")\n",
+    "\n",
+    "# Reward per step\n",
+    "colors = [\"#2ecc71\" if r >= 0 else \"#e74c3c\" for r in rewards]\n",
+    "ax1.bar(steps, rewards, color=colors, edgecolor=\"white\", linewidth=0.5)\n",
+    "ax1.axhline(0, color=\"grey\", linewidth=0.8, linestyle=\"--\")\n",
+    "ax1.set_ylabel(\"Reward\")\n",
+    "ax1.set_title(\"Reward per Step (green = positive, red = negative)\")\n",
+    "ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter(\"%.1f\"))\n",
+    "\n",
+    "# Cycle count over time\n",
+    "ax2.plot(steps, cycle_counts, marker=\"o\", color=\"#3498db\", linewidth=2, markersize=6)\n",
+    "ax2.set_xlabel(\"Step\")\n",
+    "ax2.set_ylabel(\"CPU Cycles\")\n",
+    "ax2.set_title(\"CPU Cycle Count Over Episode (lower = better)\")\n",
+    "ax2.set_xticks(steps)\n",
+    "ax2.set_xticklabels(\n",
+    "    [a.replace(\"_\", \"\\n\") for a in actions_log],\n",
+    "    fontsize=7,\n",
+    ")\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"episode_reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(\"📈 Plot saved as episode_reward_curve.png\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "model-header",
+   "metadata": {},
+   "source": [
+    "## 🤖 Cell 5 — Load Model with Unsloth (QLoRA 4-bit)\n",
+    "> **Requires T4 GPU.** Skip to Cell 9 if you only want to test the environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "model-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "MODEL_NAME  = \"unsloth/Qwen2.5-3B-Instruct\"  # swap to 7B if VRAM allows\n",
+    "MAX_SEQ_LEN = 1024\n",
+    "LORA_RANK   = 16\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name     = MODEL_NAME,\n",
+    "    max_seq_length = MAX_SEQ_LEN,\n",
+    "    dtype          = None,       # auto bf16/fp16\n",
+    "    load_in_4bit   = True,\n",
+    ")\n",
+    "\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r              = LORA_RANK,\n",
+    "    target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
+    "                      \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    lora_alpha     = LORA_RANK * 2,\n",
+    "    lora_dropout   = 0.0,\n",
+    "    bias           = \"none\",\n",
+    "    use_gradient_checkpointing = \"unsloth\",\n",
+    "    random_state   = 42,\n",
+    ")\n",
+    "\n",
+    "print(f\"✅ Loaded {MODEL_NAME} with QLoRA rank={LORA_RANK}\")\n",
+    "print(f\"   GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "prompt-header",
+   "metadata": {},
+   "source": [
+    "## 💬 Cell 6 — System Prompt & Dataset Builder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "prompt-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import re\n",
+    "from datasets import Dataset\n",
+    "\n",
+    "\n",
+    "class Deliverable2_Formatter:\n",
+    "    'State translation + robust action-array extraction for Role 2.'\n",
+    "\n",
+    "    @staticmethod\n",
+    "    def translate_state(raw_json: list) -> str:\n",
+    "        pseudo_assembly = []\n",
+    "        for i, instruction in enumerate(raw_json):\n",
+    "            if not isinstance(instruction, dict):\n",
+    "                pseudo_assembly.append(f\"{i}. NOP\")\n",
+    "                continue\n",
+    "            op = str(instruction.get(\"op\", \"UNKNOWN\")).upper()\n",
+    "            args = \", \".join(str(arg) for arg in instruction.get(\"args\", []))\n",
+    "            dest = instruction.get(\"dest\", \"\")\n",
+    "            if dest:\n",
+    "                pseudo_assembly.append(f\"{i}. {dest} = {op} {args}\".rstrip())\n",
+    "            else:\n",
+    "                pseudo_assembly.append(f\"{i}. {op} {args}\".rstrip())\n",
+    "        return \"\\n\".join(pseudo_assembly) if pseudo_assembly else \"; (empty program)\"\n",
+    "\n",
+    "    @staticmethod\n",
+    "    def extract_action_array(llm_output: str) -> list:\n",
+    "        text = (llm_output or \"\").strip()\n",
+    "        if not text:\n",
+    "            raise ValueError(\"Invalid JSON format\")\n",
+    "\n",
+    "        try:\n",
+    "            parsed = json.loads(text)\n",
+    "            if isinstance(parsed, list):\n",
+    "                return parsed\n",
+    "        except json.JSONDecodeError:\n",
+    "            pass\n",
+    "\n",
+    "        match = re.search(r\"\\[.*?\\]\", text, re.DOTALL)\n",
+    "        if match:\n",
+    "            try:\n",
+    "                parsed = json.loads(match.group(0))\n",
+    "                if isinstance(parsed, list):\n",
+    "                    return parsed\n",
+    "            except json.JSONDecodeError:\n",
+    "                pass\n",
+    "\n",
+    "        raise ValueError(\"Invalid JSON format\")\n",
+    "\n",
+    "\n",
+    "def build_system_prompt(passes: dict) -> str:\n",
+    "    lines: list = []\n",
+    "    for name in sorted(passes.keys()):\n",
+    "        if name == \"expand_constant\":\n",
+    "            lines.append(\n",
+    "                \"  - expand_constant: Splits a constant assignment into runtime arithmetic (e.g., CONST 8 becomes CONST 3 + ADD 5). May temporarily increase cycles but can enable forward passes to find better optimization chains. Use sparingly when standard passes seem stuck.\"\n",
+    "            )\n",
+    "        elif name == \"duplicate_computation\":\n",
+    "            lines.append(\n",
+    "                \"  - duplicate_computation: Creates a redundant copy of a binary operation with a fresh variable. May temporarily increase cycles but can enable alternative dead code elimination paths. Use when redundancy might unlock further simplification.\"\n",
+    "            )\n",
+    "        else:\n",
+    "            lines.append(f\"  - {name}\")\n",
+    "    action_block = \"\\n\".join(lines)\n",
+    "    return (\n",
+    "        \"You are a compiler optimization agent. Your goal is to reduce \"\n",
+    "        \"CPU cycle count by applying optimization passes to the program below.\\n\\n\"\n",
+    "        f\"Available actions:\\n{action_block}\\n\"\n",
+    "        \"  - STOP  (emit this in the array to stop; optional alias: \\\"done\\\". Stop early if no further improvement is possible.)\\n\\n\"\n",
+    "        \"Rules:\\n\"\n",
+    "        \"  • Output only a JSON array of pass names (example: [\\\"constant_folding\\\"]).\\n\"\n",
+    "        \"  • No explanation, no markdown, no extra text.\\n\"\n",
+    "        \"  • Do not invent actions not listed above.\\n\"\n",
+    "        \"  • Maximum 5 passes per response (hard episode cap: 5 optimization steps or STOP).\\n\"\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def build_dataset(\n",
+    "    programs: list,\n",
+    "    engine,\n",
+    "    passes: dict,\n",
+    ") -> Dataset:\n",
+    "    # Each row = one episode; GRPO samples K completions per row\n",
+    "    env = CompilerOptimizationEnv(\n",
+    "        engine, passes, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n",
+    "    )\n",
+    "    system_prompt = build_system_prompt(passes)\n",
+    "\n",
+    "    rows = []\n",
+    "    for prog in programs:\n",
+    "        translated_state = Deliverable2_Formatter.translate_state(prog)\n",
+    "        prompt = [\n",
+    "            {\"role\": \"system\", \"content\": system_prompt},\n",
+    "            {\"role\": \"user\", \"content\": f\"Current program:\\n{translated_state}\\n\\nChoose optimization passes:\"},\n",
+    "        ]\n",
+    "        rows.append({\"prompt\": prompt, \"program_json\": json.dumps(prog)})\n",
+    "\n",
+    "    return Dataset.from_list(rows)\n",
+    "\n",
+    "\n",
+    "# --- Demo: build dataset from mock programs -----------------------------------\n",
+    "def make_mock_programs(n: int = 20) -> list:\n",
+    "    # Generate N random mock IR programs\n",
+    "    import random\n",
+    "    ops = [\"add\", \"mul\", \"sub\", \"const\", \"load\"]\n",
+    "    progs = []\n",
+    "    for _ in range(n):\n",
+    "        length = random.randint(4, 12)\n",
+    "        prog = [\n",
+    "            {\n",
+    "                \"op\": random.choice(ops),\n",
+    "                \"dest\": f\"v{i}\",\n",
+    "                \"args\": [f\"v{max(0, i-1)}\"],\n",
+    "                \"type\": \"int\",\n",
+    "            }\n",
+    "            for i in range(length)\n",
+    "        ]\n",
+    "        prog.append({\"op\": \"ret\", \"args\": [f\"v{length-1}\"]})\n",
+    "        progs.append(prog)\n",
+    "    return progs\n",
+    "\n",
+    "\n",
+    "mock_programs = make_mock_programs(n=30)\n",
+    "train_dataset = build_dataset(mock_programs, MockEngine(), MOCK_PASSES)\n",
+    "\n",
+    "print(f\"✅ Dataset built: {len(train_dataset)} episodes\")\n",
+    "print(f\"   Sample prompt keys: {list(train_dataset[0].keys())}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "reward-fn-header",
+   "metadata": {},
+   "source": [
+    "## 🎯 Cell 7 — Reward Function Factory (for GRPO)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "reward-fn-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from typing import Any\n",
+    "\n",
+    "# Rollout + terminal weighting; reverse-pass W&B window logging\n",
+    "from toyir_rl_support import (\n",
+    "    CompilerOptimizationEnv,\n",
+    "    MOCK_PASSES,\n",
+    "    rollout_shaped_return,\n",
+    "    log_reverse_pass_stats_for_completion,\n",
+    "    SAMPLE_PROGRAM,\n",
+    ")\n",
+    "\n",
+    "EP_CAP = CompilerOptimizationEnv.MAX_EPISODE_STEPS\n",
+    "\n",
+    "\n",
+    "def _normalize_action_seq(actions: list, cap: int, passes: dict) -> list:\n",
+    "    out: list = []\n",
+    "    for x in actions[:cap]:\n",
+    "        raw = str(x).strip()\n",
+    "        s_low = raw.lower()\n",
+    "        if s_low in (\"stop\", \"done\"):\n",
+    "            out.append(\"done\" if s_low == \"done\" else \"STOP\")\n",
+    "            continue\n",
+    "        key = None\n",
+    "        for k in passes:\n",
+    "            if k.lower() == s_low:\n",
+    "                key = k\n",
+    "                break\n",
+    "        if key is None:\n",
+    "            out.append(raw)\n",
+    "        else:\n",
+    "            out.append(key)\n",
+    "    return out\n",
+    "\n",
+    "\n",
+    "def make_reward_fn(engine, passes, max_steps: int = EP_CAP):\n",
+    "    # TRL GRPO reward: terminal-weighted `rollout_shaped_return` + reverse-pass logging\n",
+    "    _cap = min(int(max_steps), EP_CAP)\n",
+    "\n",
+    "    def reward_fn(prompts, completions, **kwargs):\n",
+    "        programs = kwargs.get(\"program_json\", [None] * len(completions))\n",
+    "        rewards: list = []\n",
+    "\n",
+    "        for completion, prog_json in zip(completions, programs):\n",
+    "            raw = completion if isinstance(completion, str) else completion[0][\"content\"]\n",
+    "            if prog_json is None:\n",
+    "                rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n",
+    "                continue\n",
+    "            program = (\n",
+    "                json.loads(prog_json) if isinstance(prog_json, str) else list(prog_json)\n",
+    "            )\n",
+    "            if not program:\n",
+    "                rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n",
+    "                continue\n",
+    "            try:\n",
+    "                actions = Deliverable2_Formatter.extract_action_array(raw)\n",
+    "            except ValueError:\n",
+    "                rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n",
+    "                continue\n",
+    "            if not actions:\n",
+    "                rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n",
+    "                continue\n",
+    "            norm = _normalize_action_seq(actions, _cap, passes)\n",
+    "            tr, n_e, n_d, _hs, _hc, ok = rollout_shaped_return(\n",
+    "                program, norm, engine, passes\n",
+    "            )\n",
+    "            if not ok:\n",
+    "                rewards.append(-1000.0)\n",
+    "            else:\n",
+    "                log_reverse_pass_stats_for_completion(n_e, n_d)\n",
+    "                rewards.append(tr)\n",
+    "        return rewards\n",
+    "\n",
+    "    return reward_fn\n",
+    "\n",
+    "\n",
+    "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=EP_CAP)\n",
+    "print(\"✅ Reward function factory ready\")\n",
+    "\n",
+    "# Quick sanity check\n",
+    "test_completions = [\n",
+    "    '[\"constant_folding\"]',\n",
+    "    'Here is my plan: [\"dead_code_elimination\", \"peephole_optimization\", \"done\"]',\n",
+    "    \"hallucinated_pass\",\n",
+    "]\n",
+    "test_programs = [json.dumps(SAMPLE_PROGRAM)] * 3\n",
+    "test_rewards = reward_fn(\n",
+    "    prompts=[\"\"] * 3,\n",
+    "    completions=test_completions,\n",
+    "    program_json=test_programs,\n",
+    ")\n",
+    "print(\"\\nReward sanity check:\")\n",
+    "for act, rew in zip(test_completions, test_rewards):\n",
+    "    print(f\"  {act!r} → {rew:+.2f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "trainer-header",
+   "metadata": {},
+   "source": [
+    "## 🚀 Cell 8 — GRPO Trainer Config & Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "trainer-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "\n",
+    "# Optional W&B — comment out if not using\n",
+    "try:\n",
+    "    import wandb\n",
+    "    wandb.init(project=\"openenv-compiler-opt\", name=\"grpo-qwen2.5-3b\")\n",
+    "    REPORT_TO = \"wandb\"\n",
+    "except Exception:\n",
+    "    REPORT_TO = \"none\"\n",
+    "\n",
+    "\n",
+    "grpo_config = GRPOConfig(\n",
+    "    # ── Generation ────────────────────────────────────────────────────────\n",
+    "    num_generations             = 4,     # K rollouts per prompt for group-relative advantage\n",
+    "    max_new_tokens              = 16,    # Actions are 1 word; don't waste context\n",
+    "    temperature                 = 0.9,\n",
+    "    top_p                       = 0.95,\n",
+    "\n",
+    "    # ── Optimisation ──────────────────────────────────────────────────────\n",
+    "    learning_rate               = 5e-6,\n",
+    "    per_device_train_batch_size = 2,\n",
+    "    gradient_accumulation_steps = 4,    # effective batch = 8\n",
+    "    num_train_epochs            = 3,\n",
+    "    max_grad_norm               = 0.5,\n",
+    "\n",
+    "    # ── GRPO-specific ─────────────────────────────────────────────────────\n",
+    "    beta                        = 0.04,  # KL penalty; keeps policy near reference\n",
+    "\n",
+    "    # ── Logging / checkpointing ───────────────────────────────────────────\n",
+    "    output_dir                  = \"./grpo_compiler_checkpoints\",\n",
+    "    logging_steps               = 10,\n",
+    "    save_steps                  = 100,\n",
+    "    report_to                   = REPORT_TO,\n",
+    "\n",
+    "    # ── Reproducibility ───────────────────────────────────────────────────\n",
+    "    seed                        = 42,\n",
+    ")\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model         = model,\n",
+    "    tokenizer     = tokenizer,\n",
+    "    config        = grpo_config,\n",
+    "    train_dataset = train_dataset,\n",
+    "    reward_funcs  = reward_fn,\n",
+    ")\n",
+    "\n",
+    "print(\"✅ Trainer configured\")\n",
+    "print(f\"   num_generations (K)  = {grpo_config.num_generations}\")\n",
+    "print(f\"   effective batch size  = \"\n",
+    "      f\"{grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n",
+    "print(f\"   KL beta              = {grpo_config.beta}\")\n",
+    "print()\n",
+    "print(\"Starting training... (this will take a while on T4)\")\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "save-header",
+   "metadata": {},
+   "source": [
+    "## 💾 Cell 9 — Save Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "save-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAVE_PATH = \"./grpo_compiler_final\"\n",
+    "\n",
+    "model.save_pretrained(SAVE_PATH)\n",
+    "tokenizer.save_pretrained(SAVE_PATH)\n",
+    "\n",
+    "print(f\"✅ Model saved to {SAVE_PATH}\")\n",
+    "\n",
+    "# Optional: push to HuggingFace Hub\n",
+    "# model.push_to_hub(\"your-hf-username/compiler-opt-grpo\")\n",
+    "# tokenizer.push_to_hub(\"your-hf-username/compiler-opt-grpo\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "curriculum-header",
+   "metadata": {},
+   "source": [
+    "## 📈 Cell 10 — Curriculum Callback & Reward Tracking"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "curriculum-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CurriculumCallback:\n",
+    "    \"\"\"\n",
+    "    Tracks rolling mean reward and promotes curriculum level\n",
+    "    when the agent has mastered the current difficulty.\n",
+    "\n",
+    "    Usage: call .record(reward) after every episode.\n",
+    "    Read .level to get current difficulty (1=easy, 2=medium, 3=hard).\n",
+    "    \"\"\"\n",
+    "    def __init__(self, reward_threshold: float = 5.0, window: int = 50):\n",
+    "        self.threshold = reward_threshold\n",
+    "        self.window    = window\n",
+    "        self._history  = []\n",
+    "        self.level     = 1\n",
+    "        self._promotions = []\n",
+    "\n",
+    "    def record(self, reward: float, step: int = None):\n",
+    "        self._history.append(reward)\n",
+    "        if len(self._history) >= self.window:\n",
+    "            mean = sum(self._history[-self.window:]) / self.window\n",
+    "            if mean >= self.threshold and self.level < 3:\n",
+    "                self.level += 1\n",
+    "                self._promotions.append((step or len(self._history), self.level))\n",
+    "                print(f\"[Curriculum] ▲ Promoted to level {self.level} \"\n",
+    "                      f\"(rolling mean={mean:.2f})\")\n",
+    "\n",
+    "    def plot(self):\n",
+    "        import matplotlib.pyplot as plt\n",
+    "        import numpy as np\n",
+    "\n",
+    "        history = self._history\n",
+    "        steps   = list(range(len(history)))\n",
+    "        window  = self.window\n",
+    "        rolling = [\n",
+    "            sum(history[max(0,i-window):i+1]) / min(i+1, window)\n",
+    "            for i in steps\n",
+    "        ]\n",
+    "\n",
+    "        fig, ax = plt.subplots(figsize=(10, 4))\n",
+    "        ax.plot(steps, history, alpha=0.3, color=\"#3498db\", label=\"Episode reward\")\n",
+    "        ax.plot(steps, rolling, color=\"#e74c3c\", linewidth=2,\n",
+    "                label=f\"Rolling mean (w={window})\")\n",
+    "        ax.axhline(self.threshold, linestyle=\"--\", color=\"grey\",\n",
+    "                   linewidth=1, label=f\"Promotion threshold ({self.threshold})\")\n",
+    "        for step, level in self._promotions:\n",
+    "            ax.axvline(step, color=\"green\", linewidth=1.5, linestyle=\":\")\n",
+    "            ax.text(step, ax.get_ylim()[1]*0.9, f\" L{level}\",\n",
+    "                    color=\"green\", fontsize=9)\n",
+    "        ax.set_xlabel(\"Episode\")\n",
+    "        ax.set_ylabel(\"Reward\")\n",
+    "        ax.set_title(\"Training Reward + Curriculum Progression\")\n",
+    "        ax.legend()\n",
+    "        plt.tight_layout()\n",
+    "        plt.savefig(\"curriculum_reward_curve.png\", dpi=150)\n",
+    "        plt.show()\n",
+    "        print(\"📈 Saved curriculum_reward_curve.png\")\n",
+    "\n",
+    "\n",
+    "# ── Demo: simulate 200 episodes of improving reward ──────────────────────────\n",
+    "import random\n",
+    "cb = CurriculumCallback(reward_threshold=5.0, window=50)\n",
+    "for ep in range(200):\n",
+    "    # Simulate reward slowly improving\n",
+    "    synthetic_reward = -5 + ep * 0.08 + random.gauss(0, 2)\n",
+    "    cb.record(synthetic_reward, step=ep)\n",
+    "\n",
+    "cb.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inference-header",
+   "metadata": {},
+   "source": [
+    "## 🔍 Cell 11 — Inference: Before vs After Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "inference-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def run_inference_episode(model, tokenizer, engine, passes, program, max_steps=5):\n",
+    "    \"\"\"Run inference and execute parsed pass arrays until episode termination.\"\"\"\n",
+    "    FastLanguageModel.for_inference(model)\n",
+    "\n",
+    "    _ms = min(int(max_steps), CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n",
+    "    env = CompilerOptimizationEnv(engine, passes, max_steps=_ms)\n",
+    "    obs = env.reset(program)\n",
+    "    system_prompt = build_system_prompt(passes)\n",
+    "\n",
+    "    actions_chosen, rewards_earned = [], []\n",
+    "    done = False\n",
+    "\n",
+    "    print(f\"\\nBaseline cycles: {env.previous_cycles}\")\n",
+    "    print(f\"Initial state:\\n{obs}\\n\")\n",
+    "\n",
+    "    while not done:\n",
+    "        messages = [\n",
+    "            {\"role\": \"system\", \"content\": system_prompt},\n",
+    "            {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose optimization passes:\"},\n",
+    "        ]\n",
+    "        inputs = tokenizer.apply_chat_template(\n",
+    "            messages,\n",
+    "            tokenize=True,\n",
+    "            add_generation_prompt=True,\n",
+    "            return_tensors=\"pt\",\n",
+    "        ).to(model.device)\n",
+    "\n",
+    "        with torch.no_grad():\n",
+    "            outputs = model.generate(\n",
+    "                input_ids=inputs,\n",
+    "                max_new_tokens=32,\n",
+    "                temperature=0.1,\n",
+    "                do_sample=True,\n",
+    "            )\n",
+    "\n",
+    "        raw_output = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip()\n",
+    "\n",
+    "        try:\n",
+    "            parsed_actions = Deliverable2_Formatter.extract_action_array(raw_output)\n",
+    "        except ValueError:\n",
+    "            parsed_actions = []\n",
+    "\n",
+    "        if not parsed_actions:\n",
+    "            result = env.step(\"__invalid__\")\n",
+    "            actions_chosen.append(\"__invalid__\")\n",
+    "            rewards_earned.append(result.reward)\n",
+    "            obs = result.observation\n",
+    "            done = result.done\n",
+    "            print(f\"Step {len(actions_chosen)}: invalid output '{raw_output}' → reward={result.reward:+.2f}\")\n",
+    "            continue\n",
+    "\n",
+    "        for action in parsed_actions[:CompilerOptimizationEnv.MAX_EPISODE_STEPS]:\n",
+    "            a = str(action).strip()\n",
+    "            al = a.lower()\n",
+    "            if al in (\"done\", \"stop\"):\n",
+    "                result = env.step(\"done\" if al == \"done\" else \"STOP\")\n",
+    "                actions_chosen.append(al)\n",
+    "                rewards_earned.append(result.reward)\n",
+    "                obs = result.observation\n",
+    "                done = result.done\n",
+    "                print(f\"Step {len(actions_chosen)}: stop ('{a}') → reward={result.reward:+.2f}\")\n",
+    "                break\n",
+    "            action = al\n",
+    "            if action not in passes:\n",
+    "                for k in passes:\n",
+    "                    if k.lower() == al:\n",
+    "                        action = k\n",
+    "                        break\n",
+    "            result = env.step(action)\n",
+    "            actions_chosen.append(action)\n",
+    "            rewards_earned.append(result.reward)\n",
+    "            obs = result.observation\n",
+    "            done = result.done\n",
+    "            print(f\"Step {len(actions_chosen)}: '{action}' → reward={result.reward:+.2f}\")\n",
+    "            if done:\n",
+    "                break\n",
+    "\n",
+    "    summary = env._episode_summary()\n",
+    "    print(f\"\\n{'─'*40}\")\n",
+    "    print(f\"Total improvement: {summary['total_improvement_pct']:.1f}%\")\n",
+    "    print(f\"Final cycles: {summary['final_cycles']} (was {summary['baseline_cycles']})\")\n",
+    "    return summary\n",
+    "\n",
+    "\n",
+    "# Uncomment after training:\n",
+    "# summary = run_inference_episode(\n",
+    "#     model, tokenizer, MockEngine(), MOCK_PASSES, SAMPLE_PROGRAM\n",
+    "# )\n",
+    "\n",
+    "print(\"✅ Inference cell ready. Uncomment the last block after training to run.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "tips-header",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 📝 Notes & Tips\n",
+    "\n",
+    "| What | Why it matters |\n",
+    "|------|----------------|\n",
+    "| `num_generations=4` | GRPO needs K≥2 rollouts per prompt to compute group-relative advantage. K=4 balances diversity vs. compute. |\n",
+    "| `beta=0.04` | KL penalty keeping policy close to reference. Too low → mode collapse. Too high → no learning. |\n",
+    "| `max_new_tokens=16` | Actions are one word. This prevents wasted computation and keeps the model from adding explanations. |\n",
+    "| Equivalence / semantic failure | The shaped reward function returns `−1000.0` when a candidate fails verification (RLVR hard gate). |\n",
+    "| Terminal reward | Final cycle savings (scaled) are given at `STOP`, step cap (5), or at end of a pass list; non-terminal steps use a small constant cost. |\n",
+    "| Soft invalid-action termination | 3 consecutive invalid actions → end. Single mistakes don't kill the episode; the agent can recover. |\n",
+    "| `deepcopy` on all pass inputs | Role 3's passes mutate dicts in-place. Without this, `original_program` gets corrupted and the verifier fails spuriously. |\n",
+    "\n",
+    "**Next steps:**\n",
+    "- Replace `MockEngine` with Role 1's real engine\n",
+    "- Replace `MOCK_PASSES` with Role 3's real passes\n",
+    "- Push environment to HuggingFace Spaces: `openenv init && openenv deploy`\n",
+    "- Add W&B sweep to tune `beta`, `learning_rate`, `num_generations`"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

space/space/space/space/space/space/space/reverse_pass/toyir_rl_support.py

@@ -0,0 +1,473 @@
+"""
+Toy-IR RL training support: PASSES, mock engine, environment, terminal-weighted shaped reward.
+Used by training notebooks. Do not import toy_vm / verifier / reward_utils here.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+# Allow notebooks in subdirs to import when cwd is set to project root
+_ROOT = os.path.dirname(os.path.abspath(__file__))
+
+
+@dataclass
+class StepResult:
+    observation: str
+    reward: float
+    done: bool
+    info: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class EpisodeStats:
+    steps_taken: int = 0
+    total_reward: float = 0.0
+    passes_applied: List[str] = field(default_factory=list)
+    invalid_actions: int = 0
+    no_ops: int = 0
+    baseline_cycles: int = 0
+    final_cycles: int = 0
+
+    @property
+    def total_improvement_pct(self) -> float:
+        if self.baseline_cycles == 0:
+            return 0.0
+        return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0
+
+
+class MCPEnvironment:
+    def reset(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def step(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def state(self):
+        raise NotImplementedError
+
+
+def _op_u(ins: dict) -> str:
+    return str(ins.get("op", "")).upper()
+
+
+def _first_numeric(ins: dict) -> Optional[int]:
+    for k in ("src1", "src2"):
+        v = ins.get(k)
+        if isinstance(v, int) and v > 0:
+            return v
+    args = ins.get("args", [])
+    if args:
+        a0 = args[0]
+        try:
+            if isinstance(a0, int):
+                return a0
+            return int(str(a0), 0)
+        except (ValueError, TypeError):
+            return None
+    v = ins.get("src1")
+    if isinstance(v, int):
+        return v
+    return None
+
+
+# --- forward / reverse passes (list[dict] IR) ---------------------------------
+
+
+def constant_folding(program: List[dict]) -> List[dict]:
+    p = copy.deepcopy(program)
+    return p[:-1] if len(p) > 1 else p
+
+
+def dead_code_elimination(program: List[dict]) -> List[dict]:
+    p = copy.deepcopy(program)
+    return p[:-1] if len(p) > 2 else p
+
+
+def peephole_optimization(program: List[dict]) -> List[dict]:
+    return copy.deepcopy(program)
+
+
+def expand_constant(program: List[dict]) -> List[dict]:
+    """
+    Deoptimization: split one CONST into two consts + ADD to raise cycle count sometimes,
+    giving forward passes a different surface (reverse exploration).
+    """
+    p = copy.deepcopy(program)
+    for i, ins in enumerate(p):
+        if _op_u(ins) not in ("CONST",):
+            continue
+        v = _first_numeric(ins) if _first_numeric(ins) is not None else None
+        if v is None and ins.get("args"):
+            try:
+                v = int(ins["args"][0])
+            except (ValueError, TypeError, IndexError, KeyError):
+                v = None
+        if v is None or v < 2:
+            continue
+        d = str(ins.get("dest", "t0"))
+        h, rest = v // 2, v - (v // 2)
+        repl = [
+            {**{k: ins[k] for k in ins if k in ("type",)}, "op": "const", "dest": f"{d}_a", "args": [str(h)]},
+            {**{k: ins[k] for k in ins if k in ("type",)}, "op": "const", "dest": f"{d}_b", "args": [str(rest)]},
+            {
+                "op": "add",
+                "dest": d,
+                "args": [f"{d}_a", f"{d}_b"],
+                **{k: ins[k] for k in ins if k in ("type",) and k not in ("op", "dest", "args")},
+            },
+        ]
+        p[i : i + 1] = repl
+        return p
+    return p
+
+
+def duplicate_computation(program: List[dict]) -> List[dict]:
+    """
+    Deoptimization: duplicate a binary op into a new destination (redundant work).
+    """
+    p = copy.deepcopy(program)
+    for i, ins in enumerate(p):
+        ou = _op_u(ins)
+        if ou in ("ADD", "MUL", "SUB", "DIV"):
+            dest = str(ins.get("dest", "t0"))
+            dup = copy.deepcopy(ins)
+            dup["dest"] = f"{dest}_d"
+            p.insert(i + 1, dup)
+            return p
+    return p
+
+
+PASSES: Dict[str, Any] = {
+    "constant_folding": constant_folding,
+    "dead_code_elimination": dead_code_elimination,
+    "peephole_optimization": peephole_optimization,
+    "expand_constant": expand_constant,
+    "duplicate_computation": duplicate_computation,
+}
+
+MOCK_PASSES = PASSES  # alias for existing notebook references
+
+
+class MockEngine:
+    """Stub: cycle count = instruction list length; always equivalent."""
+
+    def execute_and_count_cycles(self, program: List[dict]) -> int:
+        return len(program)
+
+    def verify_equivalence(self, original, candidate) -> bool:
+        return True
+
+
+# Back-compat: some notebooks use these names
+MOCK_ENGINE = MockEngine()
+
+
+# --- Shaped reward (user-specified) -------------------------------------------
+
+
+def compute_shaped_reward(
+    equivalent: bool,
+    is_terminal: bool,
+    original_cycles: int,
+    current_cycles: int,
+) -> float:
+    # RLVR hard gate — broken optimizations get massive penalty
+    if not equivalent:
+        return -1000.0
+
+    # Terminal reward: large signal proportional to final cycle reduction
+    # This is what makes reverse passes learnable — agent gets rewarded for the END state, not the path
+    if is_terminal:
+        if original_cycles == 0:
+            return 0.0
+        relative_savings = (original_cycles - current_cycles) / original_cycles
+        return relative_savings * 100.0
+
+    # Per-step: tiny constant cost to discourage infinite thrashing
+    # CRITICAL: do NOT punish cycle increases here. Per-step cycle-delta reward kills reverse-pass exploration.
+    return -0.1
+
+
+# --- Environment --------------------------------------------------------------
+
+
+class CompilerOptimizationEnv(MCPEnvironment):
+    TIME_TAX: float = 1.0
+    NO_OP_PENALTY: float = -2.0
+    INVALID_ACTION_PENALTY: float = -5.0
+    MAX_INVALID_ACTIONS: int = 3
+    TERMINAL_BONUS_SCALE: float = 0.5
+    MAX_EPISODE_STEPS: int = 5  # hard cap (Change 4)
+
+    def __init__(
+        self,
+        role1_engine: Any,
+        role3_passes: Dict[str, Any],
+        max_steps: int = 5,
+        curriculum_level: int = 1,
+    ):
+        self.engine = role1_engine
+        self.passes = role3_passes
+        self.max_steps = max(1, int(max_steps))
+        self.curriculum_level = curriculum_level
+        self._valid_actions: frozenset = frozenset(self.passes.keys()) | {"STOP", "done", "DONE"}
+        self._stats: Optional[EpisodeStats] = None
+        self.original_program: Optional[List[dict]] = None
+        self.current_program: Optional[List[dict]] = None
+        self.previous_cycles = 0
+        self._consecutive_invalid = 0
+
+    def reset(self, new_program_json: List[dict]) -> str:
+        self.original_program = copy.deepcopy(new_program_json)
+        self.current_program = copy.deepcopy(new_program_json)
+        self.previous_cycles = self._safe_count_cycles(self.current_program)
+        self._consecutive_invalid = 0
+        self._stats = EpisodeStats(
+            baseline_cycles=self.previous_cycles,
+            final_cycles=self.previous_cycles,
+        )
+        return self.state()
+
+    def state(self) -> str:
+        return self._program_to_pseudoasm(self.current_program)
+
+    def step(self, action_string: str) -> StepResult:
+        assert self._stats is not None, "Call reset() before step()."
+        aup = str(action_string).upper()
+        is_stop = aup in ("STOP", "DONE")
+        if is_stop:
+            self._stats.steps_taken += 1
+            orig = int(self._stats.baseline_cycles)
+            cur = self._safe_count_cycles(self.current_program)
+            ok = bool(self.engine.verify_equivalence(self.original_program, self.current_program))
+            reward = compute_shaped_reward(ok, True, orig, cur)
+            self._stats.total_reward += reward
+            return StepResult(
+                self.state(),
+                reward,
+                True,
+                {
+                    "reason": "stop",
+                    "is_terminal": True,
+                    "stop_token": aup,
+                    "episode_stats": self._episode_summary(),
+                },
+            )
+        if action_string not in self.passes:
+            return self._handle_invalid_action(action_string)
+
+        self._stats.steps_taken += 1
+        candidate = self.passes[action_string](copy.deepcopy(self.current_program))
+        if not self.engine.verify_equivalence(self.original_program, candidate):
+            return self._handle_semantic_violation()
+        new_cycles = self._safe_count_cycles(candidate)
+        reward, info = self._compute_reward(action_string, new_cycles)
+        self.current_program = candidate
+        self.previous_cycles = new_cycles
+        self._stats.final_cycles = new_cycles
+        self._stats.total_reward += reward
+        self._stats.passes_applied.append(action_string)
+        self._consecutive_invalid = 0
+        done = self._stats.steps_taken >= self.max_steps
+        if done:
+            info["reason"] = "max_steps_reached"
+            info["episode_stats"] = self._episode_summary()
+        return StepResult(self.state(), reward, done, info)
+
+    def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, dict]:
+        info: Dict[str, Any] = {"action": action}
+        orig = int(self._stats.baseline_cycles)
+        is_term = self._stats.steps_taken >= self.max_steps
+        if new_cycles == self.previous_cycles and not is_term:
+            self._stats.no_ops += 1
+            info["no_op"] = True
+        reward = compute_shaped_reward(True, is_term, orig, new_cycles)
+        info["is_terminal"] = is_term
+        info["new_cycles"] = new_cycles
+        return reward, info
+
+    def _compute_crash_penalty(self) -> float:
+        return -1000.0
+
+    def _handle_invalid_action(self, action: str) -> StepResult:
+        self._consecutive_invalid += 1
+        self._stats.invalid_actions += 1
+        done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS
+        info: Dict[str, Any] = {
+            "error": f"Unknown action: '{action}'",
+            "valid_actions": sorted(self._valid_actions),
+            "consecutive_invalid": self._consecutive_invalid,
+        }
+        if done:
+            info["reason"] = "too_many_invalid_actions"
+            info["episode_stats"] = self._episode_summary()
+        return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info)
+
+    def _handle_semantic_violation(self) -> StepResult:
+        return StepResult(
+            self.state(),
+            self._compute_crash_penalty(),
+            True,
+            {
+                "error": "Semantic equivalence check FAILED.",
+                "reason": "semantic_violation",
+                "episode_stats": self._episode_summary(),
+            },
+        )
+
+    @staticmethod
+    def _program_to_pseudoasm(program: Optional[List[dict]]) -> str:
+        if not program:
+            return "; (empty program)"
+        lines = []
+        for i, instr in enumerate(program):
+            op = instr.get("op", "NOP")
+            args = instr.get("args", [])
+            dest = instr.get("dest")
+            typ = instr.get("type", "")
+            arg_str = ", ".join(str(a) for a in args)
+            type_hint = f":{typ}" if typ else ""
+            if dest:
+                lines.append(f"  {i:>3}:  {dest}{type_hint} = {op}  {arg_str}")
+            else:
+                lines.append(f"  {i:>3}:  {op}  {arg_str}")
+        return "\n".join(lines)
+
+    def _safe_count_cycles(self, program: List[dict]) -> int:
+        return max(0, int(self.engine.execute_and_count_cycles(program)))
+
+    def _episode_summary(self) -> dict:
+        s = self._stats
+        return {
+            "steps": s.steps_taken,
+            "total_reward": round(s.total_reward, 3),
+            "passes_applied": s.passes_applied,
+            "invalid_actions": s.invalid_actions,
+            "no_ops": s.no_ops,
+            "baseline_cycles": s.baseline_cycles,
+            "final_cycles": s.final_cycles,
+            "total_improvement_pct": round(s.total_improvement_pct, 3),
+        }
+
+    def available_actions(self) -> List[str]:
+        return sorted(self._valid_actions)
+
+
+SAMPLE_PROGRAM: List[dict] = [
+    {"op": "const", "dest": "x", "args": ["5"], "type": "int"},
+    {"op": "const", "dest": "y", "args": ["3"], "type": "int"},
+    {"op": "add", "dest": "z", "args": ["x", "y"], "type": "int"},
+    {"op": "mul", "dest": "w", "args": ["z", "x"], "type": "int"},
+    {"op": "ret", "args": ["w"]},
+]
+
+
+# --- Rollout for GRPO reward (terminal-weighted shaping) ------------------------
+
+
+def rollout_shaped_return(
+    program: List[dict],
+    action_names: List[str],
+    engine: Any,
+    passes: Optional[Dict[str, Any]] = None,
+) -> Tuple[float, int, int, bool, bool, bool]:
+    """
+    Returns:
+      (total_shaped_reward, n_expand_constant, n_duplicate, hit_stop, hit_step_cap, parse_ok)
+    If parse_ok is False (unknown pass in sequence), first value is -1000.0.
+    Sums per-step `env.step` shaped rewards, adds a terminal `compute_shaped_reward` when
+    the action list ends without an explicit STOP (or already terminal from step cap).
+    """
+    passes = passes or PASSES
+    cap = CompilerOptimizationEnv.MAX_EPISODE_STEPS
+    env = CompilerOptimizationEnv(engine, passes, max_steps=cap)
+    env.reset(program)
+    total = 0.0
+    n_e = 0
+    n_d = 0
+    hit_stop = False
+    hit_cap = False
+    for raw in action_names:
+        a = str(raw).strip()
+        if not a:
+            continue
+        aup = a.upper()
+        if aup in ("STOP", "DONE"):
+            r = env.step(a)
+            total += r.reward
+            hit_stop = True
+            return total, n_e, n_d, hit_stop, hit_cap, True
+        if a not in passes:
+            return -1000.0, n_e, n_d, hit_stop, hit_cap, False
+        if a == "expand_constant":
+            n_e += 1
+        elif a == "duplicate_computation":
+            n_d += 1
+        r = env.step(a)
+        total += r.reward
+        if r.info.get("reason") == "semantic_violation":
+            return -1000.0, n_e, n_d, hit_stop, hit_cap, True
+        if r.done and r.info.get("reason") == "max_steps_reached":
+            hit_cap = True
+            return total, n_e, n_d, hit_stop, hit_cap, True
+    ok = bool(engine.verify_equivalence(env.original_program, env.current_program))
+    cur = int(engine.execute_and_count_cycles(env.current_program or []))
+    orig = int(env._stats.baseline_cycles)
+    total += compute_shaped_reward(ok, True, orig, cur)
+    return total, n_e, n_d, hit_stop, hit_cap, True
+
+
+# Global window for logging (50-episode / completion windows)
+_REVERSE_LOG_WINDOW: List[dict] = []
+
+
+def log_reverse_pass_stats_for_completion(
+    used_expand: int,
+    used_dup: int,
+) -> None:
+    """Call once per training completion. Logs every 50 'episodes' (completions)."""
+    global _REVERSE_LOG_WINDOW
+    any_r = (used_expand + used_dup) > 0
+    _REVERSE_LOG_WINDOW.append(
+        {
+            "any_reverse": any_r,
+            "expand": used_expand,
+            "dup": used_dup,
+        }
+    )
+    if len(_REVERSE_LOG_WINDOW) < 50:
+        return
+    w = _REVERSE_LOG_WINDOW
+    _REVERSE_LOG_WINDOW = []
+    reverse_pass_episodes = sum(1 for e in w if e["any_reverse"])
+    expand_constant_count = sum(e["expand"] for e in w)
+    duplicate_computation_count = sum(e["dup"] for e in w)
+    payload = {
+        "reverse_pass_episodes": reverse_pass_episodes,
+        "expand_constant_count": expand_constant_count,
+        "duplicate_computation_count": duplicate_computation_count,
+    }
+    try:
+        import wandb
+
+        if wandb.run is not None:
+            wandb.log(payload)
+    except Exception:
+        pass
+    print(
+        f"[reverse_pass/50] reverse_pass_episodes={reverse_pass_episodes} "
+        f"expand_constant_count={expand_constant_count} "
+        f"duplicate_computation_count={duplicate_computation_count}"
+    )
+
+
+def ensure_import_path() -> None:
+    d = _ROOT
+    if d and d not in sys.path:
+        sys.path.insert(0, d)

space/space/space/space/space/space/space/space/space/app.py

@@ -23,6 +23,7 @@ from runtime_core import (
     MockEngine,
     SAMPLE_PROGRAM,
 )
+from train import run_toy_training
 
 
 DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2)
@@ -113,8 +114,9 @@ def build_demo() -> gr.Blocks:
                 # Compiler optimization (Toy-IR) — interactive demo
 
                 This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing,
-                plus the **CompilerOptimizationEnv** smoke loop from your notebooks. Full GRPO / Unsloth
-                training belongs on Colab or a GPU Space.
+                plus the **CompilerOptimizationEnv** loop. A **toy REINFORCE** tab trains a tiny
+                stateless policy over the mock passes (see `train.py`). Full **GRPO + LLM + Unsloth**
+                still belongs on Colab or a GPU machine.
                 """
             ).strip()
         )
@@ -154,6 +156,28 @@ def build_demo() -> gr.Blocks:
                 out_ep = gr.Textbox(label="Log", lines=20)
                 gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep])
 
+            with gr.Tab("Toy training (REINFORCE)"):
+                gr.Markdown(
+                    textwrap.dedent(
+                        """
+                        Trains a **stateless** categorical policy over the three mock passes
+                        (`constant_folding`, `dead_code_elimination`, `loop_unrolling`) using
+                        **REINFORCE** in pure Python. Same logic as: `python train.py --episodes 50`.
+
+                        This is a CPU smoke run, not an LLM. For real GRPO, use your project notebooks
+                        on a GPU.
+                        """
+                    ).strip()
+                )
+                tr_ep = gr.Slider(5, 200, value=50, step=1, label="episodes")
+                tr_ms = gr.Slider(2, 20, value=8, step=1, label="max_steps per episode")
+                tr_seed = gr.Number(value=0, label="random seed", precision=0)
+                tr_lr = gr.Slider(0.01, 0.5, value=0.15, step=0.01, label="learning rate")
+                out_tr = gr.Textbox(label="Training log", lines=18)
+                gr.Button("Run training", variant="primary").click(
+                    run_toy_training, [tr_ep, tr_ms, tr_seed, tr_lr], [out_tr]
+                )
+
         gr.Markdown(
             "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, "
             "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`."

space/space/space/space/space/space/space/space/space/space/space/runtime_core.py

@@ -0,0 +1,316 @@
+"""
+Shared runtime for the Gradio Space: mock compiler env + Deliverable 2 formatting.
+Sourced from `compiler_optimization_grpo.ipynb` and
+`role2_deliverable3_training_loop (2) (1) (1).ipynb`.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+# --- Deliverable 2 (LLM-facing pseudo-asm + pass-array parsing) -----------------
+
+
+class Deliverable2_Formatter:
+    @staticmethod
+    def translate_state(raw_json: list) -> str:
+        """Translate raw JSON IR into compact pseudo-assembly."""
+        if not isinstance(raw_json, list) or not raw_json:
+            return "; (empty program — 0 instructions)"
+
+        pseudo_assembly: list[str] = []
+        for i, instruction in enumerate(raw_json):
+            if not isinstance(instruction, dict):
+                pseudo_assembly.append(f"{i}. NOP")
+                continue
+
+            op = str(instruction.get("op", "UNKNOWN")).upper()
+            args = ", ".join(str(arg) for arg in instruction.get("args", []))
+            dest = instruction.get("dest", "")
+            if dest:
+                line = f"{i}. {dest} = {op} {args}".rstrip()
+            else:
+                line = f"{i}. {op} {args}".rstrip()
+            pseudo_assembly.append(line)
+
+        return "\n".join(pseudo_assembly)
+
+    @staticmethod
+    def extract_action_array(llm_output: str) -> list:
+        """Best-effort extraction of JSON pass arrays from noisy LLM output."""
+        text = (llm_output or "").strip()
+        if not text:
+            raise ValueError("Invalid JSON format")
+
+        try:
+            parsed = json.loads(text)
+            if isinstance(parsed, list):
+                return parsed
+        except json.JSONDecodeError:
+            pass
+
+        cleaned = re.sub(r"```(?:json)?", "", text, flags=re.IGNORECASE).replace("```", "").strip()
+        if cleaned != text:
+            try:
+                parsed = json.loads(cleaned)
+                if isinstance(parsed, list):
+                    return parsed
+            except json.JSONDecodeError:
+                pass
+
+        match = re.search(r"\[.*?\]", text, re.DOTALL)
+        if match:
+            candidate = match.group(0)
+            try:
+                parsed = json.loads(candidate)
+                if isinstance(parsed, list):
+                    return parsed
+            except json.JSONDecodeError:
+                try:
+                    parsed = json.loads(candidate.replace("'", '"'))
+                    if isinstance(parsed, list):
+                        return parsed
+                except json.JSONDecodeError:
+                    pass
+
+        raise ValueError("Invalid JSON format")
+
+
+# --- OpenEnv-style compiler environment (mock engine) --------------------------
+
+
+class MCPEnvironment:
+    """Minimal stub. In production: `from openenv import MCPEnvironment`."""
+
+    def reset(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def step(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def state(self):
+        raise NotImplementedError
+
+
+@dataclass
+class StepResult:
+    observation: str
+    reward: float
+    done: bool
+    info: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class EpisodeStats:
+    steps_taken: int = 0
+    total_reward: float = 0.0
+    passes_applied: List[str] = field(default_factory=list)
+    invalid_actions: int = 0
+    no_ops: int = 0
+    baseline_cycles: int = 0
+    final_cycles: int = 0
+
+    @property
+    def total_improvement_pct(self) -> float:
+        if self.baseline_cycles == 0:
+            return 0.0
+        return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0
+
+
+class CompilerOptimizationEnv(MCPEnvironment):
+    TIME_TAX: float = 1.0
+    NO_OP_PENALTY: float = -2.0
+    INVALID_ACTION_PENALTY: float = -5.0
+    MAX_INVALID_ACTIONS: int = 3
+    TERMINAL_BONUS_SCALE: float = 0.5
+
+    def __init__(
+        self,
+        role1_engine,
+        role3_passes: Dict[str, Any],
+        max_steps: int = 10,
+        curriculum_level: int = 1,
+    ):
+        self.engine = role1_engine
+        self.passes = role3_passes
+        self.max_steps = max_steps
+        self.curriculum_level = curriculum_level
+        self._valid_actions = frozenset(self.passes.keys())
+
+        self._stats: Optional[EpisodeStats] = None
+        self.original_program = None
+        self.current_program = None
+        self.previous_cycles = 0
+        self._consecutive_invalid = 0
+
+    def reset(self, new_program_json: List[Dict]) -> str:
+        self.original_program = copy.deepcopy(new_program_json)
+        self.current_program = copy.deepcopy(new_program_json)
+        self.previous_cycles = self._safe_count_cycles(self.current_program)
+        self._consecutive_invalid = 0
+        self._stats = EpisodeStats(
+            baseline_cycles=self.previous_cycles,
+            final_cycles=self.previous_cycles,
+        )
+        return self.state()
+
+    def state(self) -> str:
+        assert self.current_program is not None
+        return self._program_to_pseudoasm(self.current_program)
+
+    def step(self, action_string: str) -> StepResult:
+        assert self._stats is not None, "Call reset() before step()."
+        self._stats.steps_taken += 1
+
+        if action_string not in self._valid_actions:
+            return self._handle_invalid_action(action_string)
+
+        candidate_program = self.passes[action_string](copy.deepcopy(self.current_program))
+
+        is_valid = self.engine.verify_equivalence(self.original_program, candidate_program)
+        if not is_valid:
+            return self._handle_semantic_violation()
+
+        new_cycles = self._safe_count_cycles(candidate_program)
+        reward, info = self._compute_reward(action_string, new_cycles)
+
+        self.current_program = candidate_program
+        self.previous_cycles = new_cycles
+        self._stats.final_cycles = new_cycles
+        self._stats.total_reward += reward
+        self._stats.passes_applied.append(action_string)
+        self._consecutive_invalid = 0
+
+        done = self._stats.steps_taken >= self.max_steps
+        if done:
+            terminal_bonus = self._terminal_bonus()
+            reward += terminal_bonus
+            info["terminal_bonus"] = terminal_bonus
+            info["reason"] = "max_steps_reached"
+            info["episode_stats"] = self._episode_summary()
+
+        return StepResult(self.state(), reward, done, info)
+
+    def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, Dict]:
+        info: Dict[str, Any] = {"action": action}
+        if self.previous_cycles == 0:
+            return -self.TIME_TAX, {**info, "note": "zero_baseline"}
+
+        old_cycles = self.previous_cycles
+        delta_pct = ((old_cycles - new_cycles) / old_cycles) * 100.0
+
+        if new_cycles == old_cycles:
+            reward = self.NO_OP_PENALTY
+            if self._stats is not None:
+                self._stats.no_ops += 1
+            info["no_op"] = True
+        else:
+            reward = delta_pct - self.TIME_TAX
+            info["delta_pct"] = round(delta_pct, 3)
+
+        info["prev_cycles"] = old_cycles
+        info["new_cycles"] = new_cycles
+        return reward, info
+
+    def _terminal_bonus(self) -> float:
+        if self._stats is None:
+            return 0.0
+        return max(0.0, self._stats.total_improvement_pct * self.TERMINAL_BONUS_SCALE)
+
+    def _compute_crash_penalty(self) -> float:
+        return -2.0 * (100.0 * self.max_steps)
+
+    def _handle_invalid_action(self, action: str) -> StepResult:
+        self._consecutive_invalid += 1
+        if self._stats is not None:
+            self._stats.invalid_actions += 1
+        done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS
+        info = {
+            "error": f"Unknown action: '{action}'",
+            "valid_actions": sorted(self._valid_actions),
+            "consecutive_invalid": self._consecutive_invalid,
+        }
+        if done:
+            info["reason"] = "too_many_invalid_actions"
+            info["episode_stats"] = self._episode_summary()
+        return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info)
+
+    def _handle_semantic_violation(self) -> StepResult:
+        return StepResult(
+            self.state(),
+            self._compute_crash_penalty(),
+            True,
+            {
+                "error": "Semantic equivalence check FAILED.",
+                "reason": "semantic_violation",
+                "episode_stats": self._episode_summary(),
+            },
+        )
+
+    @staticmethod
+    def _program_to_pseudoasm(program: List[Dict]) -> str:
+        if not program:
+            return "; (empty program)"
+        lines = []
+        for i, instr in enumerate(program):
+            op = instr.get("op", "NOP")
+            args = instr.get("args", [])
+            dest = instr.get("dest")
+            typ = instr.get("type", "")
+            arg_str = ", ".join(str(a) for a in args)
+            type_hint = f":{typ}" if typ else ""
+            if dest:
+                lines.append(f"  {i:>3}:  {dest}{type_hint} = {op}  {arg_str}")
+            else:
+                lines.append(f"  {i:>3}:  {op}  {arg_str}")
+        return "\n".join(lines)
+
+    def _safe_count_cycles(self, program: List[Dict]) -> int:
+        return max(0, int(self.engine.execute_and_count_cycles(program)))
+
+    def _episode_summary(self) -> Dict:
+        s = self._stats
+        if s is None:
+            return {}
+        return {
+            "steps": s.steps_taken,
+            "total_reward": round(s.total_reward, 3),
+            "passes_applied": s.passes_applied,
+            "invalid_actions": s.invalid_actions,
+            "no_ops": s.no_ops,
+            "baseline_cycles": s.baseline_cycles,
+            "final_cycles": s.final_cycles,
+            "total_improvement_pct": round(s.total_improvement_pct, 3),
+        }
+
+    def available_actions(self) -> List[str]:
+        return sorted(self._valid_actions)
+
+
+class MockEngine:
+    """Stub engine: cycles = instruction count, all programs semantically valid."""
+
+    def execute_and_count_cycles(self, program):
+        return len(program)
+
+    def verify_equivalence(self, original, candidate):
+        return True
+
+
+MOCK_PASSES = {
+    "constant_folding": lambda p: p[:-1] if len(p) > 1 else p,
+    "dead_code_elimination": lambda p: p[:-1] if len(p) > 2 else p,
+    "loop_unrolling": lambda p: p,
+}
+
+SAMPLE_PROGRAM = [
+    {"op": "const", "dest": "x", "args": ["5"], "type": "int"},
+    {"op": "const", "dest": "y", "args": ["3"], "type": "int"},
+    {"op": "add", "dest": "z", "args": ["x", "y"], "type": "int"},
+    {"op": "mul", "dest": "w", "args": ["z", "x"], "type": "int"},
+    {"op": "ret", "args": ["w"]},
+]

space/space/space/space/space/space/space/space/space/space/space/space/Dockerfile

@@ -0,0 +1,19 @@
+# Optional: use Hugging Face Space SDK = docker (set `sdk: docker` in README.md).
+# Default README uses Gradio SDK and does not require this image.
+
+FROM python:3.11-slim
+
+WORKDIR /app
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    GRADIO_SERVER_NAME=0.0.0.0
+
+COPY requirements.txt .
+RUN pip install --upgrade pip && pip install -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+
+CMD ["python", "app.py"]

space/space/space/space/space/space/space/space/space/space/space/space/README.md

@@ -1,9 +1,54 @@
 ---
-title: Compiler Brain OpenEnv
-emoji: 🧠
-colorFrom: indigo
-colorTo: purple
-sdk: docker
+title: Compiler Tetris — Toy-IR RL Demo
+emoji: 🧩
+colorFrom: gray
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.12.0
 app_file: app.py
 pinned: false
+license: mit
 ---
+
+# Compiler Tetris (Toy-IR) — Hugging Face Space
+
+Interactive **CPU** demo for the Meta / OpenEnv-style **compiler phase-ordering** project:
+
+- **Deliverable 2 — state translation:** JSON Toy-IR → compact pseudo-assembly for the LLM.
+- **Deliverable 2 — format enforcement:** resilient extraction of a JSON **pass array** from noisy model text.
+- **CompilerOptimizationEnv** (mock Role 1 engine + mock passes): step through passes and inspect rewards.
+
+Full **GRPO / Unsloth** training is not run here (heavy GPU + long installs). Use your notebooks on Colab or a GPU Space for training.
+
+## Source notebooks (parent repo)
+
+- `compiler_optimization_grpo.ipynb`
+- `role2_deliverable3_training_loop (2) (1).ipynb`
+- `compiler_tetris (1).ipynb`
+- `metahack1 (1).ipynb`
+
+## Deploy this folder as a new Space
+
+1. Create a new Space on Hugging Face (SDK: **Gradio**).
+2. Upload the contents of this `hf_space/` directory to the Space repository root (`app.py`, `requirements.txt`, `README.md`).
+3. Optional: add `Dockerfile` only if you switch the Space to **Docker** (see below).
+
+## Optional: Docker SDK instead of Gradio SDK
+
+If you want the Space to build from `Dockerfile`, change the YAML header to:
+
+```yaml
+sdk: docker
+```
+
+and remove Gradio-specific keys (`sdk_version`, `app_file`). The container runs `python app.py`, which listens on the `PORT` environment variable provided by Spaces.
+
+## Local run
+
+```bash
+cd hf_space
+pip install -r requirements.txt
+python app.py
+```
+
+Then open `http://127.0.0.1:7860`.

space/space/space/space/space/space/space/space/space/space/space/space/app.py

@@ -1,182 +1,167 @@
+"""
+Hugging Face Spaces entrypoint — Gradio UI for Toy-IR compiler RL demo.
+"""
+
+from __future__ import annotations
+
 import json
-import logging
-import traceback
-from typing import Any, Dict, List
-
-# OpenEnv SDK import
-from openenv import MCPEnvironment
-
-# ==============================================================================
-# ROLE 1 & 3 IMPORTS
-# TODO: Import the actual execution engine and generator from your teammates
-# ==============================================================================
-# from engine import execute_tac, verify_equivalence
-# from curriculum import generate_level_code
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("CompilerEnvServer")
-
-class CompilerEnv(MCPEnvironment):
-    """
-    The OpenEnv Server Wrapper for the Toy-IR Compiler Pass Optimizer.
-    Acts as the referee between the LLM client and Role 1's Execution Engine.
-    """
-    
-    def __init__(self):
-        super().__init__()
-        # State variables
-        self.raw_json_code = None
-        self.current_state_string = ""
-        self.initial_cycles = 0
-        self.current_step = 0
-        self.max_steps = 10 # Set a max step limit per episode
-        
-        # Cycle Weights (The Physics defined by Role 2)
-        self.cycle_weights = {
-            "ADD": 1,
-            "SUB": 1,
-            "MUL": 4,
-            "DIV": 10,
-            "MEM_LOAD": 20,
-            "STORE": 20
-        }
-
-    def reset(self) -> str:
-        """
-        Grabs unoptimized code, calculates baseline cycles, and translates the 
-        state for the LLM.
-        """
-        self.current_step = 0
-        
-        # 1. Grab new unoptimized code (Role 3 integration)
-        # TODO: Replace with real generator: self.raw_json_code = generate_level_code()
-        self.raw_json_code = self._mock_generator() 
-        
-        # 2. Get baseline cycles (Role 1 integration)
-        # TODO: Replace with real engine: self.initial_cycles, _ = execute_tac(self.raw_json_code, [])
-        self.initial_cycles = 100 
-        
-        # 3. Translate to Pseudo-Assembly to prevent Attention Dilution
-        self.current_state_string = self._translate_state(self.raw_json_code)
-        
-        logger.info(f"Environment Reset. Baseline Cycles: {self.initial_cycles}")
-        return self.state()
-
-    def step(self, action: str) -> Dict[str, Any]:
-        """
-        Executes the LLM's chosen optimization pass, verifies math equivalence, 
-        and calculates the reward.
-        """
-        self.current_step += 1
-        
-        # 1. Parse LLM Action (Regex/JSON robustness)
-        try:
-            # Assuming the LLM outputs a single pass name or a list of passes
-            action_data = json.loads(action)
-            if isinstance(action_data, str):
-                action_array = [action_data]
-            else:
-                action_array = action_data
-            format_bonus = 0.1
-        except json.JSONDecodeError:
-            # Format Trap
-            return self._build_step_response(
-                reward=-2.5, 
-                done=True, 
-                error="Invalid JSON. You must output a valid JSON array of strings."
-            )
-
-        # 2. Execute Code & Verify Equivalence (Role 1 Integration)
-        # TODO: new_cycles, optimized_code = execute_tac(self.raw_json_code, action_array)
-        # TODO: is_valid = verify_equivalence(self.raw_json_code, optimized_code)
-        new_cycles = 80 # Mock Data
-        is_valid = True # Mock Data
-        
-        # 3. Calculate Reward Physics
-        if not is_valid:
-            # Correctness Penalty + Micro-Variance for GRPO
-            penalty = -2.0 - (len(action_array) * 0.01)
-            return self._build_step_response(
-                reward=penalty, 
-                done=True, 
-                error=f"Code equivalence broken by passes: {action_array}"
-            )
-
-        # Calculate Improvement Ratio + Time Tax (-1.0)
-        cycle_improvement_ratio = (self.initial_cycles - new_cycles) / self.initial_cycles
-        time_tax = -0.05 * self.current_step # Small tax to prevent pass spamming
-        reward = cycle_improvement_ratio + format_bonus + time_tax
-        
-        # Update state if sequential, or finish if one-shot
-        # NOTE: For hackathon speed, we treat this as a One-Shot episode
-        done = True 
-        
-        return self._build_step_response(
-            reward=reward, 
-            done=done, 
-            info={"status": "success", "optimized_cycles": new_cycles}
+import os
+import sys
+import textwrap
+from pathlib import Path
+
+import gradio as gr
+
+_ROOT = Path(__file__).resolve().parent
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+from runtime_core import (
+    CompilerOptimizationEnv,
+    Deliverable2_Formatter,
+    MOCK_PASSES,
+    MockEngine,
+    SAMPLE_PROGRAM,
+)
+
+
+DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2)
+
+
+def translate_ir(ir_json: str) -> tuple[str, str]:
+    try:
+        data = json.loads(ir_json.strip() or "[]")
+    except json.JSONDecodeError as e:
+        return "", f"Invalid JSON: {e}"
+    if not isinstance(data, list):
+        return "", "JSON root must be a list of instruction objects."
+    d2 = Deliverable2_Formatter.translate_state(data)
+    env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10)
+    env.reset(data)
+    internal = env.state()
+    return d2, internal
+
+
+def parse_llm_output(llm_text: str) -> str:
+    try:
+        arr = Deliverable2_Formatter.extract_action_array(llm_text)
+        return json.dumps(arr, indent=2)
+    except ValueError as e:
+        return str(e)
+
+
+def parse_action_line(line: str) -> list[str]:
+    line = line.strip()
+    if not line:
+        return []
+    try:
+        got = Deliverable2_Formatter.extract_action_array(line)
+        return [str(x).strip() for x in got]
+    except ValueError:
+        parts = [p.strip().strip("\"'") for p in line.split(",") if p.strip()]
+        return [p.lower() for p in parts]
+
+
+def run_episode(ir_json: str, actions_multiline: str, max_steps: int) -> str:
+    try:
+        program = json.loads(ir_json.strip() or "[]")
+    except json.JSONDecodeError as e:
+        return f"Invalid program JSON: {e}"
+    if not isinstance(program, list):
+        return "Program must be a JSON list."
+
+    lines = [ln for ln in actions_multiline.splitlines() if ln.strip()]
+    actions: list[str] = []
+    for ln in lines:
+        actions.extend(parse_action_line(ln))
+    if not actions:
+        return "No actions parsed. Enter JSON arrays or comma-separated pass names."
+
+    engine = MockEngine()
+    env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=int(max_steps))
+    obs0 = env.reset(program)
+    log = [
+        f"Baseline cycles: {env.previous_cycles}",
+        f"Initial observation (env):\n{obs0}",
+        "",
+        f"Deliverable 2 pseudo-assembly:\n{Deliverable2_Formatter.translate_state(program)}",
+        "",
+        "--- steps ---",
+    ]
+    for i, act in enumerate(actions, start=1):
+        r = env.step(act)
+        log.append(
+            f"{i}. {act!r} → reward={r.reward:+.3f} done={r.done} cycles={env.previous_cycles}"
+        )
+        if r.info:
+            slim = {k: v for k, v in r.info.items() if k in ("delta_pct", "error", "no_op", "reason", "terminal_bonus")}
+            if slim:
+                log.append(f"   info: {slim}")
+        if r.done:
+            break
+    log.append("")
+    log.append("Episode summary:")
+    log.append(json.dumps(env._episode_summary(), indent=2))
+    return "\n".join(log)
+
+
+def build_demo() -> gr.Blocks:
+    with gr.Blocks(title="Compiler Tetris — Toy-IR RL Demo") as demo:
+        gr.Markdown(
+            textwrap.dedent(
+                """
+                # Compiler optimization (Toy-IR) — interactive demo
+
+                This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing,
+                plus the **CompilerOptimizationEnv** smoke loop from your notebooks. Full GRPO / Unsloth
+                training belongs on Colab or a GPU Space.
+                """
+            ).strip()
+        )
+
+        with gr.Tabs():
+            with gr.Tab("Translate IR (D2)"):
+                gr.Markdown("Paste Toy-IR as a JSON **list** of instruction dicts (`op`, `args`, `dest`, …).")
+                ir_in = gr.Textbox(label="IR JSON", value=DEFAULT_IR, lines=12, max_lines=24)
+                btn_t = gr.Button("Translate", variant="primary")
+                out_d2 = gr.Textbox(label="Deliverable2 pseudo-assembly (LLM-facing)", lines=10)
+                out_env = gr.Textbox(label="Env internal pseudo-asm (with type hints)", lines=10)
+                btn_t.click(translate_ir, [ir_in], [out_d2, out_env])
+
+            with gr.Tab("Parse LLM output (D2)"):
+                gr.Markdown(
+                    "Paste model output. A JSON array is preferred; bracket extraction handles light noise."
+                )
+                llm_in = gr.Textbox(
+                    label="LLM output",
+                    value='Here is the plan: ["constant_folding", "dead_code_elimination"]',
+                    lines=4,
+                )
+                out_parse = gr.Textbox(label="Parsed array or error", lines=6)
+                gr.Button("Parse", variant="primary").click(parse_llm_output, [llm_in], [out_parse])
+
+            with gr.Tab("Env episode (mock engine)"):
+                gr.Markdown(
+                    "One JSON program + actions: each line can be a JSON array or comma-separated names."
+                )
+                ir_ep = gr.Textbox(label="Program JSON", value=DEFAULT_IR, lines=10)
+                acts = gr.Textbox(
+                    label="Actions (one JSON array or comma-list per line)",
+                    value='["constant_folding", "dead_code_elimination"]\nloop_unrolling',
+                    lines=5,
+                )
+                ms = gr.Slider(1, 20, value=10, step=1, label="max_steps")
+                out_ep = gr.Textbox(label="Log", lines=20)
+                gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep])
+
+        gr.Markdown(
+            "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, "
+            "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`."
         )
 
-    def state(self) -> str:
-        """
-        Returns the current observation to the LLM.
-        """
-        return f"Current Step: {self.current_step}/{self.max_steps}\n\n{self.current_state_string}"
-
-    def _translate_state(self, raw_json: List[Dict]) -> str:
-        """
-        Translates raw AST JSON into clean pseudo-assembly.
-        Strips all UUIDs and AST metadata.
-        """
-        pseudo_assembly = []
-        instruction_count = 1
-        
-        for inst in raw_json:
-            op = inst.get("op", "UNKNOWN")
-            src1 = inst.get("src1", "")
-            src2 = inst.get("src2", "")
-            dest = inst.get("dest", "")
-            
-            # Format arguments cleanly
-            args = f"{src1}" if src2 is None else f"{src1}, {src2}"
-            
-            if dest:
-                line = f"{instruction_count}. {dest} = {op} {args}"
-            else:
-                line = f"{instruction_count}. {op} {args}"
-                
-            pseudo_assembly.append(line)
-            instruction_count += 1
-            
-        return "\n".join(pseudo_assembly)
-
-    def _build_step_response(self, reward: float, done: bool, error: str = None, info: dict = None) -> Dict[str, Any]:
-        """Helper to format the standard OpenEnv step return dictionary."""
-        response = {
-            "reward": reward,
-            "done": done,
-            "state": self.state()
-        }
-        if error:
-            response["error"] = error
-        if info:
-            response["info"] = info
-        return response
-
-    def _mock_generator(self):
-        """Mock data so the server runs before Role 1 integrates their engine."""
-        return [
-            {"op": "CONST", "dest": "a", "src1": 2, "src2": None},
-            {"op": "CONST", "dest": "b", "src1": 3, "src2": None},
-            {"op": "ADD", "dest": "c", "src1": "a", "src2": "b"}
-        ]
+    return demo
+
 
 if __name__ == "__main__":
-    logger.info("Initializing CompilerEnv Server...")
-    try:
-        env = CompilerEnv()
-        # openenv.run() or start() depending on the specific MCP wrapper version
-        env.start() 
-    except Exception as e:
-        logger.error(f"Failed to start environment server: {e}")
-        logger.error(traceback.format_exc())
+    port = int(os.environ.get("PORT", "7860"))
+    build_demo().launch(server_name="0.0.0.0", server_port=port)

space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt

@@ -1,24 +1,2 @@
-# OpenEnv SDK (Mandatory for the hackathon judging environment)
-openenv
-
-# Core RL Environment 
-gymnasium
-
-# Hugging Face Training Stack (Versions matched to your D3 notebook)
-trl==0.23.1
-transformers==4.57.1
-peft
-accelerate
-bitsandbytes
-
-# Unsloth for efficient 4-bit QLoRA training on T4 GPUs
-# Note: Unsloth often prefers being installed via their specific pip wheel or git, 
-# but this is standard for a requirements file.
-unsloth
-
-# PyTorch (Will default to standard compatible version if no index is specified)
-torch
-
-# Logging & Visualizations (For Day 2 judging criteria)
-wandb
-matplotlib
+# Hugging Face Spaces — Gradio SDK (CPU demo)
+gradio>=4.44.0,<6

space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_optimization_grpo.ipynb

@@ -0,0 +1,959 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  },
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "title-cell",
+   "metadata": {},
+   "source": [
+    "# 🔧 Compiler Optimization RL Environment\n",
+    "### OpenEnv Hackathon 2026 — Theme #2: Long-Horizon Planning\n",
+    "\n",
+    "**What this notebook does:**\n",
+    "1. Defines a fully OpenEnv-compliant `CompilerOptimizationEnv`\n",
+    "2. Loads `Qwen2.5-3B-Instruct` via Unsloth (4-bit QLoRA)\n",
+    "3. Trains with GRPO (TRL) — LLM learns to pick compiler passes that reduce CPU cycles\n",
+    "4. Runs a smoke-test with a mock engine so you can verify reward logic without real hardware\n",
+    "5. Plots reward curves\n",
+    "\n",
+    "---\n",
+    "**Stack:** `unsloth` · `trl` · `openenv` · `wandb` · `matplotlib`\n",
+    "\n",
+    "> **Runtime:** Google Colab T4 GPU recommended. For the smoke-test only, CPU is fine."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "install-header",
+   "metadata": {},
+   "source": [
+    "## 📦 Cell 1 — Install Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "install-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install all required packages\n",
+    "# Unsloth must be installed before trl to get the right CUDA kernels\n",
+    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n",
+    "!pip install trl datasets transformers accelerate peft bitsandbytes --quiet\n",
+    "!pip install wandb matplotlib --quiet\n",
+    "\n",
+    "# Optional: install openenv for production use\n",
+    "# !pip install openenv --quiet\n",
+    "\n",
+    "print(\"✅ All packages installed\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "env-header",
+   "metadata": {},
+   "source": [
+    "## 🌍 Cell 2 — CompilerOptimizationEnv (OpenEnv-Compliant)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "env-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "CompilerOptimizationEnv — OpenEnv-compatible RL Environment\n",
+    "OpenEnv Hackathon 2026 | Theme #2: Long-Horizon Planning\n",
+    "\n",
+    "Key improvements over naive baseline:\n",
+    "  ✓ Inherits from MCPEnvironment (OpenEnv API compliant)\n",
+    "  ✓ Dynamic crash_penalty scaled to reward range (not hardcoded -1000)\n",
+    "  ✓ Soft termination: invalid actions give 3 chances before episode ends\n",
+    "  ✓ No-op detection: penalises passes that change nothing\n",
+    "  ✓ Terminal bonus: rewards cumulative improvement, not just greedy steps\n",
+    "  ✓ StepResult / EpisodeStats dataclasses for clean interfaces\n",
+    "  ✓ Curriculum level support\n",
+    "\"\"\"\n",
+    "\n",
+    "import copy\n",
+    "import math\n",
+    "from dataclasses import dataclass, field\n",
+    "from typing import Any, Dict, List, Optional, Tuple\n",
+    "\n",
+    "\n",
+    "# ── Stub base class ───────────────────────────────────────────────────────────\n",
+    "# In production replace with: from openenv import MCPEnvironment\n",
+    "class MCPEnvironment:\n",
+    "    \"\"\"Minimal stub. In production: `from openenv import MCPEnvironment`\"\"\"\n",
+    "    def reset(self, *args, **kwargs): raise NotImplementedError\n",
+    "    def step(self, *args, **kwargs): raise NotImplementedError\n",
+    "    def state(self): raise NotImplementedError\n",
+    "\n",
+    "\n",
+    "# ── Data classes ──────────────────────────────────────────────────────────────\n",
+    "@dataclass\n",
+    "class StepResult:\n",
+    "    observation: str\n",
+    "    reward: float\n",
+    "    done: bool\n",
+    "    info: Dict[str, Any] = field(default_factory=dict)\n",
+    "\n",
+    "\n",
+    "@dataclass\n",
+    "class EpisodeStats:\n",
+    "    steps_taken: int = 0\n",
+    "    total_reward: float = 0.0\n",
+    "    passes_applied: List[str] = field(default_factory=list)\n",
+    "    invalid_actions: int = 0\n",
+    "    no_ops: int = 0\n",
+    "    baseline_cycles: int = 0\n",
+    "    final_cycles: int = 0\n",
+    "\n",
+    "    @property\n",
+    "    def total_improvement_pct(self) -> float:\n",
+    "        if self.baseline_cycles == 0:\n",
+    "            return 0.0\n",
+    "        return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0\n",
+    "\n",
+    "\n",
+    "# ── Core Environment ─────────────────────────��────────────────────────────────\n",
+    "class CompilerOptimizationEnv(MCPEnvironment):\n",
+    "    \"\"\"\n",
+    "    RL environment where an LLM agent sequentially applies compiler\n",
+    "    optimization passes to minimise CPU cycle count while preserving\n",
+    "    program semantics.\n",
+    "\n",
+    "    State  : Pseudo-assembly representation of the current IR (text)\n",
+    "    Actions: Named compiler passes (strings) from the `passes` dictionary\n",
+    "    Reward : % cycle improvement per step − time_tax, with terminal bonus\n",
+    "    Done   : Semantic violation | max steps reached | too many invalid actions\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # Hyperparameters\n",
+    "    TIME_TAX: float              = 1.0\n",
+    "    NO_OP_PENALTY: float         = -2.0\n",
+    "    INVALID_ACTION_PENALTY: float = -5.0\n",
+    "    MAX_INVALID_ACTIONS: int     = 3\n",
+    "    TERMINAL_BONUS_SCALE: float  = 0.5\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        role1_engine,\n",
+    "        role3_passes: Dict[str, Any],\n",
+    "        max_steps: int = 10,\n",
+    "        curriculum_level: int = 1,\n",
+    "    ):\n",
+    "        self.engine           = role1_engine\n",
+    "        self.passes           = role3_passes\n",
+    "        self.max_steps        = max_steps\n",
+    "        self.curriculum_level = curriculum_level\n",
+    "        self._valid_actions   = frozenset(self.passes.keys())\n",
+    "\n",
+    "        # Episode state\n",
+    "        self._stats: Optional[EpisodeStats] = None\n",
+    "        self.original_program = None\n",
+    "        self.current_program  = None\n",
+    "        self.previous_cycles  = 0\n",
+    "        self._consecutive_invalid = 0\n",
+    "\n",
+    "    # ── OpenEnv API ───────────────────────────────────────────────────────────\n",
+    "    def reset(self, new_program_json: List[Dict]) -> str:\n",
+    "        self.original_program     = copy.deepcopy(new_program_json)\n",
+    "        self.current_program      = copy.deepcopy(new_program_json)\n",
+    "        self.previous_cycles      = self._safe_count_cycles(self.current_program)\n",
+    "        self._consecutive_invalid = 0\n",
+    "        self._stats = EpisodeStats(\n",
+    "            baseline_cycles=self.previous_cycles,\n",
+    "            final_cycles=self.previous_cycles,\n",
+    "        )\n",
+    "        return self.state()\n",
+    "\n",
+    "    def state(self) -> str:\n",
+    "        return self._program_to_pseudoasm(self.current_program)\n",
+    "\n",
+    "    def step(self, action_string: str) -> StepResult:\n",
+    "        assert self._stats is not None, \"Call reset() before step().\"\n",
+    "        self._stats.steps_taken += 1\n",
+    "\n",
+    "        if action_string not in self._valid_actions:\n",
+    "            return self._handle_invalid_action(action_string)\n",
+    "\n",
+    "        candidate_program = self.passes[action_string](\n",
+    "            copy.deepcopy(self.current_program)\n",
+    "        )\n",
+    "\n",
+    "        is_valid = self.engine.verify_equivalence(self.original_program, candidate_program)\n",
+    "        if not is_valid:\n",
+    "            return self._handle_semantic_violation()\n",
+    "\n",
+    "        new_cycles = self._safe_count_cycles(candidate_program)\n",
+    "        reward, info = self._compute_reward(action_string, new_cycles)\n",
+    "\n",
+    "        self.current_program          = candidate_program\n",
+    "        self.previous_cycles          = new_cycles\n",
+    "        self._stats.final_cycles      = new_cycles\n",
+    "        self._stats.total_reward      += reward\n",
+    "        self._stats.passes_applied.append(action_string)\n",
+    "        self._consecutive_invalid     = 0\n",
+    "\n",
+    "        done = self._stats.steps_taken >= self.max_steps\n",
+    "        if done:\n",
+    "            terminal_bonus            = self._terminal_bonus()\n",
+    "            reward                   += terminal_bonus\n",
+    "            info[\"terminal_bonus\"]   = terminal_bonus\n",
+    "            info[\"reason\"]           = \"max_steps_reached\"\n",
+    "            info[\"episode_stats\"]    = self._episode_summary()\n",
+    "\n",
+    "        return StepResult(self.state(), reward, done, info)\n",
+    "\n",
+    "    # ── Reward logic ──────────────────────────────────────────────────────────\n",
+    "    def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, Dict]:\n",
+    "        info: Dict[str, Any] = {\"action\": action}\n",
+    "        if self.previous_cycles == 0:\n",
+    "            return -self.TIME_TAX, {**info, \"note\": \"zero_baseline\"}\n",
+    "\n",
+    "        delta_pct = ((self.previous_cycles - new_cycles) / self.previous_cycles) * 100.0\n",
+    "\n",
+    "        if new_cycles == self.previous_cycles:\n",
+    "            reward = self.NO_OP_PENALTY\n",
+    "            self._stats.no_ops += 1\n",
+    "            info[\"no_op\"] = True\n",
+    "        else:\n",
+    "            reward = delta_pct - self.TIME_TAX\n",
+    "            info[\"delta_pct\"] = round(delta_pct, 3)\n",
+    "\n",
+    "        info[\"prev_cycles\"] = self.previous_cycles\n",
+    "        info[\"new_cycles\"]  = new_cycles\n",
+    "        return reward, info\n",
+    "\n",
+    "    def _terminal_bonus(self) -> float:\n",
+    "        return max(0.0, self._stats.total_improvement_pct * self.TERMINAL_BONUS_SCALE)\n",
+    "\n",
+    "    def _compute_crash_penalty(self) -> float:\n",
+    "        # 2× best possible episode reward — always catastrophic, never overwhelming\n",
+    "        return -2.0 * (100.0 * self.max_steps)\n",
+    "\n",
+    "    # ── Error handlers ────────────────────────────────────────────────────────\n",
+    "    def _handle_invalid_action(self, action: str) -> StepResult:\n",
+    "        self._consecutive_invalid += 1\n",
+    "        self._stats.invalid_actions += 1\n",
+    "        done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS\n",
+    "        info = {\n",
+    "            \"error\": f\"Unknown action: '{action}'\",\n",
+    "            \"valid_actions\": sorted(self._valid_actions),\n",
+    "            \"consecutive_invalid\": self._consecutive_invalid,\n",
+    "        }\n",
+    "        if done:\n",
+    "            info[\"reason\"] = \"too_many_invalid_actions\"\n",
+    "            info[\"episode_stats\"] = self._episode_summary()\n",
+    "        return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info)\n",
+    "\n",
+    "    def _handle_semantic_violation(self) -> StepResult:\n",
+    "        return StepResult(\n",
+    "            self.state(),\n",
+    "            self._compute_crash_penalty(),\n",
+    "            True,\n",
+    "            {\n",
+    "                \"error\": \"Semantic equivalence check FAILED.\",\n",
+    "                \"reason\": \"semantic_violation\",\n",
+    "                \"episode_stats\": self._episode_summary(),\n",
+    "            },\n",
+    "        )\n",
+    "\n",
+    "    # ── State representation ──────────────────────────────────────────────────\n",
+    "    @staticmethod\n",
+    "    def _program_to_pseudoasm(program: List[Dict]) -> str:\n",
+    "        if not program:\n",
+    "            return \"; (empty program)\"\n",
+    "        lines = []\n",
+    "        for i, instr in enumerate(program):\n",
+    "            op       = instr.get(\"op\", \"NOP\")\n",
+    "            args     = instr.get(\"args\", [])\n",
+    "            dest     = instr.get(\"dest\")\n",
+    "            typ      = instr.get(\"type\", \"\")\n",
+    "            arg_str  = \", \".join(str(a) for a in args)\n",
+    "            type_hint = f\":{typ}\" if typ else \"\"\n",
+    "            if dest:\n",
+    "                lines.append(f\"  {i:>3}:  {dest}{type_hint} = {op}  {arg_str}\")\n",
+    "            else:\n",
+    "                lines.append(f\"  {i:>3}:  {op}  {arg_str}\")\n",
+    "        return \"\\n\".join(lines)\n",
+    "\n",
+    "    # ── Utilities ─────────────────────────────────────────────────────────────\n",
+    "    def _safe_count_cycles(self, program: List[Dict]) -> int:\n",
+    "        return max(0, int(self.engine.execute_and_count_cycles(program)))\n",
+    "\n",
+    "    def _episode_summary(self) -> Dict:\n",
+    "        s = self._stats\n",
+    "        return {\n",
+    "            \"steps\":                 s.steps_taken,\n",
+    "            \"total_reward\":          round(s.total_reward, 3),\n",
+    "            \"passes_applied\":        s.passes_applied,\n",
+    "            \"invalid_actions\":       s.invalid_actions,\n",
+    "            \"no_ops\":                s.no_ops,\n",
+    "            \"baseline_cycles\":       s.baseline_cycles,\n",
+    "            \"final_cycles\":          s.final_cycles,\n",
+    "            \"total_improvement_pct\": round(s.total_improvement_pct, 3),\n",
+    "        }\n",
+    "\n",
+    "    def available_actions(self) -> List[str]:\n",
+    "        return sorted(self._valid_actions)\n",
+    "\n",
+    "\n",
+    "print(\"✅ CompilerOptimizationEnv defined\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "smoke-header",
+   "metadata": {},
+   "source": [
+    "## 🧪 Cell 3 — Smoke Test (No GPU / Real Engine Needed)\n",
+    "Validates the entire reward pipeline with a mock engine. Run this before spending compute."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "smoke-test-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MockEngine:\n",
+    "    \"\"\"Stub engine: cycles = instruction count, all programs semantically valid.\"\"\"\n",
+    "    def execute_and_count_cycles(self, program):\n",
+    "        return len(program)\n",
+    "\n",
+    "    def verify_equivalence(self, original, candidate):\n",
+    "        return True\n",
+    "\n",
+    "\n",
+    "MOCK_PASSES = {\n",
+    "    \"constant_folding\":      lambda p: p[:-1] if len(p) > 1 else p,\n",
+    "    \"dead_code_elimination\": lambda p: p[:-1] if len(p) > 2 else p,\n",
+    "    \"loop_unrolling\":        lambda p: p,   # intentional no-op for testing\n",
+    "}\n",
+    "\n",
+    "SAMPLE_PROGRAM = [\n",
+    "    {\"op\": \"const\", \"dest\": \"x\", \"args\": [\"5\"],      \"type\": \"int\"},\n",
+    "    {\"op\": \"const\", \"dest\": \"y\", \"args\": [\"3\"],      \"type\": \"int\"},\n",
+    "    {\"op\": \"add\",   \"dest\": \"z\", \"args\": [\"x\", \"y\"], \"type\": \"int\"},\n",
+    "    {\"op\": \"mul\",   \"dest\": \"w\", \"args\": [\"z\", \"x\"], \"type\": \"int\"},\n",
+    "    {\"op\": \"ret\",   \"args\": [\"w\"]},\n",
+    "]\n",
+    "\n",
+    "engine = MockEngine()\n",
+    "env    = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=6)\n",
+    "obs    = env.reset(SAMPLE_PROGRAM)\n",
+    "\n",
+    "print(\"═\" * 55)\n",
+    "print(\"SMOKE TEST\")\n",
+    "print(\"═\" * 55)\n",
+    "print(f\"Initial state:\\n{obs}\")\n",
+    "print(f\"\\nBaseline cycles : {env.previous_cycles}\")\n",
+    "print(f\"Available actions: {env.available_actions()}\")\n",
+    "print()\n",
+    "\n",
+    "actions_to_try = [\n",
+    "    \"constant_folding\",\n",
+    "    \"dead_code_elimination\",\n",
+    "    \"loop_unrolling\",       # no-op\n",
+    "    \"hallucinated_pass\",    # invalid — but won't kill episode yet\n",
+    "    \"constant_folding\",\n",
+    "    \"dead_code_elimination\",\n",
+    "]\n",
+    "\n",
+    "for action in actions_to_try:\n",
+    "    result = env.step(action)\n",
+    "    tag = \"✗\" if result.reward < 0 else \"✓\"\n",
+    "    print(f\"{tag} '{action}'\")\n",
+    "    print(f\"   reward={result.reward:+.2f}  done={result.done}\")\n",
+    "    relevant = {k: v for k, v in result.info.items()\n",
+    "                if k in (\"delta_pct\", \"error\", \"no_op\", \"reason\",\n",
+    "                         \"terminal_bonus\", \"episode_stats\")}\n",
+    "    if relevant:\n",
+    "        print(f\"   info: {relevant}\")\n",
+    "    print()\n",
+    "    if result.done:\n",
+    "        break\n",
+    "\n",
+    "print(\"✅ Smoke test passed\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "reward-plot-header",
+   "metadata": {},
+   "source": [
+    "## 📊 Cell 4 — Visualise Reward Across a Mock Episode"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "reward-plot-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.ticker as ticker\n",
+    "\n",
+    "# Run a full episode and collect data\n",
+    "env   = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10)\n",
+    "env.reset(SAMPLE_PROGRAM)\n",
+    "\n",
+    "sequence = [\n",
+    "    \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n",
+    "    \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n",
+    "    \"dead_code_elimination\", \"loop_unrolling\", \"constant_folding\",\n",
+    "    \"dead_code_elimination\",\n",
+    "]\n",
+    "\n",
+    "rewards, cycle_counts, actions_log = [], [], []\n",
+    "for act in sequence:\n",
+    "    r = env.step(act)\n",
+    "    rewards.append(r.reward)\n",
+    "    cycle_counts.append(env.previous_cycles)\n",
+    "    actions_log.append(act)\n",
+    "    if r.done:\n",
+    "        break\n",
+    "\n",
+    "steps = list(range(1, len(rewards) + 1))\n",
+    "\n",
+    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n",
+    "fig.suptitle(\"Compiler Optimization Episode — Mock Engine\", fontsize=14, fontweight=\"bold\")\n",
+    "\n",
+    "# Reward per step\n",
+    "colors = [\"#2ecc71\" if r >= 0 else \"#e74c3c\" for r in rewards]\n",
+    "ax1.bar(steps, rewards, color=colors, edgecolor=\"white\", linewidth=0.5)\n",
+    "ax1.axhline(0, color=\"grey\", linewidth=0.8, linestyle=\"--\")\n",
+    "ax1.set_ylabel(\"Reward\")\n",
+    "ax1.set_title(\"Reward per Step (green = positive, red = negative)\")\n",
+    "ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter(\"%.1f\"))\n",
+    "\n",
+    "# Cycle count over time\n",
+    "ax2.plot(steps, cycle_counts, marker=\"o\", color=\"#3498db\", linewidth=2, markersize=6)\n",
+    "ax2.set_xlabel(\"Step\")\n",
+    "ax2.set_ylabel(\"CPU Cycles\")\n",
+    "ax2.set_title(\"CPU Cycle Count Over Episode (lower = better)\")\n",
+    "ax2.set_xticks(steps)\n",
+    "ax2.set_xticklabels(\n",
+    "    [a.replace(\"_\", \"\\n\") for a in actions_log],\n",
+    "    fontsize=7,\n",
+    ")\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig(\"episode_reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(\"📈 Plot saved as episode_reward_curve.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "model-header",
+   "metadata": {},
+   "source": [
+    "## 🤖 Cell 5 — Load Model with Unsloth (QLoRA 4-bit)\n",
+    "> **Requires T4 GPU.** Skip to Cell 9 if you only want to test the environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "model-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "MODEL_NAME  = \"unsloth/Qwen2.5-3B-Instruct\"  # swap to 7B if VRAM allows\n",
+    "MAX_SEQ_LEN = 1024\n",
+    "LORA_RANK   = 16\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name     = MODEL_NAME,\n",
+    "    max_seq_length = MAX_SEQ_LEN,\n",
+    "    dtype          = None,       # auto bf16/fp16\n",
+    "    load_in_4bit   = True,\n",
+    ")\n",
+    "\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r              = LORA_RANK,\n",
+    "    target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
+    "                      \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    lora_alpha     = LORA_RANK * 2,\n",
+    "    lora_dropout   = 0.0,\n",
+    "    bias           = \"none\",\n",
+    "    use_gradient_checkpointing = \"unsloth\",\n",
+    "    random_state   = 42,\n",
+    ")\n",
+    "\n",
+    "print(f\"✅ Loaded {MODEL_NAME} with QLoRA rank={LORA_RANK}\")\n",
+    "print(f\"   GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "prompt-header",
+   "metadata": {},
+   "source": [
+    "## 💬 Cell 6 — System Prompt & Dataset Builder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "prompt-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from datasets import Dataset\n",
+    "\n",
+    "\n",
+    "def build_system_prompt(env: CompilerOptimizationEnv) -> str:\n",
+    "    actions = \"\\n\".join(f\"  - {a}\" for a in env.available_actions())\n",
+    "    return (\n",
+    "        \"You are a compiler optimization agent. Your goal is to reduce \"\n",
+    "        \"CPU cycle count by applying optimization passes to the program below.\\n\\n\"\n",
+    "        f\"Available actions (respond with EXACTLY one per turn):\\n{actions}\\n\"\n",
+    "        \"  - done  (stop early if no further improvement is possible)\\n\\n\"\n",
+    "        \"Rules:\\n\"\n",
+    "        \"  • Output only the action name. No explanation, no markdown.\\n\"\n",
+    "        \"  • Do not invent actions not listed above.\\n\"\n",
+    "        \"  • Repeating a pass that does nothing wastes a step.\\n\"\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def build_dataset(\n",
+    "    programs: list,\n",
+    "    engine,\n",
+    "    passes: dict,\n",
+    ") -> Dataset:\n",
+    "    \"\"\"\n",
+    "    Each row = one episode's initial state.\n",
+    "    GRPO samples K completions (actions) per row to estimate group-relative advantage.\n",
+    "    \"\"\"\n",
+    "    env = CompilerOptimizationEnv(engine, passes, max_steps=10)\n",
+    "    system_prompt = build_system_prompt(env)\n",
+    "\n",
+    "    rows = []\n",
+    "    for prog in programs:\n",
+    "        obs = env.reset(prog)\n",
+    "        prompt = [\n",
+    "            {\"role\": \"system\", \"content\": system_prompt},\n",
+    "            {\"role\": \"user\",   \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n",
+    "        ]\n",
+    "        rows.append({\"prompt\": prompt, \"program_json\": json.dumps(prog)})\n",
+    "\n",
+    "    return Dataset.from_list(rows)\n",
+    "\n",
+    "\n",
+    "# ── Demo: build dataset from mock programs ────────────────────────────────────\n",
+    "def make_mock_programs(n: int = 20) -> list:\n",
+    "    \"\"\"Generate N random mock IR programs for demo purposes.\"\"\"\n",
+    "    import random\n",
+    "    ops   = [\"add\", \"mul\", \"sub\", \"const\", \"load\"]\n",
+    "    progs = []\n",
+    "    for _ in range(n):\n",
+    "        length = random.randint(4, 12)\n",
+    "        prog = [\n",
+    "            {\"op\": random.choice(ops),\n",
+    "             \"dest\": f\"v{i}\",\n",
+    "             \"args\": [f\"v{max(0,i-1)}\"],\n",
+    "             \"type\": \"int\"}\n",
+    "            for i in range(length)\n",
+    "        ]\n",
+    "        prog.append({\"op\": \"ret\", \"args\": [f\"v{length-1}\"]})\n",
+    "        progs.append(prog)\n",
+    "    return progs\n",
+    "\n",
+    "\n",
+    "mock_programs = make_mock_programs(n=30)\n",
+    "train_dataset = build_dataset(mock_programs, MockEngine(), MOCK_PASSES)\n",
+    "\n",
+    "print(f\"✅ Dataset built: {len(train_dataset)} episodes\")\n",
+    "print(f\"   Sample prompt keys: {list(train_dataset[0].keys())}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "reward-fn-header",
+   "metadata": {},
+   "source": [
+    "## 🎯 Cell 7 — Reward Function Factory (for GRPO)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "reward-fn-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "from typing import Any\n",
+    "\n",
+    "\n",
+    "def make_reward_fn(engine, passes, max_steps=10):\n",
+    "    \"\"\"\n",
+    "    Returns a reward function compatible with TRL's GRPOTrainer.\n",
+    "    Each GRPO rollout gets its own fresh env instance — no state leakage.\n",
+    "\n",
+    "    Action parsing strips punctuation the model might add\n",
+    "    (e.g. 'constant_folding.' → 'constant_folding').\n",
+    "    \"\"\"\n",
+    "    def reward_fn(prompts, completions, **kwargs):\n",
+    "        programs = kwargs.get(\"program_json\", [None] * len(completions))\n",
+    "        rewards  = []\n",
+    "\n",
+    "        for completion, prog_json in zip(completions, programs):\n",
+    "            # Parse action from model output\n",
+    "            raw    = completion if isinstance(completion, str) else completion[0][\"content\"]\n",
+    "            action = raw.strip().lower().split()[0] if raw.strip() else \"__invalid__\"\n",
+    "            action = re.sub(r\"[^a-z0-9_]\", \"\", action)  # strip punctuation\n",
+    "\n",
+    "            # Fresh env per rollout\n",
+    "            env = CompilerOptimizationEnv(\n",
+    "                role1_engine = engine,\n",
+    "                role3_passes = passes,\n",
+    "                max_steps    = max_steps,\n",
+    "            )\n",
+    "            program = json.loads(prog_json) if prog_json else []\n",
+    "            env.reset(program)\n",
+    "\n",
+    "            if action == \"done\":\n",
+    "                rewards.append(0.0)   # neutral stop\n",
+    "            else:\n",
+    "                result = env.step(action)\n",
+    "                rewards.append(result.reward)\n",
+    "\n",
+    "        return rewards\n",
+    "\n",
+    "    return reward_fn\n",
+    "\n",
+    "\n",
+    "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=10)\n",
+    "print(\"✅ Reward function factory ready\")\n",
+    "\n",
+    "# Quick sanity check\n",
+    "test_completions = [\"constant_folding\", \"hallucinated_pass\", \"loop_unrolling\"]\n",
+    "test_programs    = [json.dumps(SAMPLE_PROGRAM)] * 3\n",
+    "test_rewards     = reward_fn(\n",
+    "    prompts     = [\"\"] * 3,\n",
+    "    completions = test_completions,\n",
+    "    program_json = test_programs,\n",
+    ")\n",
+    "print(\"\\nReward sanity check:\")\n",
+    "for act, rew in zip(test_completions, test_rewards):\n",
+    "    print(f\"  '{act}' → {rew:+.2f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "trainer-header",
+   "metadata": {},
+   "source": [
+    "## 🚀 Cell 8 — GRPO Trainer Config & Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "trainer-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "\n",
+    "# Optional W&B — comment out if not using\n",
+    "try:\n",
+    "    import wandb\n",
+    "    wandb.init(project=\"openenv-compiler-opt\", name=\"grpo-qwen2.5-3b\")\n",
+    "    REPORT_TO = \"wandb\"\n",
+    "except Exception:\n",
+    "    REPORT_TO = \"none\"\n",
+    "\n",
+    "\n",
+    "grpo_config = GRPOConfig(\n",
+    "    # ── Generation ────────────────────────────────────────────────────────\n",
+    "    num_generations             = 4,     # K rollouts per prompt for group-relative advantage\n",
+    "    max_new_tokens              = 16,    # Actions are 1 word; don't waste context\n",
+    "    temperature                 = 0.9,\n",
+    "    top_p                       = 0.95,\n",
+    "\n",
+    "    # ── Optimisation ──────────────────────────────────────────────────────\n",
+    "    learning_rate               = 5e-6,\n",
+    "    per_device_train_batch_size = 2,\n",
+    "    gradient_accumulation_steps = 4,    # effective batch = 8\n",
+    "    num_train_epochs            = 3,\n",
+    "    max_grad_norm               = 0.5,\n",
+    "\n",
+    "    # ── GRPO-specific ─────────────────────────────────────────────────────\n",
+    "    beta                        = 0.04,  # KL penalty; keeps policy near reference\n",
+    "\n",
+    "    # ── Logging / checkpointing ───────────────────────────────────────────\n",
+    "    output_dir                  = \"./grpo_compiler_checkpoints\",\n",
+    "    logging_steps               = 10,\n",
+    "    save_steps                  = 100,\n",
+    "    report_to                   = REPORT_TO,\n",
+    "\n",
+    "    # ── Reproducibility ───────────────────────────────────────────────────\n",
+    "    seed                        = 42,\n",
+    ")\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model         = model,\n",
+    "    tokenizer     = tokenizer,\n",
+    "    config        = grpo_config,\n",
+    "    train_dataset = train_dataset,\n",
+    "    reward_funcs  = reward_fn,\n",
+    ")\n",
+    "\n",
+    "print(\"✅ Trainer configured\")\n",
+    "print(f\"   num_generations (K)  = {grpo_config.num_generations}\")\n",
+    "print(f\"   effective batch size  = \"\n",
+    "      f\"{grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n",
+    "print(f\"   KL beta              = {grpo_config.beta}\")\n",
+    "print()\n",
+    "print(\"Starting training... (this will take a while on T4)\")\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "save-header",
+   "metadata": {},
+   "source": [
+    "## 💾 Cell 9 — Save Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "save-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAVE_PATH = \"./grpo_compiler_final\"\n",
+    "\n",
+    "model.save_pretrained(SAVE_PATH)\n",
+    "tokenizer.save_pretrained(SAVE_PATH)\n",
+    "\n",
+    "print(f\"✅ Model saved to {SAVE_PATH}\")\n",
+    "\n",
+    "# Optional: push to HuggingFace Hub\n",
+    "# model.push_to_hub(\"your-hf-username/compiler-opt-grpo\")\n",
+    "# tokenizer.push_to_hub(\"your-hf-username/compiler-opt-grpo\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "curriculum-header",
+   "metadata": {},
+   "source": [
+    "## 📈 Cell 10 — Curriculum Callback & Reward Tracking"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "curriculum-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CurriculumCallback:\n",
+    "    \"\"\"\n",
+    "    Tracks rolling mean reward and promotes curriculum level\n",
+    "    when the agent has mastered the current difficulty.\n",
+    "\n",
+    "    Usage: call .record(reward) after every episode.\n",
+    "    Read .level to get current difficulty (1=easy, 2=medium, 3=hard).\n",
+    "    \"\"\"\n",
+    "    def __init__(self, reward_threshold: float = 5.0, window: int = 50):\n",
+    "        self.threshold = reward_threshold\n",
+    "        self.window    = window\n",
+    "        self._history  = []\n",
+    "        self.level     = 1\n",
+    "        self._promotions = []\n",
+    "\n",
+    "    def record(self, reward: float, step: int = None):\n",
+    "        self._history.append(reward)\n",
+    "        if len(self._history) >= self.window:\n",
+    "            mean = sum(self._history[-self.window:]) / self.window\n",
+    "            if mean >= self.threshold and self.level < 3:\n",
+    "                self.level += 1\n",
+    "                self._promotions.append((step or len(self._history), self.level))\n",
+    "                print(f\"[Curriculum] ▲ Promoted to level {self.level} \"\n",
+    "                      f\"(rolling mean={mean:.2f})\")\n",
+    "\n",
+    "    def plot(self):\n",
+    "        import matplotlib.pyplot as plt\n",
+    "        import numpy as np\n",
+    "\n",
+    "        history = self._history\n",
+    "        steps   = list(range(len(history)))\n",
+    "        window  = self.window\n",
+    "        rolling = [\n",
+    "            sum(history[max(0,i-window):i+1]) / min(i+1, window)\n",
+    "            for i in steps\n",
+    "        ]\n",
+    "\n",
+    "        fig, ax = plt.subplots(figsize=(10, 4))\n",
+    "        ax.plot(steps, history, alpha=0.3, color=\"#3498db\", label=\"Episode reward\")\n",
+    "        ax.plot(steps, rolling, color=\"#e74c3c\", linewidth=2,\n",
+    "                label=f\"Rolling mean (w={window})\")\n",
+    "        ax.axhline(self.threshold, linestyle=\"--\", color=\"grey\",\n",
+    "                   linewidth=1, label=f\"Promotion threshold ({self.threshold})\")\n",
+    "        for step, level in self._promotions:\n",
+    "            ax.axvline(step, color=\"green\", linewidth=1.5, linestyle=\":\")\n",
+    "            ax.text(step, ax.get_ylim()[1]*0.9, f\" L{level}\",\n",
+    "                    color=\"green\", fontsize=9)\n",
+    "        ax.set_xlabel(\"Episode\")\n",
+    "        ax.set_ylabel(\"Reward\")\n",
+    "        ax.set_title(\"Training Reward + Curriculum Progression\")\n",
+    "        ax.legend()\n",
+    "        plt.tight_layout()\n",
+    "        plt.savefig(\"curriculum_reward_curve.png\", dpi=150)\n",
+    "        plt.show()\n",
+    "        print(\"📈 Saved curriculum_reward_curve.png\")\n",
+    "\n",
+    "\n",
+    "# ── Demo: simulate 200 episodes of improving reward ──────────────────────────\n",
+    "import random\n",
+    "cb = CurriculumCallback(reward_threshold=5.0, window=50)\n",
+    "for ep in range(200):\n",
+    "    # Simulate reward slowly improving\n",
+    "    synthetic_reward = -5 + ep * 0.08 + random.gauss(0, 2)\n",
+    "    cb.record(synthetic_reward, step=ep)\n",
+    "\n",
+    "cb.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "inference-header",
+   "metadata": {},
+   "source": [
+    "## 🔍 Cell 11 — Inference: Before vs After Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "inference-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def run_inference_episode(model, tokenizer, engine, passes, program, max_steps=10):\n",
+    "    \"\"\"Run a greedy inference episode and return the action sequence + total improvement.\"\"\"\n",
+    "    from transformers import TextStreamer\n",
+    "    FastLanguageModel.for_inference(model)\n",
+    "\n",
+    "    env = CompilerOptimizationEnv(engine, passes, max_steps=max_steps)\n",
+    "    obs = env.reset(program)\n",
+    "    system_prompt = build_system_prompt(env)\n",
+    "\n",
+    "    actions_chosen, rewards_earned = [], []\n",
+    "    done = False\n",
+    "\n",
+    "    print(f\"\\nBaseline cycles: {env.previous_cycles}\")\n",
+    "    print(f\"Initial state:\\n{obs}\\n\")\n",
+    "\n",
+    "    while not done:\n",
+    "        messages = [\n",
+    "            {\"role\": \"system\", \"content\": system_prompt},\n",
+    "            {\"role\": \"user\",   \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n",
+    "        ]\n",
+    "        inputs = tokenizer.apply_chat_template(\n",
+    "            messages,\n",
+    "            tokenize=True,\n",
+    "            add_generation_prompt=True,\n",
+    "            return_tensors=\"pt\",\n",
+    "        ).to(model.device)\n",
+    "\n",
+    "        with torch.no_grad():\n",
+    "            outputs = model.generate(\n",
+    "                input_ids  = inputs,\n",
+    "                max_new_tokens = 16,\n",
+    "                temperature = 0.1,  # greedy-ish for inference\n",
+    "                do_sample   = True,\n",
+    "            )\n",
+    "\n",
+    "        raw_action = tokenizer.decode(\n",
+    "            outputs[0][inputs.shape[-1]:], skip_special_tokens=True\n",
+    "        ).strip()\n",
+    "        action = re.sub(r\"[^a-z0-9_]\", \"\", raw_action.lower().split()[0])\n",
+    "\n",
+    "        result = env.step(action)\n",
+    "        actions_chosen.append(action)\n",
+    "        rewards_earned.append(result.reward)\n",
+    "        obs  = result.observation\n",
+    "        done = result.done\n",
+    "\n",
+    "        print(f\"Step {len(actions_chosen)}: '{action}' → reward={result.reward:+.2f}\")\n",
+    "\n",
+    "    summary = env._episode_summary()\n",
+    "    print(f\"\\n{'─'*40}\")\n",
+    "    print(f\"Total improvement: {summary['total_improvement_pct']:.1f}%\")\n",
+    "    print(f\"Final cycles: {summary['final_cycles']} (was {summary['baseline_cycles']})\")\n",
+    "    return summary\n",
+    "\n",
+    "\n",
+    "# Uncomment after training:\n",
+    "# summary = run_inference_episode(\n",
+    "#     model, tokenizer, MockEngine(), MOCK_PASSES, SAMPLE_PROGRAM\n",
+    "# )\n",
+    "\n",
+    "print(\"✅ Inference cell ready. Uncomment the last block after training to run.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "tips-header",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## 📝 Notes & Tips\n",
+    "\n",
+    "| What | Why it matters |\n",
+    "|------|----------------|\n",
+    "| `num_generations=4` | GRPO needs K≥2 rollouts per prompt to compute group-relative advantage. K=4 balances diversity vs. compute. |\n",
+    "| `beta=0.04` | KL penalty keeping policy close to reference. Too low → mode collapse. Too high → no learning. |\n",
+    "| `max_new_tokens=16` | Actions are one word. This prevents wasted computation and keeps the model from adding explanations. |\n",
+    "| Dynamic crash penalty | Computed as `−2 × (100 × max_steps)` so it always dominates the best possible episode reward without washing out all other reward signal. |\n",
+    "| Terminal bonus | Rewards cumulative improvement, not just greedy single-step gains. Critical for long-horizon tasks. |\n",
+    "| Soft invalid-action termination | 3 consecutive invalid actions → end. Single mistakes don't kill the episode; the agent can recover. |\n",
+    "| `deepcopy` on all pass inputs | Role 3's passes mutate dicts in-place. Without this, `original_program` gets corrupted and the verifier fails spuriously. |\n",
+    "\n",
+    "**Next steps:**\n",
+    "- Replace `MockEngine` with Role 1's real engine\n",
+    "- Replace `MOCK_PASSES` with Role 3's real passes\n",
+    "- Push environment to HuggingFace Spaces: `openenv init && openenv deploy`\n",
+    "- Add W&B sweep to tune `beta`, `learning_rate`, `num_generations`"
+   ]
+  }
+ ]
+}"

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/DockerFile

@@ -0,0 +1,34 @@
+# Use a lightweight Python 3.10 image as the base
+FROM python:3.10-slim
+
+# Hugging Face Spaces requires a non-root user for security.
+# We create a user named 'user' with user ID 1000.
+RUN useradd -m -u 1000 user
+
+# Set environment variables to ensure Python output is logged immediately
+# and to add the local bin directory to the PATH for pip installations.
+ENV PATH="/home/user/.local/bin:$PATH" \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1
+
+# Switch away from root to the new user
+USER user
+
+# Set the working directory inside the container
+WORKDIR /home/user/app
+
+# Copy the requirements file first to leverage Docker cache layers
+COPY --chown=user requirements.txt .
+
+# Upgrade pip and install the dependencies defined by Role 2
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application files (app.py, openenv.yaml, engine.py, etc.)
+COPY --chown=user . .
+
+# Expose port 7860, which is the default port Hugging Face routes traffic to
+EXPOSE 7860
+
+# Start the OpenEnv server using the entrypoint defined in app.py
+CMD ["python", "app.py"]

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md

@@ -1 +1,9 @@
-# MetaHackathon2026Finals
+---
+title: Compiler Brain OpenEnv
+emoji: 🧠
+colorFrom: indigo
+colorTo: purple
+sdk: docker
+app_file: app.py
+pinned: false
+---

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/requirements.txt

@@ -0,0 +1,24 @@
+# OpenEnv SDK (Mandatory for the hackathon judging environment)
+openenv
+
+# Core RL Environment 
+gymnasium
+
+# Hugging Face Training Stack (Versions matched to your D3 notebook)
+trl==0.23.1
+transformers==4.57.1
+peft
+accelerate
+bitsandbytes
+
+# Unsloth for efficient 4-bit QLoRA training on T4 GPUs
+# Note: Unsloth often prefers being installed via their specific pip wheel or git, 
+# but this is standard for a requirements file.
+unsloth
+
+# PyTorch (Will default to standard compatible version if no index is specified)
+torch
+
+# Logging & Visualizations (For Day 2 judging criteria)
+wandb
+matplotlib

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/app.py

@@ -0,0 +1,182 @@
+import json
+import logging
+import traceback
+from typing import Any, Dict, List
+
+# OpenEnv SDK import
+from openenv import MCPEnvironment
+
+# ==============================================================================
+# ROLE 1 & 3 IMPORTS
+# TODO: Import the actual execution engine and generator from your teammates
+# ==============================================================================
+# from engine import execute_tac, verify_equivalence
+# from curriculum import generate_level_code
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("CompilerEnvServer")
+
+class CompilerEnv(MCPEnvironment):
+    """
+    The OpenEnv Server Wrapper for the Toy-IR Compiler Pass Optimizer.
+    Acts as the referee between the LLM client and Role 1's Execution Engine.
+    """
+    
+    def __init__(self):
+        super().__init__()
+        # State variables
+        self.raw_json_code = None
+        self.current_state_string = ""
+        self.initial_cycles = 0
+        self.current_step = 0
+        self.max_steps = 10 # Set a max step limit per episode
+        
+        # Cycle Weights (The Physics defined by Role 2)
+        self.cycle_weights = {
+            "ADD": 1,
+            "SUB": 1,
+            "MUL": 4,
+            "DIV": 10,
+            "MEM_LOAD": 20,
+            "STORE": 20
+        }
+
+    def reset(self) -> str:
+        """
+        Grabs unoptimized code, calculates baseline cycles, and translates the 
+        state for the LLM.
+        """
+        self.current_step = 0
+        
+        # 1. Grab new unoptimized code (Role 3 integration)
+        # TODO: Replace with real generator: self.raw_json_code = generate_level_code()
+        self.raw_json_code = self._mock_generator() 
+        
+        # 2. Get baseline cycles (Role 1 integration)
+        # TODO: Replace with real engine: self.initial_cycles, _ = execute_tac(self.raw_json_code, [])
+        self.initial_cycles = 100 
+        
+        # 3. Translate to Pseudo-Assembly to prevent Attention Dilution
+        self.current_state_string = self._translate_state(self.raw_json_code)
+        
+        logger.info(f"Environment Reset. Baseline Cycles: {self.initial_cycles}")
+        return self.state()
+
+    def step(self, action: str) -> Dict[str, Any]:
+        """
+        Executes the LLM's chosen optimization pass, verifies math equivalence, 
+        and calculates the reward.
+        """
+        self.current_step += 1
+        
+        # 1. Parse LLM Action (Regex/JSON robustness)
+        try:
+            # Assuming the LLM outputs a single pass name or a list of passes
+            action_data = json.loads(action)
+            if isinstance(action_data, str):
+                action_array = [action_data]
+            else:
+                action_array = action_data
+            format_bonus = 0.1
+        except json.JSONDecodeError:
+            # Format Trap
+            return self._build_step_response(
+                reward=-2.5, 
+                done=True, 
+                error="Invalid JSON. You must output a valid JSON array of strings."
+            )
+
+        # 2. Execute Code & Verify Equivalence (Role 1 Integration)
+        # TODO: new_cycles, optimized_code = execute_tac(self.raw_json_code, action_array)
+        # TODO: is_valid = verify_equivalence(self.raw_json_code, optimized_code)
+        new_cycles = 80 # Mock Data
+        is_valid = True # Mock Data
+        
+        # 3. Calculate Reward Physics
+        if not is_valid:
+            # Correctness Penalty + Micro-Variance for GRPO
+            penalty = -2.0 - (len(action_array) * 0.01)
+            return self._build_step_response(
+                reward=penalty, 
+                done=True, 
+                error=f"Code equivalence broken by passes: {action_array}"
+            )
+
+        # Calculate Improvement Ratio + Time Tax (-1.0)
+        cycle_improvement_ratio = (self.initial_cycles - new_cycles) / self.initial_cycles
+        time_tax = -0.05 * self.current_step # Small tax to prevent pass spamming
+        reward = cycle_improvement_ratio + format_bonus + time_tax
+        
+        # Update state if sequential, or finish if one-shot
+        # NOTE: For hackathon speed, we treat this as a One-Shot episode
+        done = True 
+        
+        return self._build_step_response(
+            reward=reward, 
+            done=done, 
+            info={"status": "success", "optimized_cycles": new_cycles}
+        )
+
+    def state(self) -> str:
+        """
+        Returns the current observation to the LLM.
+        """
+        return f"Current Step: {self.current_step}/{self.max_steps}\n\n{self.current_state_string}"
+
+    def _translate_state(self, raw_json: List[Dict]) -> str:
+        """
+        Translates raw AST JSON into clean pseudo-assembly.
+        Strips all UUIDs and AST metadata.
+        """
+        pseudo_assembly = []
+        instruction_count = 1
+        
+        for inst in raw_json:
+            op = inst.get("op", "UNKNOWN")
+            src1 = inst.get("src1", "")
+            src2 = inst.get("src2", "")
+            dest = inst.get("dest", "")
+            
+            # Format arguments cleanly
+            args = f"{src1}" if src2 is None else f"{src1}, {src2}"
+            
+            if dest:
+                line = f"{instruction_count}. {dest} = {op} {args}"
+            else:
+                line = f"{instruction_count}. {op} {args}"
+                
+            pseudo_assembly.append(line)
+            instruction_count += 1
+            
+        return "\n".join(pseudo_assembly)
+
+    def _build_step_response(self, reward: float, done: bool, error: str = None, info: dict = None) -> Dict[str, Any]:
+        """Helper to format the standard OpenEnv step return dictionary."""
+        response = {
+            "reward": reward,
+            "done": done,
+            "state": self.state()
+        }
+        if error:
+            response["error"] = error
+        if info:
+            response["info"] = info
+        return response
+
+    def _mock_generator(self):
+        """Mock data so the server runs before Role 1 integrates their engine."""
+        return [
+            {"op": "CONST", "dest": "a", "src1": 2, "src2": None},
+            {"op": "CONST", "dest": "b", "src1": 3, "src2": None},
+            {"op": "ADD", "dest": "c", "src1": "a", "src2": "b"}
+        ]
+
+if __name__ == "__main__":
+    logger.info("Initializing CompilerEnv Server...")
+    try:
+        env = CompilerEnv()
+        # openenv.run() or start() depending on the specific MCP wrapper version
+        env.start() 
+    except Exception as e:
+        logger.error(f"Failed to start environment server: {e}")
+        logger.error(traceback.format_exc())

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/compiler_tetris (1).ipynb

@@ -0,0 +1,651 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "b7fdb9ab",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== CELL 1: Setup ====\n",
+        "!pip install -q pytest\n",
+        "import os, sys\n",
+        "\n",
+        "LOCAL_DIR = \"/content/toyir\"\n",
+        "DRIVE_DIR = \"/content/drive/MyDrive/env/toyir\"\n",
+        "\n",
+        "os.makedirs(LOCAL_DIR, exist_ok=True)\n",
+        "sys.path.insert(0, LOCAL_DIR)\n",
+        "print(\"Workspace ready at\", LOCAL_DIR)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "d0337367",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== CELL 1b: Mount Drive + sync workspace ====\n",
+        "from google.colab import drive\n",
+        "drive.mount('/content/drive')\n",
+        "\n",
+        "import shutil\n",
+        "os.makedirs(DRIVE_DIR, exist_ok=True)\n",
+        "\n",
+        "# On fresh runtime: pull existing files from Drive -> local\n",
+        "for f in os.listdir(DRIVE_DIR):\n",
+        "    src = os.path.join(DRIVE_DIR, f)\n",
+        "    if os.path.isfile(src) and f.endswith(\".py\"):\n",
+        "        shutil.copy2(src, os.path.join(LOCAL_DIR, f))\n",
+        "print(\"Synced from Drive:\", sorted(os.listdir(LOCAL_DIR)))\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "2c83c529",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%%writefile /content/toyir/toy_vm.py\n",
+        "\"\"\"Toy-IR VM: deterministic, dependency-light execution engine for RLVR.\"\"\"\n",
+        "from __future__ import annotations\n",
+        "from copy import deepcopy\n",
+        "from typing import Any\n",
+        "\n",
+        "CYCLE_COST: dict[str, int] = {\n",
+        "    \"CONST\": 1, \"ADD\": 1, \"SUB\": 1, \"MUL\": 3, \"DIV\": 5,\n",
+        "    \"COPY\": 1, \"LOAD\": 4, \"STORE\": 4, \"NOP\": 0,\n",
+        "}\n",
+        "OP_ORDER = (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"COPY\", \"LOAD\", \"STORE\", \"NOP\")\n",
+        "VALID_OPS = frozenset(CYCLE_COST)\n",
+        "REQUIRED_KEYS = (\"op\", \"dest\", \"src1\", \"src2\")\n",
+        "\n",
+        "\n",
+        "def _is_str_or_int_or_none(x: Any) -> bool:\n",
+        "    return x is None or isinstance(x, str) or (isinstance(x, int) and not isinstance(x, bool))\n",
+        "\n",
+        "\n",
+        "def _is_str_or_none(x: Any) -> bool:\n",
+        "    return x is None or isinstance(x, str)\n",
+        "\n",
+        "\n",
+        "def validate_ir(program: list[dict]) -> dict:\n",
+        "    errors: list[str] = []\n",
+        "    for i, ins in enumerate(program):\n",
+        "        if not isinstance(ins, dict):\n",
+        "            errors.append(f\"[{i}] not a dict\")\n",
+        "            continue\n",
+        "        for k in REQUIRED_KEYS:\n",
+        "            if k not in ins:\n",
+        "                errors.append(f\"[{i}] missing key '{k}'\")\n",
+        "        if \"op\" not in ins:\n",
+        "            continue\n",
+        "        op = ins[\"op\"]\n",
+        "        if op not in VALID_OPS:\n",
+        "            errors.append(f\"[{i}] invalid op '{op}'\")\n",
+        "            continue\n",
+        "        if not _is_str_or_none(ins.get(\"dest\")):\n",
+        "            errors.append(f\"[{i}] dest must be str|None\")\n",
+        "        if not _is_str_or_int_or_none(ins.get(\"src1\")):\n",
+        "            errors.append(f\"[{i}] src1 must be str|int|None\")\n",
+        "        if not _is_str_or_int_or_none(ins.get(\"src2\")):\n",
+        "            errors.append(f\"[{i}] src2 must be str|int|None\")\n",
+        "\n",
+        "        # per-op contracts\n",
+        "        if op in (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"COPY\", \"LOAD\"):\n",
+        "            if not isinstance(ins.get(\"dest\"), str):\n",
+        "                errors.append(f\"[{i}] {op} requires str dest\")\n",
+        "        if op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n",
+        "            if ins.get(\"src1\") is None or ins.get(\"src2\") is None:\n",
+        "                errors.append(f\"[{i}] {op} requires src1 and src2\")\n",
+        "        if op in (\"CONST\", \"COPY\"):\n",
+        "            if ins.get(\"src1\") is None:\n",
+        "                errors.append(f\"[{i}] {op} requires src1\")\n",
+        "        if op in (\"CONST\", \"COPY\", \"LOAD\", \"STORE\"):\n",
+        "            if ins.get(\"src2\") is not None:\n",
+        "                errors.append(f\"[{i}] {op} requires src2=None\")\n",
+        "        if op == \"LOAD\":\n",
+        "            if not isinstance(ins.get(\"src1\"), str):\n",
+        "                errors.append(f\"[{i}] LOAD src1 must be str (variable name)\")\n",
+        "        if op == \"STORE\":\n",
+        "            if not isinstance(ins.get(\"dest\"), str):\n",
+        "                errors.append(f\"[{i}] STORE dest must be str (variable name)\")\n",
+        "            if ins.get(\"src1\") is None:\n",
+        "                errors.append(f\"[{i}] STORE requires src1\")\n",
+        "        if op == \"NOP\":\n",
+        "            if any(ins.get(k) is not None for k in (\"dest\", \"src1\", \"src2\")):\n",
+        "                errors.append(f\"[{i}] NOP fields must all be None\")\n",
+        "\n",
+        "    return {\"valid\": len(errors) == 0, \"errors\": errors}\n",
+        "\n",
+        "\n",
+        "def _resolve(operand: Any, variables: dict) -> int:\n",
+        "    if isinstance(operand, str):\n",
+        "        if operand not in variables:\n",
+        "            raise KeyError(f\"undefined variable '{operand}'\")\n",
+        "        return variables[operand]\n",
+        "    if isinstance(operand, int) and not isinstance(operand, bool):\n",
+        "        return operand\n",
+        "    raise TypeError(f\"bad operand: {operand!r}\")\n",
+        "\n",
+        "\n",
+        "def _exec_one(ins: dict, variables: dict, memory: dict) -> int:\n",
+        "    op = ins[\"op\"]\n",
+        "    dest = ins[\"dest\"]\n",
+        "    s1, s2 = ins[\"src1\"], ins[\"src2\"]\n",
+        "\n",
+        "    if op == \"CONST\":\n",
+        "        variables[dest] = _resolve(s1, variables)\n",
+        "    elif op == \"ADD\":\n",
+        "        variables[dest] = _resolve(s1, variables) + _resolve(s2, variables)\n",
+        "    elif op == \"SUB\":\n",
+        "        variables[dest] = _resolve(s1, variables) - _resolve(s2, variables)\n",
+        "    elif op == \"MUL\":\n",
+        "        variables[dest] = _resolve(s1, variables) * _resolve(s2, variables)\n",
+        "    elif op == \"DIV\":\n",
+        "        b = _resolve(s2, variables)\n",
+        "        if b == 0:\n",
+        "            raise ZeroDivisionError(\"DIV by zero\")\n",
+        "        variables[dest] = _resolve(s1, variables) // b\n",
+        "    elif op == \"COPY\":\n",
+        "        variables[dest] = _resolve(s1, variables)\n",
+        "    elif op == \"LOAD\":\n",
+        "        if not isinstance(s1, str):\n",
+        "            raise TypeError(\"LOAD src1 must be a variable name\")\n",
+        "        addr = variables.get(s1)\n",
+        "        if not (isinstance(addr, int) and not isinstance(addr, bool)):\n",
+        "            raise TypeError(f\"LOAD address from '{s1}' is not int: {addr!r}\")\n",
+        "        if addr not in memory:\n",
+        "            raise KeyError(f\"LOAD from uninitialized addr {addr}\")\n",
+        "        variables[dest] = memory[addr]\n",
+        "    elif op == \"STORE\":\n",
+        "        if not isinstance(dest, str):\n",
+        "            raise TypeError(\"STORE dest must be a variable name\")\n",
+        "        addr = variables.get(dest)\n",
+        "        if not (isinstance(addr, int) and not isinstance(addr, bool)):\n",
+        "            raise TypeError(f\"STORE address from '{dest}' is not int: {addr!r}\")\n",
+        "        memory[addr] = _resolve(s1, variables)\n",
+        "    elif op == \"NOP\":\n",
+        "        pass\n",
+        "    else:\n",
+        "        raise ValueError(f\"non-executable op {op}\")\n",
+        "    return CYCLE_COST[op]\n",
+        "\n",
+        "\n",
+        "def execute(program: list[dict], initial_state: dict) -> dict:\n",
+        "    variables: dict = {}\n",
+        "    memory: dict = {}\n",
+        "    cycles = 0\n",
+        "    try:\n",
+        "        if not isinstance(initial_state, dict):\n",
+        "            raise TypeError(\"initial_state must be a dict\")\n",
+        "        state = deepcopy(initial_state)\n",
+        "        variables = state.get(\"variables\", {})\n",
+        "        memory = state.get(\"memory\", {})\n",
+        "        if not isinstance(variables, dict):\n",
+        "            raise TypeError(\"initial_state['variables'] must be a dict\")\n",
+        "        if not isinstance(memory, dict):\n",
+        "            raise TypeError(\"initial_state['memory'] must be a dict\")\n",
+        "\n",
+        "        for ins in program:\n",
+        "            cycles += _exec_one(ins, variables, memory)\n",
+        "    except Exception as e:\n",
+        "        return {\"memory\": memory, \"variables\": variables,\n",
+        "                \"cycles\": None, \"success\": False, \"error\": str(e)}\n",
+        "    return {\"memory\": memory, \"variables\": variables,\n",
+        "            \"cycles\": cycles, \"success\": True, \"error\": None}\n",
+        "\n",
+        "\n",
+        "def count_cycles(program: list[dict]) -> int:\n",
+        "    total = 0\n",
+        "    for i, ins in enumerate(program):\n",
+        "        op = ins.get(\"op\") if isinstance(ins, dict) else None\n",
+        "        if op not in CYCLE_COST:\n",
+        "            raise ValueError(f\"count_cycles: invalid or missing op at index {i}: {op!r}\")\n",
+        "        total += CYCLE_COST[op]\n",
+        "    return total\n",
+        "\n",
+        "\n",
+        "def profile(program: list[dict]) -> dict:\n",
+        "    out: dict[str, int] = {\"n_instructions\": len(program)}\n",
+        "    for op in OP_ORDER:\n",
+        "        out[f\"n_{op.lower()}\"] = 0\n",
+        "    for ins in program:\n",
+        "        key = f\"n_{ins['op'].lower()}\"\n",
+        "        if key in out:\n",
+        "            out[key] += 1\n",
+        "    return out\n",
+        "\n",
+        "\n",
+        "def _fmt_operand(x: Any) -> str:\n",
+        "    if x is None:\n",
+        "        return \"\"\n",
+        "    return str(x)\n",
+        "\n",
+        "\n",
+        "def dump_ir(program: list[dict]) -> str:\n",
+        "    stores = [ins[\"dest\"] for ins in program if ins[\"op\"] == \"STORE\"]\n",
+        "    seen, ordered = set(), []\n",
+        "    for s in stores:\n",
+        "        if s not in seen:\n",
+        "            seen.add(s); ordered.append(s)\n",
+        "    header = \"// OBSERVABLE OUT: \" + (\", \".join(f\"mem[{s}]\" for s in ordered) if ordered else \"(none)\")\n",
+        "    lines = [header]\n",
+        "    for ins in program:\n",
+        "        op = ins[\"op\"]\n",
+        "        dest, s1, s2 = ins[\"dest\"], ins[\"src1\"], ins[\"src2\"]\n",
+        "        if op == \"STORE\":\n",
+        "            lines.append(f\"STORE [{dest}] {_fmt_operand(s1)}\")\n",
+        "        elif op == \"LOAD\":\n",
+        "            lines.append(f\"{dest} = LOAD [{_fmt_operand(s1)}]\")\n",
+        "        elif op == \"NOP\":\n",
+        "            lines.append(\"NOP\")\n",
+        "        else:\n",
+        "            parts = [op, _fmt_operand(s1)]\n",
+        "            if s2 is not None:\n",
+        "                parts.append(_fmt_operand(s2))\n",
+        "            lines.append(f\"{dest} = {' '.join(p for p in parts if p)}\")\n",
+        "    return \"\\n\".join(lines)\n",
+        "\n",
+        "\n",
+        "def clean_nops(program: list[dict]) -> list[dict]:\n",
+        "    return [deepcopy(ins) for ins in program if ins[\"op\"] != \"NOP\"]\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "3c85a58b",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%%writefile /content/toyir/verifier.py\n",
+        "\"\"\"Strict equivalence verifier - anti-cheat for RL reward hacking.\"\"\"\n",
+        "from __future__ import annotations\n",
+        "from toy_vm import execute\n",
+        "\n",
+        "\n",
+        "def verify(original_program: list[dict],\n",
+        "           optimized_program: list[dict],\n",
+        "           initial_states: list[dict]) -> dict:\n",
+        "    if not initial_states:\n",
+        "        raise ValueError(\"verify requires at least one initial_state\")\n",
+        "\n",
+        "    results = []\n",
+        "    all_match = True\n",
+        "    for st in initial_states:\n",
+        "        r_o = execute(original_program, st)\n",
+        "        r_p = execute(optimized_program, st)\n",
+        "        mem_o = r_o[\"memory\"] if r_o[\"success\"] else None\n",
+        "        mem_p = r_p[\"memory\"] if r_p[\"success\"] else None\n",
+        "\n",
+        "        execution_errors = {}\n",
+        "        if not r_o[\"success\"]:\n",
+        "            execution_errors[\"original_error\"] = r_o[\"error\"]\n",
+        "        if not r_p[\"success\"]:\n",
+        "            execution_errors[\"optimized_error\"] = r_p[\"error\"]\n",
+        "\n",
+        "        if execution_errors:\n",
+        "            match = False\n",
+        "            mismatches = [\"__execution_error__\"]\n",
+        "        else:\n",
+        "            keys = set(mem_o.keys()) | set(mem_p.keys())\n",
+        "            mismatches = [k for k in keys if mem_o.get(k) != mem_p.get(k)]\n",
+        "            match = len(mismatches) == 0\n",
+        "\n",
+        "        if not match:\n",
+        "            all_match = False\n",
+        "        results.append({\n",
+        "            \"initial_state\": st,\n",
+        "            \"original_memory\": mem_o,\n",
+        "            \"optimized_memory\": mem_p,\n",
+        "            \"match\": match,\n",
+        "            \"mismatches\": sorted(mismatches, key=lambda x: (isinstance(x, str), x)),\n",
+        "            \"execution_errors\": execution_errors or None,\n",
+        "        })\n",
+        "    return {\"equivalent\": all_match, \"results\": results}\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "cd7f8fa2",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%%writefile /content/toyir/reward_utils.py\n",
+        "\"\"\"Reward signal: cycle-reduction metric for RL agent.\"\"\"\n",
+        "from __future__ import annotations\n",
+        "from toy_vm import execute\n",
+        "\n",
+        "\n",
+        "def compute_reduction(original_program: list[dict],\n",
+        "                      optimized_program: list[dict],\n",
+        "                      initial_state: dict) -> dict:\n",
+        "    r_o = execute(original_program, initial_state)\n",
+        "    r_p = execute(optimized_program, initial_state)\n",
+        "\n",
+        "    if not (r_o[\"success\"] and r_p[\"success\"]):\n",
+        "        return {\n",
+        "            \"original_cycles\": r_o[\"cycles\"],\n",
+        "            \"optimized_cycles\": r_p[\"cycles\"],\n",
+        "            \"absolute_savings\": None,\n",
+        "            \"relative_savings\": None,\n",
+        "            \"success\": False,\n",
+        "            \"error\": r_o[\"error\"] or r_p[\"error\"],\n",
+        "        }\n",
+        "\n",
+        "    oc, pc = r_o[\"cycles\"], r_p[\"cycles\"]\n",
+        "    rel = 0.0 if oc == 0 else (oc - pc) / oc\n",
+        "    rel = max(-1.0, min(1.0, rel))\n",
+        "    return {\n",
+        "        \"original_cycles\": oc,\n",
+        "        \"optimized_cycles\": pc,\n",
+        "        \"absolute_savings\": oc - pc,\n",
+        "        \"relative_savings\": rel,\n",
+        "        \"success\": True,\n",
+        "        \"error\": None,\n",
+        "    }\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "4f2a941c",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%%writefile /content/toyir/test_vm.py\n",
+        "\"\"\"Pytest suite covering VM, verifier, reward, and RL exploit guardrails.\"\"\"\n",
+        "from __future__ import annotations\n",
+        "import random\n",
+        "import pytest\n",
+        "from toy_vm import (\n",
+        "    validate_ir, execute, count_cycles, profile, dump_ir, clean_nops, OP_ORDER,\n",
+        ")\n",
+        "from verifier import verify\n",
+        "from reward_utils import compute_reduction\n",
+        "\n",
+        "\n",
+        "def I(op, dest=None, src1=None, src2=None):\n",
+        "    return {\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2}\n",
+        "\n",
+        "\n",
+        "def test_const_fold_equivalence():\n",
+        "    orig = [I(\"CONST\", \"t0\", 5), I(\"CONST\", \"t1\", 7), I(\"ADD\", \"t2\", \"t0\", \"t1\"),\n",
+        "            I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"t2\")]\n",
+        "    opt = [I(\"CONST\", \"t2\", 12), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"t2\")]\n",
+        "    res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"]\n",
+        "    red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"absolute_savings\"] > 0\n",
+        "\n",
+        "\n",
+        "def test_dead_code_invisible():\n",
+        "    orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 3),\n",
+        "            I(\"MUL\", \"dead\", \"x\", 99), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 3), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"]\n",
+        "\n",
+        "\n",
+        "def test_fold_then_dce():\n",
+        "    orig = [I(\"CONST\", \"a\", 2), I(\"CONST\", \"b\", 3), I(\"ADD\", \"c\", \"a\", \"b\"),\n",
+        "            I(\"MUL\", \"d\", \"c\", 10), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"d\")]\n",
+        "    opt = [I(\"CONST\", \"d\", 50), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"d\")]\n",
+        "    res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"]\n",
+        "    red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"original_cycles\"] > red[\"optimized_cycles\"]\n",
+        "\n",
+        "\n",
+        "def test_strength_reduction():\n",
+        "    orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 7),\n",
+        "            I(\"MUL\", \"y\", \"x\", 2), I(\"STORE\", \"addr\", \"y\")]\n",
+        "    opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 7),\n",
+        "           I(\"ADD\", \"y\", \"x\", \"x\"), I(\"STORE\", \"addr\", \"y\")]\n",
+        "    res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"]\n",
+        "    red = compute_reduction(orig, opt, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"absolute_savings\"] > 0\n",
+        "\n",
+        "\n",
+        "def test_already_optimal():\n",
+        "    p = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    red = compute_reduction(p, p, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"absolute_savings\"] == 0\n",
+        "    assert red[\"relative_savings\"] == 0.0\n",
+        "\n",
+        "\n",
+        "def test_div_by_zero():\n",
+        "    p = [I(\"CONST\", \"a\", 10), I(\"CONST\", \"b\", 0), I(\"DIV\", \"c\", \"a\", \"b\")]\n",
+        "    r = execute(p, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert r[\"success\"] is False\n",
+        "    assert \"zero\" in r[\"error\"].lower()\n",
+        "\n",
+        "\n",
+        "def test_hardcoded_exploit_caught():\n",
+        "    orig = [I(\"CONST\", \"addr\", 0), I(\"LOAD\", \"x\", \"in_addr\"),\n",
+        "            I(\"MUL\", \"y\", \"x\", 3), I(\"STORE\", \"addr\", \"y\")]\n",
+        "    opt = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"y\", 42), I(\"STORE\", \"addr\", \"y\")]\n",
+        "    rng = random.Random(0)\n",
+        "    states = [{\"variables\": {\"in_addr\": 1}, \"memory\": {1: rng.randint(0, 100)}} for _ in range(5)]\n",
+        "    res = verify(orig, opt, states)\n",
+        "    assert res[\"equivalent\"] is False\n",
+        "    assert any(not r[\"match\"] for r in res[\"results\"])\n",
+        "\n",
+        "\n",
+        "def test_nop_zero_cost():\n",
+        "    p = [I(\"NOP\"), I(\"CONST\", \"x\", 1), I(\"NOP\")]\n",
+        "    assert count_cycles(p) == 1\n",
+        "    r = execute(p, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert r[\"success\"] and r[\"cycles\"] == 1\n",
+        "\n",
+        "\n",
+        "def test_load_store_non_int_address():\n",
+        "    p = [I(\"COPY\", \"addr\", \"bad\"), I(\"LOAD\", \"x\", \"addr\")]\n",
+        "    r = execute(p, {\"variables\": {\"bad\": \"notint\"}, \"memory\": {}})\n",
+        "    assert r[\"success\"] is False\n",
+        "\n",
+        "\n",
+        "def test_validate_ir_malformed():\n",
+        "    assert not validate_ir([I(\"FOOBAR\", \"t\", 1, 2)])[\"valid\"]\n",
+        "    assert not validate_ir([{\"op\": \"ADD\", \"dest\": \"x\"}])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"CONST\", 123, 1)])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"STORE\", None, \"x\")])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"LOAD\", \"x\", 42)])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"NOP\", \"x\", 1)])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"ADD\", \"x\", \"a\", None)])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"LOAD\", \"x\", \"addr\", 9)])[\"valid\"]\n",
+        "    assert not validate_ir([I(\"STORE\", \"addr\", \"x\", 9)])[\"valid\"]\n",
+        "    assert validate_ir([I(\"CONST\", \"x\", 1)])[\"valid\"]\n",
+        "\n",
+        "\n",
+        "def test_dump_ir_format():\n",
+        "    p = [I(\"CONST\", \"t0\", 5), I(\"ADD\", \"t1\", \"t0\", \"a\"),\n",
+        "         I(\"MUL\", \"t2\", \"t1\", 3), I(\"STORE\", \"addr0\", \"t2\")]\n",
+        "    s = dump_ir(p)\n",
+        "    lines = s.split(\"\\n\")\n",
+        "    assert lines[0] == \"// OBSERVABLE OUT: mem[addr0]\"\n",
+        "    assert lines[1] == \"t0 = CONST 5\"\n",
+        "    assert lines[2] == \"t1 = ADD t0 a\"\n",
+        "    assert lines[3] == \"t2 = MUL t1 3\"\n",
+        "    assert lines[4] == \"STORE [addr0] t2\"\n",
+        "    assert len(lines) == 5\n",
+        "\n",
+        "\n",
+        "def test_dump_ir_multi_store_header_order():\n",
+        "    p = [I(\"CONST\", \"a\", 0), I(\"CONST\", \"b\", 1),\n",
+        "         I(\"STORE\", \"a\", 5), I(\"STORE\", \"b\", 6), I(\"STORE\", \"a\", 7)]\n",
+        "    s = dump_ir(p)\n",
+        "    assert s.split(\"\\n\")[0] == \"// OBSERVABLE OUT: mem[a], mem[b]\"\n",
+        "\n",
+        "\n",
+        "def test_clean_nops():\n",
+        "    p = [I(\"NOP\"), I(\"CONST\", \"x\", 1), I(\"NOP\"), I(\"CONST\", \"y\", 2), I(\"NOP\")]\n",
+        "    out = clean_nops(p)\n",
+        "    assert len(out) == 2\n",
+        "    assert all(ins[\"op\"] != \"NOP\" for ins in out)\n",
+        "\n",
+        "\n",
+        "def test_verify_mismatches():\n",
+        "    orig = [I(\"CONST\", \"a0\", 0), I(\"CONST\", \"a1\", 1),\n",
+        "            I(\"CONST\", \"x\", 5), I(\"CONST\", \"y\", 9),\n",
+        "            I(\"STORE\", \"a0\", \"x\"), I(\"STORE\", \"a1\", \"y\")]\n",
+        "    opt = [I(\"CONST\", \"a0\", 0), I(\"CONST\", \"a1\", 1),\n",
+        "           I(\"CONST\", \"x\", 5), I(\"CONST\", \"y\", 8),\n",
+        "           I(\"STORE\", \"a0\", \"x\"), I(\"STORE\", \"a1\", \"y\")]\n",
+        "    res = verify(orig, opt, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"] is False\n",
+        "    assert res[\"results\"][0][\"mismatches\"] == [1]\n",
+        "\n",
+        "\n",
+        "def test_profile_counts_and_ordering():\n",
+        "    p = [I(\"CONST\", \"x\", 1), I(\"ADD\", \"y\", \"x\", 2), I(\"MUL\", \"z\", \"y\", 3)]\n",
+        "    pr = profile(p)\n",
+        "    assert pr[\"n_instructions\"] == 3\n",
+        "    assert pr[\"n_const\"] == 1 and pr[\"n_add\"] == 1 and pr[\"n_mul\"] == 1\n",
+        "    expected = [\"n_instructions\"] + [f\"n_{op.lower()}\" for op in OP_ORDER]\n",
+        "    assert list(pr.keys()) == expected\n",
+        "\n",
+        "\n",
+        "def test_no_input_mutation():\n",
+        "    p = [I(\"CONST\", \"x\", 1), I(\"CONST\", \"addr\", 0), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    st = {\"variables\": {}, \"memory\": {}}\n",
+        "    snap = {\"variables\": dict(st[\"variables\"]), \"memory\": dict(st[\"memory\"])}\n",
+        "    execute(p, st)\n",
+        "    assert st == snap\n",
+        "\n",
+        "\n",
+        "def test_clean_nops_no_aliasing():\n",
+        "    p = [I(\"CONST\", \"x\", 1), I(\"NOP\")]\n",
+        "    out = clean_nops(p)\n",
+        "    out[0][\"op\"] = \"MUL\"\n",
+        "    assert p[0][\"op\"] == \"CONST\"\n",
+        "\n",
+        "\n",
+        "def test_verify_rejects_empty_states():\n",
+        "    with pytest.raises(ValueError):\n",
+        "        verify([I(\"NOP\")], [I(\"NOP\")], [])\n",
+        "\n",
+        "\n",
+        "def test_verify_surfaces_execution_errors():\n",
+        "    good = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    bad = [I(\"LOAD\", \"x\", \"missing_addr\")]\n",
+        "    res = verify(good, bad, [{\"variables\": {}, \"memory\": {}}])\n",
+        "    assert res[\"equivalent\"] is False\n",
+        "    assert res[\"results\"][0][\"mismatches\"] == [\"__execution_error__\"]\n",
+        "    assert res[\"results\"][0][\"execution_errors\"][\"optimized_error\"] is not None\n",
+        "    assert \"original_error\" not in res[\"results\"][0][\"execution_errors\"]\n",
+        "\n",
+        "\n",
+        "def test_execute_rejects_bad_initial_state_shape():\n",
+        "    r = execute([I(\"CONST\", \"x\", 1)], \"not-a-dict\")\n",
+        "    assert r[\"success\"] is False\n",
+        "    assert \"initial_state\" in r[\"error\"]\n",
+        "\n",
+        "\n",
+        "def test_compute_reduction_failure_path():\n",
+        "    orig = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    bad = [I(\"CONST\", \"a\", 1), I(\"CONST\", \"b\", 0), I(\"DIV\", \"c\", \"a\", \"b\")]\n",
+        "    red = compute_reduction(orig, bad, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"success\"] is False\n",
+        "    assert red[\"relative_savings\"] is None\n",
+        "    assert red[\"absolute_savings\"] is None\n",
+        "    assert \"zero\" in red[\"error\"].lower()\n",
+        "\n",
+        "\n",
+        "def test_relative_savings_clamped_on_slowdown():\n",
+        "    fast = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1), I(\"STORE\", \"addr\", \"x\")]\n",
+        "    # slower variant: extra dead MULs balloon the cycle count\n",
+        "    slow = [I(\"CONST\", \"addr\", 0), I(\"CONST\", \"x\", 1)]\n",
+        "    slow += [I(\"MUL\", \"junk\", \"x\", 2)] * 50\n",
+        "    slow += [I(\"STORE\", \"addr\", \"x\")]\n",
+        "    red = compute_reduction(fast, slow, {\"variables\": {}, \"memory\": {}})\n",
+        "    assert red[\"absolute_savings\"] < 0\n",
+        "    assert red[\"relative_savings\"] == -1.0  # clamped\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "e6b83243",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== CELL 6: Push to Drive ====\n",
+        "import shutil, os\n",
+        "for f in os.listdir(LOCAL_DIR):\n",
+        "    if f.endswith(\".py\"):\n",
+        "        shutil.copy2(os.path.join(LOCAL_DIR, f), os.path.join(DRIVE_DIR, f))\n",
+        "print(\"Saved to Drive:\", sorted(os.listdir(DRIVE_DIR)))\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "5675f9a0",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== CELL 7: Run tests ====\n",
+        "import subprocess\n",
+        "res = subprocess.run(\n",
+        "    [\"python\", \"-m\", \"pytest\", \"-v\", \"test_vm.py\"],\n",
+        "    cwd=LOCAL_DIR, capture_output=True, text=True,\n",
+        ")\n",
+        "print(res.stdout)\n",
+        "print(res.stderr)\n",
+        "assert res.returncode == 0, \"Tests failed\"\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1a500bb3",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== CELL 8: Smoke demo ====\n",
+        "import importlib, toy_vm, verifier, reward_utils\n",
+        "importlib.reload(toy_vm); importlib.reload(verifier); importlib.reload(reward_utils)\n",
+        "from toy_vm import dump_ir\n",
+        "from reward_utils import compute_reduction\n",
+        "from verifier import verify\n",
+        "\n",
+        "orig = [\n",
+        "    {\"op\":\"CONST\",\"dest\":\"addr\",\"src1\":0,\"src2\":None},\n",
+        "    {\"op\":\"CONST\",\"dest\":\"a\",\"src1\":2,\"src2\":None},\n",
+        "    {\"op\":\"CONST\",\"dest\":\"b\",\"src1\":3,\"src2\":None},\n",
+        "    {\"op\":\"ADD\",\"dest\":\"c\",\"src1\":\"a\",\"src2\":\"b\"},\n",
+        "    {\"op\":\"MUL\",\"dest\":\"d\",\"src1\":\"c\",\"src2\":2},\n",
+        "    {\"op\":\"MUL\",\"dest\":\"dead\",\"src1\":\"d\",\"src2\":99},\n",
+        "    {\"op\":\"STORE\",\"dest\":\"addr\",\"src1\":\"d\",\"src2\":None},\n",
+        "]\n",
+        "opt = [\n",
+        "    {\"op\":\"CONST\",\"dest\":\"addr\",\"src1\":0,\"src2\":None},\n",
+        "    {\"op\":\"CONST\",\"dest\":\"d\",\"src1\":10,\"src2\":None},\n",
+        "    {\"op\":\"STORE\",\"dest\":\"addr\",\"src1\":\"d\",\"src2\":None},\n",
+        "]\n",
+        "\n",
+        "print(dump_ir(orig)); print(\"---\"); print(dump_ir(opt)); print(\"---\")\n",
+        "print(compute_reduction(orig, opt, {\"variables\":{},\"memory\":{}}))\n",
+        "print(\"equivalent:\", verify(orig, opt, [{\"variables\":{},\"memory\":{}}])[\"equivalent\"])\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1 (1).ipynb

@@ -0,0 +1,1379 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "## Design Assumptions (do not violate)\n",
+        "\n",
+        "# 1. **DCE never eliminates STOREs.** They define program output (final mem state).\n",
+        "# 2. **Addresses are distinct by construction.** Generator allocates each address variable to a unique integer; no aliasing.\n",
+        "# 3. **CF refuses to fold DIV by zero.** If src2 == 0 on a DIV op, leave instruction unchanged.\n",
+        "# 4. **Generator never emits DIV by literal zero.** When DIV is generated, src2 is always a non-zero constant or a variable known to be non-zero.\n",
+        "# # 5. **Integer arithmetic only.** No floats anywhere — avoids equivalence-check precision issues.\n",
+        "# 6. Generator declares `observable_addrs` per program — verifier compares only these mem entries.\n",
+        "# 7. State translator annotates observable outputs at top of dump.\n",
+        "# 8. Reward distinguishes broken (-1000) from valid-but-worse (small negative) — Harshal's formula.\n",
+        "# 9. Multi-input verification: 3-5 random initial states, all must match.\n",
+        "# 10. Integer division uses Python floor division (//). Aarush's VM must match."
+      ],
+      "metadata": {
+        "id": "_XI5jT2Ibvrf"
+      },
+      "execution_count": 2,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "id": "j6CQ327KWXwr"
+      },
+      "outputs": [],
+      "source": [
+        "# === TAC Schema v1.0 (LOCKED with Role 1 / Aarush) ===\n",
+        "# Reverse passes deferred to stretch goal — not in initial action space.\n",
+        "\n",
+        "OPS = [\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\", \"STORE\", \"NOP\"]\n",
+        "\n",
+        "CYCLE_COSTS = {\n",
+        "    \"CONST\": 1,\n",
+        "    \"ADD\":   1,\n",
+        "    \"SUB\":   1,\n",
+        "    \"MUL\":   3,\n",
+        "    \"DIV\":   5,\n",
+        "    \"LOAD\":  4,\n",
+        "    \"STORE\": 4,\n",
+        "    \"NOP\":   0,\n",
+        "}\n",
+        "\n",
+        "# Instruction shape: {\"op\": str, \"dest\": str|None, \"src1\": Any, \"src2\": Any}\n",
+        "# Operands: str = variable name, int = literal constant, None = unused\n",
+        "#\n",
+        "# Op semantics:\n",
+        "#   CONST: dest = src1                      (src1 is int literal, src2 = None)\n",
+        "#   ADD/SUB/MUL/DIV: dest = src1 OP src2    (src1, src2 are var names or int literals)\n",
+        "#   LOAD:  dest = mem[src1]                 (src1 is a var holding an address)\n",
+        "#   STORE: mem[dest] = src1                 (dest is a var holding an address)\n",
+        "#   NOP:   no-op                            (all fields None)\n",
+        "#\n",
+        "# Program output (for equivalence check) = final memory state (mem dict).\n",
+        "# Programs ship as: (initial_vars: dict, initial_mem: dict, instructions: list[dict])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import random\n",
+        "\n",
+        "def generate_level_1():\n",
+        "    \"\"\"\n",
+        "    Generate a Level 1 Toy-IR program.\n",
+        "\n",
+        "    Characteristics:\n",
+        "    - 4-6 instructions\n",
+        "    - 2-3 CONST ops with literal values\n",
+        "    - 1-2 arithmetic ops on those constants (foldable by CF)\n",
+        "    - 0-1 dead variables (killable by DCE)\n",
+        "    - Exactly 1 STORE at the end so the program has an observable output\n",
+        "\n",
+        "    Returns:\n",
+        "        dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n",
+        "    \"\"\"\n",
+        "    instructions = []\n",
+        "    var_counter = 0\n",
+        "\n",
+        "    def new_var():\n",
+        "        nonlocal var_counter\n",
+        "        name = f\"v{var_counter}\"\n",
+        "        var_counter += 1\n",
+        "        return name\n",
+        "\n",
+        "    # Step 1: Generate 2-3 constant assignments\n",
+        "    num_consts = random.randint(2, 3)\n",
+        "    const_vars = []\n",
+        "    for _ in range(num_consts):\n",
+        "        var = new_var()\n",
+        "        value = random.randint(1, 10)\n",
+        "        instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n",
+        "        const_vars.append(var)\n",
+        "\n",
+        "    # Step 2: Generate 1-2 arithmetic ops using those constants\n",
+        "    num_arith = random.randint(1, 2)\n",
+        "    last_result = None\n",
+        "    for _ in range(num_arith):\n",
+        "        op = random.choice([\"ADD\", \"MUL\"])\n",
+        "        src1 = random.choice(const_vars)\n",
+        "        src2 = random.choice(const_vars)\n",
+        "        dest = new_var()\n",
+        "        instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n",
+        "        last_result = dest\n",
+        "        const_vars.append(dest)\n",
+        "\n",
+        "    # Step 3: Optionally add 1 dead variable (50% chance)\n",
+        "    if random.random() < 0.5:\n",
+        "        dead_var = new_var()\n",
+        "        dead_value = random.randint(1, 10)\n",
+        "        instructions.append({\"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None})\n",
+        "        # Note: dead_var is intentionally never used — DCE should catch it.\n",
+        "\n",
+        "    # Step 4: Add a STORE at the end so the program has observable output\n",
+        "    initial_vars = {\"addr0\": 0}\n",
+        "    instructions.append({\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None})\n",
+        "\n",
+        "    initial_mem = {}\n",
+        "\n",
+        "    return {\n",
+        "        \"initial_vars\":     initial_vars,\n",
+        "        \"initial_mem\":      initial_mem,\n",
+        "        \"instructions\":     instructions,\n",
+        "        \"observable_addrs\": [0],   # Aarush's verifier compares only these mem entries\n",
+        "    }\n",
+        "\n",
+        "\n",
+        "# Sanity-check: generate a few programs and print them\n",
+        "for seed in [42, 1, 7, 99]:\n",
+        "    random.seed(seed)\n",
+        "    prog = generate_level_1()\n",
+        "    print(f\"\\n=== seed={seed} ===\")\n",
+        "    print(f\"initial_vars     : {prog['initial_vars']}\")\n",
+        "    print(f\"initial_mem      : {prog['initial_mem']}\")\n",
+        "    print(f\"observable_addrs : {prog['observable_addrs']}\")\n",
+        "    print(f\"instructions ({len(prog['instructions'])}):\")\n",
+        "    for i, instr in enumerate(prog['instructions']):\n",
+        "        print(f\"  {i}: {instr}\")"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "a86vD1OyWcB9",
+        "outputId": "2d6d4604-6dc1-4c29-af64-e41cf5717c0d"
+      },
+      "execution_count": 4,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\n",
+            "=== seed=42 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (4):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "\n",
+            "=== seed=1 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (5):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 10, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 2, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v1', 'src2': 'v1'}\n",
+            "  3: {'op': 'MUL', 'dest': 'v3', 'src1': 'v2', 'src2': 'v1'}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "\n",
+            "=== seed=7 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (6):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n",
+            "  3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n",
+            "  4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n",
+            "  5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "\n",
+            "=== seed=99 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (5):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 7, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 4, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 10, 'src2': None}\n",
+            "  3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def _resolve_operand(operand, known_constants):\n",
+        "    \"\"\"\n",
+        "    Given a TAC operand (string variable name or int literal),\n",
+        "    return its concrete integer value if known, or None if unknown.\n",
+        "    \"\"\"\n",
+        "    if isinstance(operand, int):\n",
+        "        return operand\n",
+        "    if isinstance(operand, str) and operand in known_constants:\n",
+        "        return known_constants[operand]\n",
+        "    return None\n",
+        "\n",
+        "\n",
+        "def _compute(op, v1, v2):\n",
+        "    \"\"\"Compute the result of a binary arithmetic op on two known integers.\"\"\"\n",
+        "    if op == \"ADD\": return v1 + v2\n",
+        "    if op == \"SUB\": return v1 - v2\n",
+        "    if op == \"MUL\": return v1 * v2\n",
+        "    if op == \"DIV\": return v1 // v2   # floor division (locked semantics)\n",
+        "    raise ValueError(f\"_compute called with non-arithmetic op: {op}\")\n",
+        "\n",
+        "\n",
+        "def constant_folding(program):\n",
+        "    \"\"\"\n",
+        "    Forward pass that folds constant arithmetic into CONST ops, and\n",
+        "    propagates known constants into instruction operands.\n",
+        "\n",
+        "    Behavior:\n",
+        "    - If both sources of an arithmetic op resolve to known integers,\n",
+        "      replaces the instruction with a CONST holding the computed result.\n",
+        "    - If only one source is known, still substitutes that known value\n",
+        "      into the instruction (constant propagation), enabling downstream\n",
+        "      passes (e.g., peephole) to recognize patterns like ADD-with-0 or MUL-by-1.\n",
+        "    - Refuses to fold DIV by zero (Design Assumption #3).\n",
+        "    - Always returns fresh dicts; never aliases input instructions.\n",
+        "\n",
+        "    Args:\n",
+        "        program: list of TAC instruction dicts (per locked schema).\n",
+        "\n",
+        "    Returns:\n",
+        "        new list of TAC instruction dicts. Always semantics-preserving.\n",
+        "        Never raises, never returns None.\n",
+        "    \"\"\"\n",
+        "    known_constants = {}\n",
+        "    new_program = []\n",
+        "\n",
+        "    for instr in program:\n",
+        "        # Always work on a copy — never alias input dicts\n",
+        "        instr = instr.copy()\n",
+        "        op   = instr[\"op\"]\n",
+        "        dest = instr[\"dest\"]\n",
+        "\n",
+        "        if op == \"CONST\":\n",
+        "            known_constants[dest] = instr[\"src1\"]\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        elif op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n",
+        "            # === Constant propagation: substitute known constants into operands ===\n",
+        "            if isinstance(instr[\"src1\"], str) and instr[\"src1\"] in known_constants:\n",
+        "                instr[\"src1\"] = known_constants[instr[\"src1\"]]\n",
+        "            if isinstance(instr[\"src2\"], str) and instr[\"src2\"] in known_constants:\n",
+        "                instr[\"src2\"] = known_constants[instr[\"src2\"]]\n",
+        "\n",
+        "            # === Try to fold ===\n",
+        "            v1 = _resolve_operand(instr[\"src1\"], known_constants)\n",
+        "            v2 = _resolve_operand(instr[\"src2\"], known_constants)\n",
+        "\n",
+        "            if v1 is not None and v2 is not None:\n",
+        "                # Both operands are known integers\n",
+        "                if op == \"DIV\" and v2 == 0:\n",
+        "                    # Refuse to fold DIV by zero\n",
+        "                    new_program.append(instr)\n",
+        "                    known_constants.pop(dest, None)\n",
+        "                else:\n",
+        "                    # Fold: replace with CONST\n",
+        "                    result = _compute(op, v1, v2)\n",
+        "                    new_program.append({\n",
+        "                        \"op\":   \"CONST\",\n",
+        "                        \"dest\": dest,\n",
+        "                        \"src1\": result,\n",
+        "                        \"src2\": None,\n",
+        "                    })\n",
+        "                    known_constants[dest] = result\n",
+        "            else:\n",
+        "                # Can't fold (at least one operand unknown).\n",
+        "                # Instruction may still have been mutated by propagation above.\n",
+        "                new_program.append(instr)\n",
+        "                known_constants.pop(dest, None)\n",
+        "\n",
+        "        elif op == \"LOAD\":\n",
+        "            # Memory reads aren't statically resolvable\n",
+        "            new_program.append(instr)\n",
+        "            known_constants.pop(dest, None)\n",
+        "\n",
+        "        elif op in (\"STORE\", \"NOP\"):\n",
+        "            # No dest tracking needed.\n",
+        "            # Note: we COULD propagate src1 into a STORE for downstream readability,\n",
+        "            # but the cycle cost is unchanged and the executor handles vars fine.\n",
+        "            # Leave STORE alone — keeps the code minimal.\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        else:\n",
+        "            # Unknown op — defensive pass-through\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "    return new_program"
+      ],
+      "metadata": {
+        "id": "CGk8Fz4CZA3t"
+      },
+      "execution_count": 5,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for constant_folding ===\n",
+        "\n",
+        "# Test 1: simple fold — ADD of two CONSTs\n",
+        "test1 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result1 = constant_folding(test1)\n",
+        "print(\"Test 1 (simple ADD fold):\")\n",
+        "for i, instr in enumerate(result1): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 becomes CONST c = 8\n",
+        "\n",
+        "# Test 2: chained fold — second op uses first op's folded result\n",
+        "test2 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},   # → 8\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"d\", \"src1\": \"c\", \"src2\": \"b\"},   # → 8 * 5 = 40\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"d\", \"src2\": None},\n",
+        "]\n",
+        "result2 = constant_folding(test2)\n",
+        "print(\"\\nTest 2 (chained fold):\")\n",
+        "for i, instr in enumerate(result2): print(f\"  {i}: {instr}\")\n",
+        "# Expected: c becomes CONST 8, d becomes CONST 40\n",
+        "\n",
+        "# Test 3: DIV by zero refused\n",
+        "test3 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 10, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0,  \"src2\": None},\n",
+        "    {\"op\": \"DIV\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},   # 10/0 — must NOT fold\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result3 = constant_folding(test3)\n",
+        "print(\"\\nTest 3 (DIV by zero refused):\")\n",
+        "for i, instr in enumerate(result3): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 unchanged (still DIV, not CONST)\n",
+        "\n",
+        "# Test 4: unknown source can't fold\n",
+        "test4 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None},   # b is unknown\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},        # can't fold\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result4 = constant_folding(test4)\n",
+        "print(\"\\nTest 4 (LOAD makes b unknown, ADD not folded):\")\n",
+        "for i, instr in enumerate(result4): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 unchanged (still ADD)\n",
+        "\n",
+        "# Test 5: idempotence — running CF on a generated program\n",
+        "random.seed(42)\n",
+        "prog = generate_level_1()\n",
+        "folded = constant_folding(prog[\"instructions\"])\n",
+        "print(\"\\nTest 5 (CF on generated Level 1 program, seed=42):\")\n",
+        "print(\"Before:\")\n",
+        "for i, instr in enumerate(prog[\"instructions\"]): print(f\"  {i}: {instr}\")\n",
+        "print(\"After:\")\n",
+        "for i, instr in enumerate(folded): print(f\"  {i}: {instr}\")\n",
+        "\n",
+        "# Test 6: constant propagation — only one source is known\n",
+        "test6 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 0, \"src2\": None},\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None},   # b unknown\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},        # a known (=0), b unknown\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result6 = constant_folding(test6)\n",
+        "print(\"\\nTest 6 (propagation: a=0 substituted into ADD even though b unknown):\")\n",
+        "for i, instr in enumerate(result6): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 is still ADD (can't fold — b unknown), but src1 is now literal 0, not 'a'\n",
+        "# This sets up peephole to recognize \"ADD with 0\" later"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "WLKSTVa8fRF8",
+        "outputId": "e9480e1e-5d59-4afe-a884-bae0fa4a2192"
+      },
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Test 1 (simple ADD fold):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 2 (chained fold):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n",
+            "  3: {'op': 'CONST', 'dest': 'd', 'src1': 40, 'src2': None}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'd', 'src2': None}\n",
+            "\n",
+            "Test 3 (DIV by zero refused):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 10, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 0, 'src2': None}\n",
+            "  2: {'op': 'DIV', 'dest': 'c', 'src1': 10, 'src2': 0}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 4 (LOAD makes b unknown, ADD not folded):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'c', 'src1': 3, 'src2': 'b'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 5 (CF on generated Level 1 program, seed=42):\n",
+            "Before:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "After:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "\n",
+            "Test 6 (propagation: a=0 substituted into ADD even though b unknown):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 0, 'src2': None}\n",
+            "  1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'c', 'src1': 0, 'src2': 'b'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def dead_code_elimination(program):\n",
+        "    \"\"\"\n",
+        "    Backward liveness analysis. Replace dead instructions with NOPs,\n",
+        "    then strip NOPs.\n",
+        "\n",
+        "    Rules:\n",
+        "    - STOREs are always live (define program output).\n",
+        "    - Address variables used in STORE/LOAD are always live.\n",
+        "    - Any instruction whose dest is never read later is dead.\n",
+        "\n",
+        "    Returns a fresh list of instruction dicts. Never raises, never returns None.\n",
+        "    \"\"\"\n",
+        "    # Walk backward, build new program in reverse, then re-reverse at the end\n",
+        "    live = set()\n",
+        "    new_program_reversed = []\n",
+        "\n",
+        "    for instr in reversed(program):\n",
+        "        instr = instr.copy()  # never alias input\n",
+        "        op   = instr[\"op\"]\n",
+        "        dest = instr[\"dest\"]\n",
+        "        src1 = instr[\"src1\"]\n",
+        "        src2 = instr[\"src2\"]\n",
+        "\n",
+        "        if op == \"STORE\":\n",
+        "            # STOREs are always kept; mark sources live\n",
+        "            if isinstance(src1, str): live.add(src1)\n",
+        "            if isinstance(dest, str): live.add(dest)  # address variable\n",
+        "            new_program_reversed.append(instr)\n",
+        "\n",
+        "        elif op == \"NOP\":\n",
+        "            # Pass through; will be stripped at the end\n",
+        "            new_program_reversed.append(instr)\n",
+        "\n",
+        "        elif op in (\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\"):\n",
+        "            if dest in live:\n",
+        "                # Live instruction — keep it, mark its sources live\n",
+        "                live.discard(dest)\n",
+        "                if isinstance(src1, str): live.add(src1)\n",
+        "                if isinstance(src2, str): live.add(src2)\n",
+        "                new_program_reversed.append(instr)\n",
+        "            else:\n",
+        "                # Dead — replace with NOP\n",
+        "                new_program_reversed.append({\n",
+        "                    \"op\": \"NOP\", \"dest\": None, \"src1\": None, \"src2\": None,\n",
+        "                })\n",
+        "\n",
+        "        else:\n",
+        "            # Unknown op — defensive pass-through\n",
+        "            new_program_reversed.append(instr)\n",
+        "\n",
+        "    # Reverse back to forward order, then strip NOPs\n",
+        "    new_program = list(reversed(new_program_reversed))\n",
+        "    new_program = [instr for instr in new_program if instr[\"op\"] != \"NOP\"]\n",
+        "\n",
+        "    return new_program"
+      ],
+      "metadata": {
+        "id": "5vo-KgaffTFR"
+      },
+      "execution_count": 12,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for dead_code_elimination ===\n",
+        "\n",
+        "# Test 1: simple dead variable\n",
+        "test1 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\",    \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None},   # never used\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n",
+        "]\n",
+        "result1 = dead_code_elimination(test1)\n",
+        "print(\"Test 1 (kill unused CONST):\")\n",
+        "for i, instr in enumerate(result1): print(f\"  {i}: {instr}\")\n",
+        "# Expected: 'dead' instruction removed, 2 instructions remain\n",
+        "\n",
+        "# Test 2: chain of dead computation\n",
+        "test2 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"x\", \"src1\": \"a\", \"src2\": \"b\"},   # x never used → dead\n",
+        "    {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 7, \"src2\": None},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result2 = dead_code_elimination(test2)\n",
+        "print(\"\\nTest 2 (kill unused ADD and its feeders... but only if feeders are also unused):\")\n",
+        "for i, instr in enumerate(result2): print(f\"  {i}: {instr}\")\n",
+        "# Expected: x's ADD killed. a and b also killed (only fed into x, which is dead).\n",
+        "# Final: just CONST c=7, STORE.\n",
+        "\n",
+        "# Test 3: STORE always preserved\n",
+        "test3 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n",
+        "]\n",
+        "result3 = dead_code_elimination(test3)\n",
+        "print(\"\\nTest 3 (STORE preserved, feeder kept live):\")\n",
+        "for i, instr in enumerate(result3): print(f\"  {i}: {instr}\")\n",
+        "# Expected: both instructions unchanged\n",
+        "\n",
+        "# Test 4: variable used by STORE is live, even if defined far above\n",
+        "test4 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},   # used by STORE → live\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 9, \"src2\": None},   # never used → dead\n",
+        "    {\"op\": \"CONST\", \"dest\": \"c\", \"src1\": 1, \"src2\": None},   # never used → dead\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n",
+        "]\n",
+        "result4 = dead_code_elimination(test4)\n",
+        "print(\"\\nTest 4 (only 'a' is live, b/c killed):\")\n",
+        "for i, instr in enumerate(result4): print(f\"  {i}: {instr}\")\n",
+        "# Expected: a's CONST + STORE remain. b, c stripped.\n",
+        "\n",
+        "# Test 5: combine CF + DCE on a generated program\n",
+        "random.seed(42)\n",
+        "prog = generate_level_1()\n",
+        "print(\"\\nTest 5 (CF then DCE on seed=42):\")\n",
+        "print(\"Original:\")\n",
+        "for i, instr in enumerate(prog[\"instructions\"]): print(f\"  {i}: {instr}\")\n",
+        "\n",
+        "after_cf = constant_folding(prog[\"instructions\"])\n",
+        "print(\"After CF:\")\n",
+        "for i, instr in enumerate(after_cf): print(f\"  {i}: {instr}\")\n",
+        "\n",
+        "after_dce = dead_code_elimination(after_cf)\n",
+        "print(\"After CF + DCE:\")\n",
+        "for i, instr in enumerate(after_dce): print(f\"  {i}: {instr}\")\n",
+        "# Expected: original 4 instructions become much shorter — dead 'v1' eliminated, v2 folded\n",
+        "\n",
+        "# Test 6: idempotence — running DCE twice gives same result\n",
+        "test6 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"dead\", \"src1\": 99, \"src2\": None},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"a\", \"src2\": None},\n",
+        "]\n",
+        "once = dead_code_elimination(test6)\n",
+        "twice = dead_code_elimination(once)\n",
+        "print(\"\\nTest 6 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "vv8QRFRlvzdq",
+        "outputId": "e1d35874-ff1a-4dec-f43f-e22867dd3911"
+      },
+      "execution_count": 8,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Test 1 (kill unused CONST):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n",
+            "\n",
+            "Test 2 (kill unused ADD and its feeders... but only if feeders are also unused):\n",
+            "  0: {'op': 'CONST', 'dest': 'c', 'src1': 7, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 3 (STORE preserved, feeder kept live):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n",
+            "\n",
+            "Test 4 (only 'a' is live, b/c killed):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 5, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'a', 'src2': None}\n",
+            "\n",
+            "Test 5 (CF then DCE on seed=42):\n",
+            "Original:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "After CF:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "After CF + DCE:\n",
+            "  0: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "\n",
+            "Test 6 (idempotence): PASS\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def peephole_optimization(program):\n",
+        "    \"\"\"\n",
+        "    Single-instruction peephole rewrites. Replaces expensive ops with\n",
+        "    cheaper equivalents when operands hit special values (0, 1, 2, self).\n",
+        "\n",
+        "    Patterns (all preserve semantics):\n",
+        "      MUL by 2      -> ADD x + x\n",
+        "      MUL by 1      -> CONST (the other operand)\n",
+        "      MUL by 0      -> CONST 0\n",
+        "      ADD with 0    -> CONST (the other operand)\n",
+        "      SUB x - x     -> CONST 0\n",
+        "      DIV by 1      -> CONST (the dividend)\n",
+        "\n",
+        "    Notes:\n",
+        "    - Operates on individual instructions; no cross-instruction state.\n",
+        "    - Relies on constant_folding having propagated literal values into operands.\n",
+        "    - Returns fresh instruction dicts; never aliases inputs.\n",
+        "\n",
+        "    Returns a new list. Never raises, never returns None.\n",
+        "    \"\"\"\n",
+        "    new_program = []\n",
+        "\n",
+        "    for instr in program:\n",
+        "        instr = instr.copy()\n",
+        "        op   = instr[\"op\"]\n",
+        "        dest = instr[\"dest\"]\n",
+        "        src1 = instr[\"src1\"]\n",
+        "        src2 = instr[\"src2\"]\n",
+        "\n",
+        "        # === MUL patterns ===\n",
+        "        if op == \"MUL\":\n",
+        "            # MUL by 0 -> CONST 0\n",
+        "            if src1 == 0 or src2 == 0:\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n",
+        "                continue\n",
+        "            # MUL by 1 -> CONST (other operand) if the other operand is a literal,\n",
+        "            #            otherwise leave alone (we don't want to introduce a useless\n",
+        "            #            \"CONST dest = some_var_name\" — that's not a valid CONST).\n",
+        "            if src1 == 1 and isinstance(src2, int):\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n",
+        "                continue\n",
+        "            if src2 == 1 and isinstance(src1, int):\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n",
+        "                continue\n",
+        "            # MUL by 2 -> ADD x + x   (when one operand is literal 2 and the other is a var)\n",
+        "            if src1 == 2 and isinstance(src2, str):\n",
+        "                new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src2, \"src2\": src2})\n",
+        "                continue\n",
+        "            if src2 == 2 and isinstance(src1, str):\n",
+        "                new_program.append({\"op\": \"ADD\", \"dest\": dest, \"src1\": src1, \"src2\": src1})\n",
+        "                continue\n",
+        "            # No pattern matched — keep as is\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        # === ADD patterns ===\n",
+        "        elif op == \"ADD\":\n",
+        "            # ADD with 0 -> CONST (other operand) if the other operand is a literal\n",
+        "            if src1 == 0 and isinstance(src2, int):\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src2, \"src2\": None})\n",
+        "                continue\n",
+        "            if src2 == 0 and isinstance(src1, int):\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n",
+        "                continue\n",
+        "            # If \"ADD x + 0\" where x is a variable, we'd want a copy — but our schema\n",
+        "            # has no MOV/COPY op. Leave alone; CF + DCE handle the rest in practice.\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        # === SUB patterns ===\n",
+        "        elif op == \"SUB\":\n",
+        "            # SUB x - x -> CONST 0 (only if both sources are the same string variable)\n",
+        "            if isinstance(src1, str) and isinstance(src2, str) and src1 == src2:\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": 0, \"src2\": None})\n",
+        "                continue\n",
+        "            # SUB x - 0 with literal 0 on src2 — we'd want a copy; skip (no COPY op)\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        # === DIV patterns ===\n",
+        "        elif op == \"DIV\":\n",
+        "            # DIV by 1 -> CONST (dividend) if dividend is a literal\n",
+        "            if src2 == 1 and isinstance(src1, int):\n",
+        "                new_program.append({\"op\": \"CONST\", \"dest\": dest, \"src1\": src1, \"src2\": None})\n",
+        "                continue\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        else:\n",
+        "            # CONST, LOAD, STORE, NOP — pass through\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "    return new_program"
+      ],
+      "metadata": {
+        "id": "IaczDNeLv5PW"
+      },
+      "execution_count": 9,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for peephole_optimization ===\n",
+        "\n",
+        "# Test 1: MUL by 2 -> ADD x+x\n",
+        "test1 = [\n",
+        "    {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n",
+        "    {\"op\": \"MUL\",  \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result1 = peephole_optimization(test1)\n",
+        "print(\"Test 1 (MUL by 2 -> ADD self):\")\n",
+        "for i, instr in enumerate(result1): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 1 becomes ADD y+y\n",
+        "\n",
+        "# Test 2: MUL by 0 -> CONST 0\n",
+        "test2 = [\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"x\", \"src1\": \"y\",  \"src2\": 0},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result2 = peephole_optimization(test2)\n",
+        "print(\"\\nTest 2 (MUL by 0 -> CONST 0):\")\n",
+        "for i, instr in enumerate(result2): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 1 becomes CONST x = 0\n",
+        "\n",
+        "# Test 3: MUL by 1 (both literals)\n",
+        "test3 = [\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"x\", \"src1\": 7, \"src2\": 1},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result3 = peephole_optimization(test3)\n",
+        "print(\"\\nTest 3 (MUL 7*1 -> CONST 7):\")\n",
+        "for i, instr in enumerate(result3): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 0 becomes CONST x = 7\n",
+        "\n",
+        "# Test 4: SUB x - x -> CONST 0\n",
+        "test4 = [\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n",
+        "    {\"op\": \"SUB\",   \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result4 = peephole_optimization(test4)\n",
+        "print(\"\\nTest 4 (SUB y-y -> CONST 0):\")\n",
+        "for i, instr in enumerate(result4): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 1 becomes CONST x = 0\n",
+        "\n",
+        "# Test 5: ADD with 0\n",
+        "test5 = [\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"x\", \"src1\": 0, \"src2\": 5},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result5 = peephole_optimization(test5)\n",
+        "print(\"\\nTest 5 (ADD 0+5 -> CONST 5):\")\n",
+        "for i, instr in enumerate(result5): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 0 becomes CONST x = 5\n",
+        "\n",
+        "# Test 6: full pipeline — CF then peephole on a generated program\n",
+        "random.seed(7)\n",
+        "prog = generate_level_1()\n",
+        "print(\"\\nTest 6 (CF then peephole on seed=7):\")\n",
+        "print(\"Original:\")\n",
+        "for i, instr in enumerate(prog[\"instructions\"]): print(f\"  {i}: {instr}\")\n",
+        "after_cf = constant_folding(prog[\"instructions\"])\n",
+        "print(\"After CF:\")\n",
+        "for i, instr in enumerate(after_cf): print(f\"  {i}: {instr}\")\n",
+        "after_peep = peephole_optimization(after_cf)\n",
+        "print(\"After CF + peephole:\")\n",
+        "for i, instr in enumerate(after_peep): print(f\"  {i}: {instr}\")\n",
+        "\n",
+        "# Test 7: idempotence\n",
+        "test7 = [\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"x\", \"src1\": \"y\", \"src2\": 2},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "once = peephole_optimization(test7)\n",
+        "twice = peephole_optimization(once)\n",
+        "print(\"\\nTest 7 (idempotence):\", \"PASS\" if once == twice else \"FAIL\")\n",
+        "\n",
+        "# Test 8: no false fires — vanilla program shouldn't get rewritten\n",
+        "test8 = [\n",
+        "    {\"op\": \"LOAD\", \"dest\": \"y\", \"src1\": \"addr0\", \"src2\": None},\n",
+        "    {\"op\": \"MUL\",  \"dest\": \"x\", \"src1\": \"y\", \"src2\": \"y\"},   # y*y, no peephole pattern\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"x\", \"src2\": None},\n",
+        "]\n",
+        "result8 = peephole_optimization(test8)\n",
+        "print(\"\\nTest 8 (no false rewrite on y*y):\")\n",
+        "for i, instr in enumerate(result8): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 1 unchanged (still MUL y*y)"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "DWDUZUQvwXye",
+        "outputId": "1d1e327f-5f26-43e4-a6ac-f7a286da3781"
+      },
+      "execution_count": 10,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Test 1 (MUL by 2 -> ADD self):\n",
+            "  0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n",
+            "  1: {'op': 'ADD', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n",
+            "  2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n",
+            "\n",
+            "Test 2 (MUL by 0 -> CONST 0):\n",
+            "  0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n",
+            "  2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n",
+            "\n",
+            "Test 3 (MUL 7*1 -> CONST 7):\n",
+            "  0: {'op': 'CONST', 'dest': 'x', 'src1': 7, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n",
+            "\n",
+            "Test 4 (SUB y-y -> CONST 0):\n",
+            "  0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'x', 'src1': 0, 'src2': None}\n",
+            "  2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n",
+            "\n",
+            "Test 5 (ADD 0+5 -> CONST 5):\n",
+            "  0: {'op': 'CONST', 'dest': 'x', 'src1': 5, 'src2': None}\n",
+            "  1: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n",
+            "\n",
+            "Test 6 (CF then peephole on seed=7):\n",
+            "Original:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n",
+            "  3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n",
+            "  4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n",
+            "  5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "After CF:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n",
+            "  3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n",
+            "  4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n",
+            "  5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "After CF + peephole:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n",
+            "  3: {'op': 'CONST', 'dest': 'v3', 'src1': 8, 'src2': None}\n",
+            "  4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n",
+            "  5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "\n",
+            "Test 7 (idempotence): PASS\n",
+            "\n",
+            "Test 8 (no false rewrite on y*y):\n",
+            "  0: {'op': 'LOAD', 'dest': 'y', 'src1': 'addr0', 'src2': None}\n",
+            "  1: {'op': 'MUL', 'dest': 'x', 'src1': 'y', 'src2': 'y'}\n",
+            "  2: {'op': 'STORE', 'dest': 'addr0', 'src1': 'x', 'src2': None}\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def dump_ir(program_data):\n",
+        "    \"\"\"\n",
+        "    Convert a TAC program into a readable string for the LLM observation.\n",
+        "\n",
+        "    Format:\n",
+        "        // OBSERVABLE OUT: mem[0], mem[1]\n",
+        "        0: v0 = 3                       # CONST  [1 cycle]\n",
+        "        1: v1 = 5                       # CONST  [1 cycle]\n",
+        "        2: v2 = v0 + v1                 # ADD    [1 cycle]\n",
+        "        3: mem[addr0] = v2              # STORE  [4 cycles]\n",
+        "        // TOTAL: 7 cycles\n",
+        "\n",
+        "    Args:\n",
+        "        program_data: dict with keys 'instructions' and 'observable_addrs',\n",
+        "                      OR a raw list of instructions (legacy).\n",
+        "\n",
+        "    Returns:\n",
+        "        A multi-line string suitable for inclusion in an LLM prompt.\n",
+        "    \"\"\"\n",
+        "    # Accept both forms — full program dict or raw instruction list\n",
+        "    if isinstance(program_data, dict):\n",
+        "        instructions      = program_data[\"instructions\"]\n",
+        "        observable_addrs  = program_data.get(\"observable_addrs\", [])\n",
+        "    else:\n",
+        "        instructions      = program_data\n",
+        "        observable_addrs  = []\n",
+        "\n",
+        "    lines = []\n",
+        "\n",
+        "    # Header: observable outputs\n",
+        "    if observable_addrs:\n",
+        "        addr_str = \", \".join(f\"mem[{a}]\" for a in observable_addrs)\n",
+        "        lines.append(f\"// OBSERVABLE OUT: {addr_str}\")\n",
+        "\n",
+        "    # Body: each instruction in human-readable form\n",
+        "    total_cycles = 0\n",
+        "    for i, instr in enumerate(instructions):\n",
+        "        op   = instr[\"op\"]\n",
+        "        dest = instr[\"dest\"]\n",
+        "        src1 = instr[\"src1\"]\n",
+        "        src2 = instr[\"src2\"]\n",
+        "        cost = CYCLE_COSTS.get(op, 0)\n",
+        "        total_cycles += cost\n",
+        "\n",
+        "        # Render the instruction body\n",
+        "        if op == \"CONST\":\n",
+        "            body = f\"{dest} = {src1}\"\n",
+        "        elif op == \"ADD\":\n",
+        "            body = f\"{dest} = {src1} + {src2}\"\n",
+        "        elif op == \"SUB\":\n",
+        "            body = f\"{dest} = {src1} - {src2}\"\n",
+        "        elif op == \"MUL\":\n",
+        "            body = f\"{dest} = {src1} * {src2}\"\n",
+        "        elif op == \"DIV\":\n",
+        "            body = f\"{dest} = {src1} // {src2}\"\n",
+        "        elif op == \"LOAD\":\n",
+        "            body = f\"{dest} = mem[{src1}]\"\n",
+        "        elif op == \"STORE\":\n",
+        "            body = f\"mem[{dest}] = {src1}\"\n",
+        "        elif op == \"NOP\":\n",
+        "            body = \"nop\"\n",
+        "        else:\n",
+        "            body = f\"<unknown op: {op}>\"\n",
+        "\n",
+        "        cost_label = f\"{cost} cycle\" if cost == 1 else f\"{cost} cycles\"\n",
+        "        lines.append(f\"{i}: {body:<35} # {op:<6} [{cost_label}]\")\n",
+        "\n",
+        "    # Footer: total cycles\n",
+        "    lines.append(f\"// TOTAL: {total_cycles} cycles\")\n",
+        "\n",
+        "    return \"\\n\".join(lines)"
+      ],
+      "metadata": {
+        "id": "yB_yMNOZwaPT"
+      },
+      "execution_count": 13,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for state translator ===\n",
+        "\n",
+        "# Test 1: full generated program\n",
+        "random.seed(1)\n",
+        "prog = generate_level_1()\n",
+        "print(\"Test 1 (raw seed=1 program):\")\n",
+        "print(dump_ir(prog))\n",
+        "\n",
+        "# Test 2: after CF\n",
+        "print(\"\\nTest 2 (after CF):\")\n",
+        "folded = constant_folding(prog[\"instructions\"])\n",
+        "prog_after_cf = {**prog, \"instructions\": folded}\n",
+        "print(dump_ir(prog_after_cf))\n",
+        "\n",
+        "# Test 3: after CF + DCE — should show fewer instructions, lower total\n",
+        "print(\"\\nTest 3 (after CF + DCE):\")\n",
+        "optimized = dead_code_elimination(folded)\n",
+        "prog_optimized = {**prog, \"instructions\": optimized}\n",
+        "print(dump_ir(prog_optimized))\n",
+        "\n",
+        "# Test 4: every op type at least once\n",
+        "test4_instructions = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\",     \"src1\": 5,        \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"b\",     \"src1\": \"a\",      \"src2\": 3},\n",
+        "    {\"op\": \"SUB\",   \"dest\": \"c\",     \"src1\": \"b\",      \"src2\": \"a\"},\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"d\",     \"src1\": \"c\",      \"src2\": 2},\n",
+        "    {\"op\": \"DIV\",   \"dest\": \"e\",     \"src1\": \"d\",      \"src2\": 4},\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"f\",     \"src1\": \"addr0\",  \"src2\": None},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"e\",      \"src2\": None},\n",
+        "    {\"op\": \"NOP\",   \"dest\": None,    \"src1\": None,     \"src2\": None},\n",
+        "]\n",
+        "test4_data = {\n",
+        "    \"instructions\":     test4_instructions,\n",
+        "    \"observable_addrs\": [0, 1],\n",
+        "}\n",
+        "print(\"\\nTest 4 (every op type):\")\n",
+        "print(dump_ir(test4_data))\n",
+        "\n",
+        "# Test 5: empty program shouldn't crash\n",
+        "print(\"\\nTest 5 (empty program):\")\n",
+        "print(dump_ir({\"instructions\": [], \"observable_addrs\": [0]}))"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "SUqa7odp3wB_",
+        "outputId": "c3f44089-213a-4a69-f29f-42b755ed142a"
+      },
+      "execution_count": 14,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Test 1 (raw seed=1 program):\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 10                             # CONST  [1 cycle]\n",
+            "1: v1 = 2                              # CONST  [1 cycle]\n",
+            "2: v2 = v1 + v1                        # ADD    [1 cycle]\n",
+            "3: v3 = v2 * v1                        # MUL    [3 cycles]\n",
+            "4: mem[addr0] = v3                     # STORE  [4 cycles]\n",
+            "// TOTAL: 10 cycles\n",
+            "\n",
+            "Test 2 (after CF):\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 10                             # CONST  [1 cycle]\n",
+            "1: v1 = 2                              # CONST  [1 cycle]\n",
+            "2: v2 = 4                              # CONST  [1 cycle]\n",
+            "3: v3 = 8                              # CONST  [1 cycle]\n",
+            "4: mem[addr0] = v3                     # STORE  [4 cycles]\n",
+            "// TOTAL: 8 cycles\n",
+            "\n",
+            "Test 3 (after CF + DCE):\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v3 = 8                              # CONST  [1 cycle]\n",
+            "1: mem[addr0] = v3                     # STORE  [4 cycles]\n",
+            "// TOTAL: 5 cycles\n",
+            "\n",
+            "Test 4 (every op type):\n",
+            "// OBSERVABLE OUT: mem[0], mem[1]\n",
+            "0: a = 5                               # CONST  [1 cycle]\n",
+            "1: b = a + 3                           # ADD    [1 cycle]\n",
+            "2: c = b - a                           # SUB    [1 cycle]\n",
+            "3: d = c * 2                           # MUL    [3 cycles]\n",
+            "4: e = d // 4                          # DIV    [5 cycles]\n",
+            "5: f = mem[addr0]                      # LOAD   [4 cycles]\n",
+            "6: mem[addr0] = e                      # STORE  [4 cycles]\n",
+            "7: nop                                 # NOP    [0 cycles]\n",
+            "// TOTAL: 19 cycles\n",
+            "\n",
+            "Test 5 (empty program):\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "// TOTAL: 0 cycles\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def generate_level_2():\n",
+        "    \"\"\"\n",
+        "    Generate a Level 2 Toy-IR program.\n",
+        "\n",
+        "    Characteristics:\n",
+        "    - 8-12 instructions\n",
+        "    - 3-5 CONSTs (some literal, some used in arithmetic)\n",
+        "    - 3-5 arithmetic ops (ADD, SUB, MUL, DIV) with chaining\n",
+        "    - 1-3 dead variables (DCE opportunities)\n",
+        "    - 1 LOAD from initial memory (introduces non-constant variable)\n",
+        "    - 1 STORE at the end (observable output)\n",
+        "    - Mix of operands designed to hit peephole patterns occasionally:\n",
+        "        - MUL by literal 2 / 1 / 0 (not always — diversity matters)\n",
+        "        - DIV by non-zero divisor only (Design Assumption #4)\n",
+        "\n",
+        "    Returns:\n",
+        "        dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n",
+        "    \"\"\"\n",
+        "    instructions = []\n",
+        "    var_counter = 0\n",
+        "\n",
+        "    def new_var():\n",
+        "        nonlocal var_counter\n",
+        "        name = f\"v{var_counter}\"\n",
+        "        var_counter += 1\n",
+        "        return name\n",
+        "\n",
+        "    available_vars = []  # vars that exist and can be used as sources\n",
+        "\n",
+        "    # Step 1: 3-5 CONSTs\n",
+        "    num_consts = random.randint(3, 5)\n",
+        "    for _ in range(num_consts):\n",
+        "        var = new_var()\n",
+        "        value = random.randint(1, 10)\n",
+        "        instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n",
+        "        available_vars.append(var)\n",
+        "\n",
+        "    # Step 2: 1 LOAD from initial memory\n",
+        "    # We'll seed initial_mem with something at a fresh address.\n",
+        "    load_addr_var = \"addr_in\"\n",
+        "    loaded_var    = new_var()\n",
+        "    instructions.append({\n",
+        "        \"op\": \"LOAD\", \"dest\": loaded_var, \"src1\": load_addr_var, \"src2\": None,\n",
+        "    })\n",
+        "    available_vars.append(loaded_var)\n",
+        "\n",
+        "    # Step 3: 3-5 arithmetic ops, chained\n",
+        "    num_arith = random.randint(3, 5)\n",
+        "    last_result = None\n",
+        "    for _ in range(num_arith):\n",
+        "        op = random.choice([\"ADD\", \"SUB\", \"MUL\", \"DIV\"])\n",
+        "\n",
+        "        # 30% chance to use a literal as src2 to expose peephole opportunities\n",
+        "        # (MUL by 2/1/0, etc.)\n",
+        "        use_literal_src2 = random.random() < 0.3\n",
+        "\n",
+        "        src1 = random.choice(available_vars)\n",
+        "        if use_literal_src2:\n",
+        "            # For DIV, never emit literal 0 (Design Assumption #4)\n",
+        "            if op == \"DIV\":\n",
+        "                src2 = random.choice([1, 2, 3, 4])  # safe non-zero divisors\n",
+        "            else:\n",
+        "                src2 = random.choice([0, 1, 2, 3])  # 0/1/2 hit peephole patterns\n",
+        "        else:\n",
+        "            src2 = random.choice(available_vars)\n",
+        "            # If we sampled a variable for DIV's src2, we can't be sure it's non-zero\n",
+        "            # at runtime. To stay safe per Assumption #4, restrict DIV to literal src2.\n",
+        "            if op == \"DIV\":\n",
+        "                src2 = random.choice([1, 2, 3, 4])\n",
+        "\n",
+        "        dest = new_var()\n",
+        "        instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n",
+        "        available_vars.append(dest)\n",
+        "        last_result = dest\n",
+        "\n",
+        "    # Step 4: 1-3 extra dead CONSTs sprinkled in (DCE targets)\n",
+        "    num_dead = random.randint(1, 3)\n",
+        "    for _ in range(num_dead):\n",
+        "        dead_var   = new_var()\n",
+        "        dead_value = random.randint(1, 20)\n",
+        "        # Insert at a random position before the (eventual) STORE\n",
+        "        insert_pos = random.randint(0, len(instructions))\n",
+        "        instructions.insert(insert_pos, {\n",
+        "            \"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None,\n",
+        "        })\n",
+        "        # Note: dead_var intentionally never used afterward\n",
+        "\n",
+        "    # Step 5: STORE the final result\n",
+        "    initial_vars = {\n",
+        "        \"addr0\":   0,\n",
+        "        \"addr_in\": 1,   # address that the LOAD reads from\n",
+        "    }\n",
+        "    initial_mem = {1: random.randint(1, 20)}  # seed mem[1] with a random value\n",
+        "\n",
+        "    instructions.append({\n",
+        "        \"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None,\n",
+        "    })\n",
+        "\n",
+        "    return {\n",
+        "        \"initial_vars\":     initial_vars,\n",
+        "        \"initial_mem\":      initial_mem,\n",
+        "        \"instructions\":     instructions,\n",
+        "        \"observable_addrs\": [0],\n",
+        "    }"
+      ],
+      "metadata": {
+        "id": "OL6QbL0H3xuP"
+      },
+      "execution_count": 15,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for Level 2 generator ===\n",
+        "\n",
+        "# Test 1: spot-check a few seeds\n",
+        "for seed in [42, 1, 7]:\n",
+        "    random.seed(seed)\n",
+        "    prog = generate_level_2()\n",
+        "    print(f\"\\n=== Level 2, seed={seed} ===\")\n",
+        "    print(dump_ir(prog))\n",
+        "    print(f\"initial_vars: {prog['initial_vars']}\")\n",
+        "    print(f\"initial_mem : {prog['initial_mem']}\")\n",
+        "\n",
+        "# Test 2: full pipeline (CF + DCE + peephole) on a Level 2 program\n",
+        "print(\"\\n=== Full optimization pipeline on Level 2 (seed=42) ===\")\n",
+        "random.seed(42)\n",
+        "prog = generate_level_2()\n",
+        "print(\"ORIGINAL:\")\n",
+        "print(dump_ir(prog))\n",
+        "\n",
+        "after_cf = constant_folding(prog[\"instructions\"])\n",
+        "print(\"\\nAFTER CF:\")\n",
+        "print(dump_ir({**prog, \"instructions\": after_cf}))\n",
+        "\n",
+        "after_dce = dead_code_elimination(after_cf)\n",
+        "print(\"\\nAFTER CF + DCE:\")\n",
+        "print(dump_ir({**prog, \"instructions\": after_dce}))\n",
+        "\n",
+        "after_peep = peephole_optimization(after_dce)\n",
+        "print(\"\\nAFTER CF + DCE + PEEPHOLE:\")\n",
+        "print(dump_ir({**prog, \"instructions\": after_peep}))"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "D0XjNetU5gIo",
+        "outputId": "1c35cdfc-a95d-48ff-c7ce-c4c236bbbbdb"
+      },
+      "execution_count": 16,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\n",
+            "=== Level 2, seed=42 ===\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 2                              # CONST  [1 cycle]\n",
+            "1: v1 = 1                              # CONST  [1 cycle]\n",
+            "2: v2 = 5                              # CONST  [1 cycle]\n",
+            "3: v3 = 4                              # CONST  [1 cycle]\n",
+            "4: v11 = 19                            # CONST  [1 cycle]\n",
+            "5: v4 = 4                              # CONST  [1 cycle]\n",
+            "6: v5 = mem[addr_in]                   # LOAD   [4 cycles]\n",
+            "7: v9 = 18                             # CONST  [1 cycle]\n",
+            "8: v10 = 8                             # CONST  [1 cycle]\n",
+            "9: v6 = v4 + v0                        # ADD    [1 cycle]\n",
+            "10: v7 = v0 // 2                        # DIV    [5 cycles]\n",
+            "11: v8 = v0 - v3                        # SUB    [1 cycle]\n",
+            "12: mem[addr0] = v8                     # STORE  [4 cycles]\n",
+            "// TOTAL: 23 cycles\n",
+            "initial_vars: {'addr0': 0, 'addr_in': 1}\n",
+            "initial_mem : {1: 1}\n",
+            "\n",
+            "=== Level 2, seed=1 ===\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 10                             # CONST  [1 cycle]\n",
+            "1: v7 = 19                             # CONST  [1 cycle]\n",
+            "2: v1 = 2                              # CONST  [1 cycle]\n",
+            "3: v2 = 5                              # CONST  [1 cycle]\n",
+            "4: v3 = mem[addr_in]                   # LOAD   [4 cycles]\n",
+            "5: v4 = v3 // 2                        # DIV    [5 cycles]\n",
+            "6: v5 = v3 + v3                        # ADD    [1 cycle]\n",
+            "7: v6 = v2 + v5                        # ADD    [1 cycle]\n",
+            "8: mem[addr0] = v6                     # STORE  [4 cycles]\n",
+            "// TOTAL: 19 cycles\n",
+            "initial_vars: {'addr0': 0, 'addr_in': 1}\n",
+            "initial_mem : {1: 11}\n",
+            "\n",
+            "=== Level 2, seed=7 ===\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 3                              # CONST  [1 cycle]\n",
+            "1: v1 = 7                              # CONST  [1 cycle]\n",
+            "2: v2 = 1                              # CONST  [1 cycle]\n",
+            "3: v3 = 2                              # CONST  [1 cycle]\n",
+            "4: v4 = mem[addr_in]                   # LOAD   [4 cycles]\n",
+            "5: v5 = v0 + v4                        # ADD    [1 cycle]\n",
+            "6: v6 = v3 - 3                         # SUB    [1 cycle]\n",
+            "7: v7 = v4 + 3                         # ADD    [1 cycle]\n",
+            "8: v10 = 2                             # CONST  [1 cycle]\n",
+            "9: v8 = v1 + v3                        # ADD    [1 cycle]\n",
+            "10: v9 = v6 + v0                        # ADD    [1 cycle]\n",
+            "11: mem[addr0] = v9                     # STORE  [4 cycles]\n",
+            "// TOTAL: 18 cycles\n",
+            "initial_vars: {'addr0': 0, 'addr_in': 1}\n",
+            "initial_mem : {1: 5}\n",
+            "\n",
+            "=== Full optimization pipeline on Level 2 (seed=42) ===\n",
+            "ORIGINAL:\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 2                              # CONST  [1 cycle]\n",
+            "1: v1 = 1                              # CONST  [1 cycle]\n",
+            "2: v2 = 5                              # CONST  [1 cycle]\n",
+            "3: v3 = 4                              # CONST  [1 cycle]\n",
+            "4: v11 = 19                            # CONST  [1 cycle]\n",
+            "5: v4 = 4                              # CONST  [1 cycle]\n",
+            "6: v5 = mem[addr_in]                   # LOAD   [4 cycles]\n",
+            "7: v9 = 18                             # CONST  [1 cycle]\n",
+            "8: v10 = 8                             # CONST  [1 cycle]\n",
+            "9: v6 = v4 + v0                        # ADD    [1 cycle]\n",
+            "10: v7 = v0 // 2                        # DIV    [5 cycles]\n",
+            "11: v8 = v0 - v3                        # SUB    [1 cycle]\n",
+            "12: mem[addr0] = v8                     # STORE  [4 cycles]\n",
+            "// TOTAL: 23 cycles\n",
+            "\n",
+            "AFTER CF:\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v0 = 2                              # CONST  [1 cycle]\n",
+            "1: v1 = 1                              # CONST  [1 cycle]\n",
+            "2: v2 = 5                              # CONST  [1 cycle]\n",
+            "3: v3 = 4                              # CONST  [1 cycle]\n",
+            "4: v11 = 19                            # CONST  [1 cycle]\n",
+            "5: v4 = 4                              # CONST  [1 cycle]\n",
+            "6: v5 = mem[addr_in]                   # LOAD   [4 cycles]\n",
+            "7: v9 = 18                             # CONST  [1 cycle]\n",
+            "8: v10 = 8                             # CONST  [1 cycle]\n",
+            "9: v6 = 6                              # CONST  [1 cycle]\n",
+            "10: v7 = 1                              # CONST  [1 cycle]\n",
+            "11: v8 = -2                             # CONST  [1 cycle]\n",
+            "12: mem[addr0] = v8                     # STORE  [4 cycles]\n",
+            "// TOTAL: 19 cycles\n",
+            "\n",
+            "AFTER CF + DCE:\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v8 = -2                             # CONST  [1 cycle]\n",
+            "1: mem[addr0] = v8                     # STORE  [4 cycles]\n",
+            "// TOTAL: 5 cycles\n",
+            "\n",
+            "AFTER CF + DCE + PEEPHOLE:\n",
+            "// OBSERVABLE OUT: mem[0]\n",
+            "0: v8 = -2                             # CONST  [1 cycle]\n",
+            "1: mem[addr0] = v8                     # STORE  [4 cycles]\n",
+            "// TOTAL: 5 cycles\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "Vo-9MitD5hwM"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/metahack1.ipynb

@@ -0,0 +1,484 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "## Design Assumptions (do not violate)\n",
+        "\n",
+        "# 1. **DCE never eliminates STOREs.** They define program output (final mem state).\n",
+        "# 2. **Addresses are distinct by construction.** Generator allocates each address variable to a unique integer; no aliasing.\n",
+        "# 3. **CF refuses to fold DIV by zero.** If src2 == 0 on a DIV op, leave instruction unchanged.\n",
+        "# 4. **Generator never emits DIV by literal zero.** When DIV is generated, src2 is always a non-zero constant or a variable known to be non-zero.\n",
+        "# # 5. **Integer arithmetic only.** No floats anywhere — avoids equivalence-check precision issues.\n",
+        "# 6. Generator declares `observable_addrs` per program — verifier compares only these mem entries.\n",
+        "# 7. State translator annotates observable outputs at top of dump.\n",
+        "# 8. Reward distinguishes broken (-1000) from valid-but-worse (small negative) — Harshal's formula.\n",
+        "# 9. Multi-input verification: 3-5 random initial states, all must match.\n",
+        "# 10. Integer division uses Python floor division (//). Aarush's VM must match."
+      ],
+      "metadata": {
+        "id": "_XI5jT2Ibvrf"
+      },
+      "execution_count": 17,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "id": "j6CQ327KWXwr"
+      },
+      "outputs": [],
+      "source": [
+        "# === TAC Schema v1.0 (LOCKED with Role 1 / Aarush) ===\n",
+        "# Reverse passes deferred to stretch goal — not in initial action space.\n",
+        "\n",
+        "OPS = [\"CONST\", \"ADD\", \"SUB\", \"MUL\", \"DIV\", \"LOAD\", \"STORE\", \"NOP\"]\n",
+        "\n",
+        "CYCLE_COSTS = {\n",
+        "    \"CONST\": 1,\n",
+        "    \"ADD\":   1,\n",
+        "    \"SUB\":   1,\n",
+        "    \"MUL\":   3,\n",
+        "    \"DIV\":   5,\n",
+        "    \"LOAD\":  4,\n",
+        "    \"STORE\": 4,\n",
+        "    \"NOP\":   0,\n",
+        "}\n",
+        "\n",
+        "# Instruction shape: {\"op\": str, \"dest\": str|None, \"src1\": Any, \"src2\": Any}\n",
+        "# Operands: str = variable name, int = literal constant, None = unused\n",
+        "#\n",
+        "# Op semantics:\n",
+        "#   CONST: dest = src1                      (src1 is int literal, src2 = None)\n",
+        "#   ADD/SUB/MUL/DIV: dest = src1 OP src2    (src1, src2 are var names or int literals)\n",
+        "#   LOAD:  dest = mem[src1]                 (src1 is a var holding an address)\n",
+        "#   STORE: mem[dest] = src1                 (dest is a var holding an address)\n",
+        "#   NOP:   no-op                            (all fields None)\n",
+        "#\n",
+        "# Program output (for equivalence check) = final memory state (mem dict).\n",
+        "# Programs ship as: (initial_vars: dict, initial_mem: dict, instructions: list[dict])\n",
+        "# 10. Integer division uses Python floor division (//). Aarush's VM must match."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import random\n",
+        "\n",
+        "def generate_level_1():\n",
+        "    \"\"\"\n",
+        "    Generate a Level 1 Toy-IR program.\n",
+        "\n",
+        "    Characteristics:\n",
+        "    - 4-6 instructions\n",
+        "    - 2-3 CONST ops with literal values\n",
+        "    - 1-2 arithmetic ops on those constants (foldable by CF)\n",
+        "    - 0-1 dead variables (killable by DCE)\n",
+        "    - Exactly 1 STORE at the end so the program has an observable output\n",
+        "\n",
+        "    Returns:\n",
+        "        dict with keys: initial_vars, initial_mem, instructions, observable_addrs\n",
+        "    \"\"\"\n",
+        "    instructions = []\n",
+        "    var_counter = 0\n",
+        "\n",
+        "    def new_var():\n",
+        "        nonlocal var_counter\n",
+        "        name = f\"v{var_counter}\"\n",
+        "        var_counter += 1\n",
+        "        return name\n",
+        "\n",
+        "    # Step 1: Generate 2-3 constant assignments\n",
+        "    num_consts = random.randint(2, 3)\n",
+        "    const_vars = []\n",
+        "    for _ in range(num_consts):\n",
+        "        var = new_var()\n",
+        "        value = random.randint(1, 10)\n",
+        "        instructions.append({\"op\": \"CONST\", \"dest\": var, \"src1\": value, \"src2\": None})\n",
+        "        const_vars.append(var)\n",
+        "\n",
+        "    # Step 2: Generate 1-2 arithmetic ops using those constants\n",
+        "    num_arith = random.randint(1, 2)\n",
+        "    last_result = None\n",
+        "    for _ in range(num_arith):\n",
+        "        op = random.choice([\"ADD\", \"MUL\"])\n",
+        "        src1 = random.choice(const_vars)\n",
+        "        src2 = random.choice(const_vars)\n",
+        "        dest = new_var()\n",
+        "        instructions.append({\"op\": op, \"dest\": dest, \"src1\": src1, \"src2\": src2})\n",
+        "        last_result = dest\n",
+        "        const_vars.append(dest)\n",
+        "\n",
+        "    # Step 3: Optionally add 1 dead variable (50% chance)\n",
+        "    if random.random() < 0.5:\n",
+        "        dead_var = new_var()\n",
+        "        dead_value = random.randint(1, 10)\n",
+        "        instructions.append({\"op\": \"CONST\", \"dest\": dead_var, \"src1\": dead_value, \"src2\": None})\n",
+        "        # Note: dead_var is intentionally never used — DCE should catch it.\n",
+        "\n",
+        "    # Step 4: Add a STORE at the end so the program has observable output\n",
+        "    initial_vars = {\"addr0\": 0}\n",
+        "    instructions.append({\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": last_result, \"src2\": None})\n",
+        "\n",
+        "    initial_mem = {}\n",
+        "\n",
+        "    return {\n",
+        "        \"initial_vars\":     initial_vars,\n",
+        "        \"initial_mem\":      initial_mem,\n",
+        "        \"instructions\":     instructions,\n",
+        "        \"observable_addrs\": [0],   # Aarush's verifier compares only these mem entries\n",
+        "    }\n",
+        "\n",
+        "\n",
+        "# Sanity-check: generate a few programs and print them\n",
+        "for seed in [42, 1, 7, 99]:\n",
+        "    random.seed(seed)\n",
+        "    prog = generate_level_1()\n",
+        "    print(f\"\\n=== seed={seed} ===\")\n",
+        "    print(f\"initial_vars     : {prog['initial_vars']}\")\n",
+        "    print(f\"initial_mem      : {prog['initial_mem']}\")\n",
+        "    print(f\"observable_addrs : {prog['observable_addrs']}\")\n",
+        "    print(f\"instructions ({len(prog['instructions'])}):\")\n",
+        "    for i, instr in enumerate(prog['instructions']):\n",
+        "        print(f\"  {i}: {instr}\")"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "a86vD1OyWcB9",
+        "outputId": "de3d7093-eaea-420e-a0d7-2bd9451bcf6b"
+      },
+      "execution_count": 14,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\n",
+            "=== seed=42 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (4):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "\n",
+            "=== seed=1 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (5):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 10, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 2, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v1', 'src2': 'v1'}\n",
+            "  3: {'op': 'MUL', 'dest': 'v3', 'src1': 'v2', 'src2': 'v1'}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "\n",
+            "=== seed=7 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (6):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 7, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 1, 'src2': None}\n",
+            "  3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v1', 'src2': 'v2'}\n",
+            "  4: {'op': 'CONST', 'dest': 'v4', 'src1': 9, 'src2': None}\n",
+            "  5: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n",
+            "\n",
+            "=== seed=99 ===\n",
+            "initial_vars     : {'addr0': 0}\n",
+            "initial_mem      : {}\n",
+            "observable_addrs : [0]\n",
+            "instructions (5):\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 7, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 4, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 10, 'src2': None}\n",
+            "  3: {'op': 'ADD', 'dest': 'v3', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v3', 'src2': None}\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def _resolve_operand(operand, known_constants):\n",
+        "    \"\"\"\n",
+        "    Given a TAC operand (string variable name or int literal),\n",
+        "    return its concrete integer value if known, or None if unknown.\n",
+        "    \"\"\"\n",
+        "    if isinstance(operand, int):\n",
+        "        return operand\n",
+        "    if isinstance(operand, str) and operand in known_constants:\n",
+        "        return known_constants[operand]\n",
+        "    return None\n",
+        "\n",
+        "\n",
+        "def _compute(op, v1, v2):\n",
+        "    \"\"\"Compute the result of a binary arithmetic op on two known integers.\"\"\"\n",
+        "    if op == \"ADD\": return v1 + v2\n",
+        "    if op == \"SUB\": return v1 - v2\n",
+        "    if op == \"MUL\": return v1 * v2\n",
+        "    if op == \"DIV\": return v1 // v2   # floor division (locked semantics)\n",
+        "    raise ValueError(f\"_compute called with non-arithmetic op: {op}\")\n",
+        "\n",
+        "\n",
+        "def constant_folding(program):\n",
+        "    \"\"\"\n",
+        "    Forward pass that folds constant arithmetic into CONST ops, and\n",
+        "    propagates known constants into instruction operands.\n",
+        "\n",
+        "    Behavior:\n",
+        "    - If both sources of an arithmetic op resolve to known integers,\n",
+        "      replaces the instruction with a CONST holding the computed result.\n",
+        "    - If only one source is known, still substitutes that known value\n",
+        "      into the instruction (constant propagation), enabling downstream\n",
+        "      passes (e.g., peephole) to recognize patterns like ADD-with-0 or MUL-by-1.\n",
+        "    - Refuses to fold DIV by zero (Design Assumption #3).\n",
+        "    - Always returns fresh dicts; never aliases input instructions.\n",
+        "\n",
+        "    Args:\n",
+        "        program: list of TAC instruction dicts (per locked schema).\n",
+        "\n",
+        "    Returns:\n",
+        "        new list of TAC instruction dicts. Always semantics-preserving.\n",
+        "        Never raises, never returns None.\n",
+        "    \"\"\"\n",
+        "    known_constants = {}\n",
+        "    new_program = []\n",
+        "\n",
+        "    for instr in program:\n",
+        "        # Always work on a copy — never alias input dicts\n",
+        "        instr = instr.copy()\n",
+        "        op   = instr[\"op\"]\n",
+        "        dest = instr[\"dest\"]\n",
+        "\n",
+        "        if op == \"CONST\":\n",
+        "            known_constants[dest] = instr[\"src1\"]\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        elif op in (\"ADD\", \"SUB\", \"MUL\", \"DIV\"):\n",
+        "            # === Constant propagation: substitute known constants into operands ===\n",
+        "            if isinstance(instr[\"src1\"], str) and instr[\"src1\"] in known_constants:\n",
+        "                instr[\"src1\"] = known_constants[instr[\"src1\"]]\n",
+        "            if isinstance(instr[\"src2\"], str) and instr[\"src2\"] in known_constants:\n",
+        "                instr[\"src2\"] = known_constants[instr[\"src2\"]]\n",
+        "\n",
+        "            # === Try to fold ===\n",
+        "            v1 = _resolve_operand(instr[\"src1\"], known_constants)\n",
+        "            v2 = _resolve_operand(instr[\"src2\"], known_constants)\n",
+        "\n",
+        "            if v1 is not None and v2 is not None:\n",
+        "                # Both operands are known integers\n",
+        "                if op == \"DIV\" and v2 == 0:\n",
+        "                    # Refuse to fold DIV by zero\n",
+        "                    new_program.append(instr)\n",
+        "                    known_constants.pop(dest, None)\n",
+        "                else:\n",
+        "                    # Fold: replace with CONST\n",
+        "                    result = _compute(op, v1, v2)\n",
+        "                    new_program.append({\n",
+        "                        \"op\":   \"CONST\",\n",
+        "                        \"dest\": dest,\n",
+        "                        \"src1\": result,\n",
+        "                        \"src2\": None,\n",
+        "                    })\n",
+        "                    known_constants[dest] = result\n",
+        "            else:\n",
+        "                # Can't fold (at least one operand unknown).\n",
+        "                # Instruction may still have been mutated by propagation above.\n",
+        "                new_program.append(instr)\n",
+        "                known_constants.pop(dest, None)\n",
+        "\n",
+        "        elif op == \"LOAD\":\n",
+        "            # Memory reads aren't statically resolvable\n",
+        "            new_program.append(instr)\n",
+        "            known_constants.pop(dest, None)\n",
+        "\n",
+        "        elif op in (\"STORE\", \"NOP\"):\n",
+        "            # No dest tracking needed.\n",
+        "            # Note: we COULD propagate src1 into a STORE for downstream readability,\n",
+        "            # but the cycle cost is unchanged and the executor handles vars fine.\n",
+        "            # Leave STORE alone — keeps the code minimal.\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "        else:\n",
+        "            # Unknown op — defensive pass-through\n",
+        "            new_program.append(instr)\n",
+        "\n",
+        "    return new_program"
+      ],
+      "metadata": {
+        "id": "CGk8Fz4CZA3t"
+      },
+      "execution_count": 18,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# === Sanity tests for constant_folding ===\n",
+        "\n",
+        "# Test 1: simple fold — ADD of two CONSTs\n",
+        "test1 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result1 = constant_folding(test1)\n",
+        "print(\"Test 1 (simple ADD fold):\")\n",
+        "for i, instr in enumerate(result1): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 becomes CONST c = 8\n",
+        "\n",
+        "# Test 2: chained fold — second op uses first op's folded result\n",
+        "test2 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 5, \"src2\": None},\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},   # → 8\n",
+        "    {\"op\": \"MUL\",   \"dest\": \"d\", \"src1\": \"c\", \"src2\": \"b\"},   # → 8 * 5 = 40\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"d\", \"src2\": None},\n",
+        "]\n",
+        "result2 = constant_folding(test2)\n",
+        "print(\"\\nTest 2 (chained fold):\")\n",
+        "for i, instr in enumerate(result2): print(f\"  {i}: {instr}\")\n",
+        "# Expected: c becomes CONST 8, d becomes CONST 40\n",
+        "\n",
+        "# Test 3: DIV by zero refused\n",
+        "test3 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 10, \"src2\": None},\n",
+        "    {\"op\": \"CONST\", \"dest\": \"b\", \"src1\": 0,  \"src2\": None},\n",
+        "    {\"op\": \"DIV\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},   # 10/0 — must NOT fold\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result3 = constant_folding(test3)\n",
+        "print(\"\\nTest 3 (DIV by zero refused):\")\n",
+        "for i, instr in enumerate(result3): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 unchanged (still DIV, not CONST)\n",
+        "\n",
+        "# Test 4: unknown source can't fold\n",
+        "test4 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 3, \"src2\": None},\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None},   # b is unknown\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},        # can't fold\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result4 = constant_folding(test4)\n",
+        "print(\"\\nTest 4 (LOAD makes b unknown, ADD not folded):\")\n",
+        "for i, instr in enumerate(result4): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 unchanged (still ADD)\n",
+        "\n",
+        "# Test 5: idempotence — running CF on a generated program\n",
+        "random.seed(42)\n",
+        "prog = generate_level_1()\n",
+        "folded = constant_folding(prog[\"instructions\"])\n",
+        "print(\"\\nTest 5 (CF on generated Level 1 program, seed=42):\")\n",
+        "print(\"Before:\")\n",
+        "for i, instr in enumerate(prog[\"instructions\"]): print(f\"  {i}: {instr}\")\n",
+        "print(\"After:\")\n",
+        "for i, instr in enumerate(folded): print(f\"  {i}: {instr}\")\n",
+        "\n",
+        "# Test 6: constant propagation — only one source is known\n",
+        "test6 = [\n",
+        "    {\"op\": \"CONST\", \"dest\": \"a\", \"src1\": 0, \"src2\": None},\n",
+        "    {\"op\": \"LOAD\",  \"dest\": \"b\", \"src1\": \"addr0\", \"src2\": None},   # b unknown\n",
+        "    {\"op\": \"ADD\",   \"dest\": \"c\", \"src1\": \"a\", \"src2\": \"b\"},        # a known (=0), b unknown\n",
+        "    {\"op\": \"STORE\", \"dest\": \"addr0\", \"src1\": \"c\", \"src2\": None},\n",
+        "]\n",
+        "result6 = constant_folding(test6)\n",
+        "print(\"\\nTest 6 (propagation: a=0 substituted into ADD even though b unknown):\")\n",
+        "for i, instr in enumerate(result6): print(f\"  {i}: {instr}\")\n",
+        "# Expected: instruction 2 is still ADD (can't fold — b unknown), but src1 is now literal 0, not 'a'\n",
+        "# This sets up peephole to recognize \"ADD with 0\" later"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "WLKSTVa8fRF8",
+        "outputId": "27042d38-0044-4814-bbfa-8a2c2d87ee06"
+      },
+      "execution_count": 20,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Test 1 (simple ADD fold):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 2 (chained fold):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'c', 'src1': 8, 'src2': None}\n",
+            "  3: {'op': 'CONST', 'dest': 'd', 'src1': 40, 'src2': None}\n",
+            "  4: {'op': 'STORE', 'dest': 'addr0', 'src1': 'd', 'src2': None}\n",
+            "\n",
+            "Test 3 (DIV by zero refused):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 10, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'b', 'src1': 0, 'src2': None}\n",
+            "  2: {'op': 'DIV', 'dest': 'c', 'src1': 10, 'src2': 0}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 4 (LOAD makes b unknown, ADD not folded):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 3, 'src2': None}\n",
+            "  1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'c', 'src1': 3, 'src2': 'b'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n",
+            "\n",
+            "Test 5 (CF on generated Level 1 program, seed=42):\n",
+            "Before:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'v2', 'src1': 'v0', 'src2': 'v0'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "After:\n",
+            "  0: {'op': 'CONST', 'dest': 'v0', 'src1': 1, 'src2': None}\n",
+            "  1: {'op': 'CONST', 'dest': 'v1', 'src1': 5, 'src2': None}\n",
+            "  2: {'op': 'CONST', 'dest': 'v2', 'src1': 2, 'src2': None}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'v2', 'src2': None}\n",
+            "\n",
+            "Test 6 (propagation: a=0 substituted into ADD even though b unknown):\n",
+            "  0: {'op': 'CONST', 'dest': 'a', 'src1': 0, 'src2': None}\n",
+            "  1: {'op': 'LOAD', 'dest': 'b', 'src1': 'addr0', 'src2': None}\n",
+            "  2: {'op': 'ADD', 'dest': 'c', 'src1': 0, 'src2': 'b'}\n",
+            "  3: {'op': 'STORE', 'dest': 'addr0', 'src1': 'c', 'src2': None}\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "5vo-KgaffTFR"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.github/workflows/deploy.yml

@@ -0,0 +1,30 @@
+name: Deploy to Hugging Face Space
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout Repo
+        uses: actions/checkout@v4
+
+      - name: Push to Hugging Face
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          git config --global user.email "github-actions@github.com"
+          git config --global user.name "github-actions"
+
+          git clone https://user:$HF_TOKEN@huggingface.co/spaces/greedybeserk95/Compilertetris space
+
+          rsync -av --exclude ".git" ./ space/
+
+          cd space
+          git add .
+          git commit -m "Auto deploy from GitHub" || echo "No changes"
+          git push

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md

@@ -1,11 +1 @@
----
-title: Compilertetris
-emoji: ⚡
-colorFrom: red
-colorTo: gray
-sdk: docker
-pinned: false
-short_description: an RL environment which optimizes IR code
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MetaHackathon2026Finals

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/space/README.md

@@ -0,0 +1,11 @@
+---
+title: Compilertetris
+emoji: ⚡
+colorFrom: red
+colorTo: gray
+sdk: docker
+pinned: false
+short_description: an RL environment which optimizes IR code
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

space/space/space/space/space/space/space/space/space/space/train.py

@@ -0,0 +1,208 @@
+"""
+Toy REINFORCE training for CompilerOptimizationEnv (CPU, no PyTorch).
+
+This is meant for the Hugging Face Space: fast smoke training that uses the
+same `runtime_core` mock engine and passes as the Gradio demo.
+
+For full GRPO + Unsloth + LLM, use your Colab / GPU notebooks
+(`compiler_optimization_grpo.ipynb`, `role2_deliverable3_*`).
+
+CLI:  python train.py --episodes 50 --max-steps 8 --seed 0
+Or import:  from train import run_toy_training; print(run_toy_training(20, 8, 0))
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import textwrap
+from typing import List, Sequence, Tuple
+
+from runtime_core import (
+    CompilerOptimizationEnv,
+    MOCK_PASSES,
+    MockEngine,
+    SAMPLE_PROGRAM,
+)
+
+# Additional tiny programs so the policy is not overfit to a single IR.
+TRAINING_PROGRAMS: List[List[dict]] = [
+    SAMPLE_PROGRAM,
+    [
+        {"op": "const", "dest": "a", "args": ["1"], "type": "int"},
+        {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"},
+        {"op": "ret", "args": ["b"]},
+    ],
+    [
+        {"op": "const", "dest": "x", "args": ["2"], "type": "int"},
+        {"op": "const", "dest": "y", "args": ["4"], "type": "int"},
+        {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"},
+        {"op": "const", "dest": "k", "args": ["1"], "type": "int"},
+        {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"},
+        {"op": "ret", "args": ["w"]},
+    ],
+]
+
+ACTIONS: Tuple[str, ...] = tuple(sorted(MOCK_PASSES.keys()))
+
+
+def _softmax(logits: Sequence[float]) -> List[float]:
+    m = max(logits) if logits else 0.0
+    ex = [math.exp(x - m) for x in logits]
+    s = sum(ex) or 1.0
+    return [e / s for e in ex]
+
+
+def _sample_action(rng: random.Random, logits: List[float]) -> Tuple[int, List[float], float]:
+    """Return (action index, prob vector, log prob of chosen action)."""
+    p = _softmax(logits)
+    u = rng.random()
+    acc = 0.0
+    idx = len(p) - 1
+    for i, pi in enumerate(p):
+        acc += pi
+        if u <= acc:
+            idx = i
+            break
+    log_p = math.log(p[idx] + 1e-12)
+    return idx, p, log_p
+
+
+def _rollout(
+    engine: MockEngine,
+    program: List[dict],
+    max_steps: int,
+    logits: List[float],
+    rng: random.Random,
+) -> Tuple[float, List[Tuple[int, List[float], float]]]:
+    """
+    One episode. Returns total reward and per-step (action index, prob vector, step reward)
+    for REINFORCE with returns G_t = sum of rewards from t onward.
+    """
+    env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=max_steps)
+    env.reset(program)
+    total = 0.0
+    trace: List[Tuple[int, List[float], float]] = []
+    for _ in range(max_steps):
+        a_idx, pvec, _ = _sample_action(rng, logits)
+        action = ACTIONS[a_idx]
+        step = env.step(action)
+        r = float(step.reward)
+        total += r
+        trace.append((a_idx, pvec, r))
+        if step.done:
+            break
+    return total, trace
+
+
+def _reinforce_update(
+    logits: List[float], trace: List[Tuple[int, List[float], float]], lr: float
+) -> None:
+    """REINFORCE with Monte Carlo return G_t from each step."""
+    G = 0.0
+    for t in range(len(trace) - 1, -1, -1):
+        a_idx, pvec, r = trace[t]
+        G += r
+        for i in range(len(logits)):
+            delta = 1.0 if i == a_idx else 0.0
+            logits[i] += lr * G * (delta - pvec[i])
+
+
+def run_toy_training(
+    episodes: int,
+    max_steps: int,
+    seed: int = 0,
+    lr: float = 0.15,
+) -> str:
+    """
+    Train a stateless categorical policy over the three mock passes; print-friendly report.
+    """
+    if episodes < 1:
+        return "episodes must be >= 1"
+    if max_steps < 1:
+        return "max_steps must be >= 1"
+    episodes = int(episodes)
+    max_steps = int(max_steps)
+    seed = int(seed)
+    lr = float(lr)
+
+    rng = random.Random(seed)
+    engine = MockEngine()
+    logits = [0.0 for _ in ACTIONS]
+
+    history: List[Tuple[int, float]] = []
+    for ep in range(1, episodes + 1):
+        program = TRAINING_PROGRAMS[(ep - 1) % len(TRAINING_PROGRAMS)]
+        G, trace = _rollout(engine, program, max_steps, logits, rng)
+        if trace:
+            _reinforce_update(logits, trace, lr)
+        history.append((ep, G))
+
+    final_p = _softmax(logits)
+    lines = [
+        "Toy REINFORCE (stateless policy over pass names, CPU, stdlib only)",
+        f"  episodes={episodes}  max_steps={max_steps}  seed={seed}  lr={lr}",
+        f"  actions order: {list(ACTIONS)}",
+        "",
+        f"  final logits: {[round(x, 4) for x in logits]}",
+        f"  final policy:  {', '.join(f'{a}={p:.3f}' for a, p in zip(ACTIONS, final_p))}",
+        "",
+        "  return per episode (last 10): " + ", ".join(f"{G:+.1f}" for _, G in history[-10:]),
+        "",
+    ]
+    if history:
+        mean_r = sum(G for _, G in history) / len(history)
+        lines.append(f"  mean return over all episodes: {mean_r:+.3f}")
+    lines.append("")
+    lines.append(
+        textwrap.dedent(
+            """
+            This does not train an LLM. For GRPO + Qwen + Unsloth, run the project notebooks
+            on a GPU machine (e.g. Colab), not the CPU Space.
+            """
+        ).strip()
+    )
+    return "\n".join(lines)
+
+
+def _parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Toy REINFORCE on CompilerOptimizationEnv (CPU)")
+    p.add_argument("--episodes", type=int, default=50, help="Number of training episodes")
+    p.add_argument("--max-steps", type=int, default=8, help="max_steps per env episode")
+    p.add_argument("--seed", type=int, default=0, help="RNG seed")
+    p.add_argument("--lr", type=float, default=0.15, help="REINFORCE learning rate")
+    p.add_argument(
+        "--out-json",
+        type=str,
+        default="",
+        help="If set, write a small run summary to this path (e.g. training_log.json).",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    report = run_toy_training(
+        episodes=args.episodes,
+        max_steps=args.max_steps,
+        seed=args.seed,
+        lr=args.lr,
+    )
+    print(report)
+    if args.out_json:
+        payload = {
+            "episodes": args.episodes,
+            "max_steps": args.max_steps,
+            "seed": args.seed,
+            "lr": args.lr,
+            "summary_text": report,
+        }
+        with open(args.out_json, "w", encoding="utf-8") as f:
+            json.dump(payload, f, indent=2)
+        print(f"\nWrote {args.out_json}")
+
+
+if __name__ == "__main__":
+    main()

space/space/space/space/space/space/space/space/untitled folder.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b29f9a98dd93fb2a14b1d3e8409ce51a756b578debde97af1190991bbb65d203
+size 63307

space/space/space/space/space/space/write_colab_synth_notebook.py

@@ -0,0 +1,345 @@
+"""Generate Compilertetris_GRPO_synthetic_dataset.ipynb (GRPO + random Toy-IR corpus)."""
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+OUT = ROOT / "Compilertetris_GRPO_synthetic_dataset.ipynb"
+
+
+def cell_md(s: str) -> dict:
+    return {
+        "cell_type": "markdown",
+        "metadata": {},
+        "source": [line + "\n" for line in s.strip().split("\n")],
+    }
+
+
+def cell_code(s: str) -> dict:
+    return {
+        "cell_type": "code",
+        "metadata": {},
+        "execution_count": None,
+        "outputs": [],
+        "source": [line + "\n" for line in s.rstrip().split("\n")],
+    }
+
+cells: list = []
+
+cells.append(cell_md("""
+# Compiler Tetris — GRPO with **synthetic Toy-IR corpus** (Colab)
+
+This variant builds **50–200+** training prompts from `program_generator.py` (random const/add/mul chains in the same JSON shape as `runtime_core` / Deliverable 2). The original `metahack1` notebook uses a *different* IR schema; that generator is not mixed in here unless you add a converter.
+
+| What | Value |
+|------|--------|
+| Space | `greedybeserk95/Compilertetris` |
+| Code root after clone | `/content/Compilertetris` |
+| LoRA output | `/content/compilertetris_lora` |
+| This run uses **shorter** `TRAIN_STEPS` than the 3-example notebook (see config cell). |
+| **Checkpoints** | Under `output_dir` as `checkpoint-*` (see `CHECKPOINT_EVERY` / `KEEP_LAST_N_CHECKPOINTS` in the GRPO cell) |
+
+**Tuning:** In the dataset cell, set `N_TRAIN_PROGRAMS` (e.g. 80, 120, 200).
+"""))
+
+cells.append(
+    cell_md(
+        """
+## Checkpoints
+
+The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `/content/grpo_compilertetris/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training).
+"""
+    )
+)
+
+cells.append(cell_code("""
+# --- Central path config ---
+
+HF_SPACE_REPO = "https://huggingface.co/spaces/greedybeserk95/Compilertetris"
+HF_ADAPTER_REPO = "greedybeserk95/Compilertetris-grpo-adapter"
+REPO_DIR = "/content/Compilertetris"
+"""))
+
+cells.append(cell_code("""
+import os, subprocess
+
+if os.path.isdir(REPO_DIR + "/.git"):
+    subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"], timeout=300)
+else:
+    subprocess.check_call(["git", "clone", HF_SPACE_REPO, REPO_DIR], timeout=600)
+print("Repo:", REPO_DIR)
+"""))
+
+cells.append(cell_code("""!nvidia-smi"""))
+
+cells.append(cell_code("""
+!pip install -q "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
+!pip install -q trl datasets transformers accelerate peft bitsandbytes
+"""))
+
+cells.append(cell_code("""
+import sys, json, os
+sys.path.insert(0, REPO_DIR)
+
+import torch
+import unsloth
+from unsloth import FastLanguageModel
+from datasets import Dataset
+
+from runtime_core import (
+    CompilerOptimizationEnv,
+    Deliverable2_Formatter,
+    MockEngine,
+    MOCK_PASSES,
+)
+from program_generator import build_training_program_corpus
+
+print("torch", torch.__version__, "cuda", torch.cuda.is_available())
+print("MOCK_PASSES", list(MOCK_PASSES.keys()))
+"""))
+
+
+def _cell_chatml_synth() -> str:
+    p0 = """# Qwen2.5 ChatML + synthetic PROGRAMS
+IM_END = \"<\" + \"|im_end|>\"
+"""
+    p1 = r'''
+import re
+
+# How many training programs (50–200 recommended; includes 3 builtins if include_builtins=True)
+N_TRAIN_PROGRAMS = 120
+RANDOM_SEED = 42
+
+SYSTEM_PROMPT = f"""You are a compiler optimization agent.
+You will see Toy-IR as pseudo-assembly (Deliverable-2 text).
+
+Output ONLY a valid JSON array of optimization pass names IN ORDER.
+Available passes: {", ".join(repr(p) for p in sorted(MOCK_PASSES))}
+Rules: JSON array only; max 8 passes; you may repeat passes."""
+
+
+def build_prompt(program_list: list, program_id: int) -> str:
+    obs = Deliverable2_Formatter.translate_state(program_list)
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}{IM_END}\n"
+        f"<|im_start|>user\n#PROGRAM_ID:{program_id}\n{obs}{IM_END}\n"
+        f"<|im_start|>assistant\n"
+    )
+
+
+def program_from_prompt(prompt: str) -> list:
+    m = re.search(r"#PROGRAM_ID:(\d+)", prompt)
+    if m:
+        i = int(m.group(1))
+        if 0 <= i < len(PROGRAMS):
+            return PROGRAMS[i]
+    if "#PROGRAM_JSON" in prompt:
+        tail = prompt.split("#PROGRAM_JSON", 1)[1]
+        if IM_END in tail:
+            tail = tail.split(IM_END, 1)[0]
+        raw = tail.strip()
+        if raw:
+            return json.loads(raw)
+    raise ValueError("cannot resolve program from prompt (expect #PROGRAM_ID:N or #PROGRAM_JSON)")
+
+
+PROGRAMS = build_training_program_corpus(
+    n_total=N_TRAIN_PROGRAMS,
+    seed=RANDOM_SEED,
+    include_builtins=True,
+)
+print("Corpus size:", len(PROGRAMS))
+
+train_dataset = Dataset.from_dict({
+    "prompt": [build_prompt(p, i) for i, p in enumerate(PROGRAMS)],
+})
+print("Dataset rows:", len(train_dataset))
+'''.lstrip("\n")
+    return p0 + p1
+
+
+cells.append(cell_code(_cell_chatml_synth()))
+
+cells.append(cell_code("""
+def env_reward_for_completion(prompt: str, completion: str, max_env_steps: int = 8) -> float:
+    try:
+        program = program_from_prompt(prompt)
+    except Exception:
+        return -8.0
+    try:
+        actions = Deliverable2_Formatter.extract_action_array(completion)
+    except Exception:
+        return -5.0
+    actions = [str(a).strip() for a in actions][: max_env_steps]
+    if not actions:
+        return -4.0
+    env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=max_env_steps)
+    env.reset(program)
+    total = 0.0
+    for a in actions:
+        step = env.step(a)
+        total += float(step.reward)
+        if step.done:
+            break
+    return float(total)
+
+
+def make_reward_function(max_env_steps: int = 8):
+    def reward_func(prompts: list, completions: list, **kwargs) -> list:
+        return [env_reward_for_completion(p, c, max_env_steps) for p, c in zip(prompts, completions)]
+    return reward_func
+"""))
+
+cells.append(cell_code("""
+from trl import GRPOConfig, GRPOTrainer
+
+MODEL_NAME = "unsloth/Qwen2.5-3B-Instruct"
+MAX_SEQ_LEN = 1024
+MAX_COMPLETION = 256
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=MODEL_NAME,
+    max_seq_length=MAX_SEQ_LEN,
+    dtype=None,
+    load_in_4bit=True,
+)
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=16, lora_alpha=16, lora_dropout=0.0, bias="none",
+    use_gradient_checkpointing="unsloth", random_state=0,
+)
+"""))
+
+# Shorter run than 200 steps — scale up after smoke test
+cells.append(
+    cell_code(
+        """
+NUM_GENERATIONS = 2
+LEARNING_RATE = 2e-5
+# Reduced from 200: increase after you verify loss/reward is stable
+TRAIN_STEPS = 80
+
+# --- Checkpoints (TrainingArguments / GRPO) ---
+GRPO_OUTPUT_DIR = "/content/grpo_compilertetris"
+CHECKPOINT_EVERY = 20   # save a checkpoint every N global steps
+KEEP_LAST_N_CHECKPOINTS = 5  # on disk; older folders are deleted
+
+grpo_config = GRPOConfig(
+    output_dir=GRPO_OUTPUT_DIR,
+    learning_rate=LEARNING_RATE,
+    per_device_train_batch_size=NUM_GENERATIONS,
+    gradient_accumulation_steps=1,
+    num_generations=NUM_GENERATIONS,
+    max_completion_length=MAX_COMPLETION,
+    max_prompt_length=MAX_SEQ_LEN,
+    remove_unused_columns=False,
+    temperature=0.7,
+    max_steps=TRAIN_STEPS,
+    logging_steps=5,
+    save_strategy="steps",
+    save_steps=CHECKPOINT_EVERY,
+    save_total_limit=KEEP_LAST_N_CHECKPOINTS,
+    seed=0,
+    report_to="none",
+    use_vllm=False,
+)
+
+reward_fn = make_reward_function()
+
+trainer = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=[reward_fn],
+    args=grpo_config,
+    train_dataset=train_dataset,
+)
+"""
+    )
+)
+
+cells.append(cell_code("""
+print("Starting GRPO (synthetic corpus)…")
+trainer.train()
+print("Done.")
+"""))
+
+cells.append(
+    cell_code(
+        """
+# List on-disk checkpoints (for resume or manual export)
+import glob, os
+
+ckpts = sorted(
+    glob.glob(os.path.join(GRPO_OUTPUT_DIR, "checkpoint-*")),
+    key=lambda p: int(p.split("checkpoint-")[-1]) if p.split("checkpoint-")[-1].isdigit() else 0,
+)
+print("Checkpoints in", GRPO_OUTPUT_DIR, ":", len(ckpts))
+for c in ckpts:
+    print(" ", c)
+if ckpts:
+    print("Latest:", ckpts[-1])
+"""
+    )
+)
+
+cells.append(
+    cell_md(
+        """
+**Resume after disconnect / crash** — re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of:
+
+- `trainer.train(resume_from_checkpoint=True)` — continues from the latest `checkpoint-*` in `output_dir`
+- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-40")` — specific step
+"""
+    )
+)
+
+cells.append(
+    cell_code(
+        """
+# Uncomment to resume from the latest checkpoint (run after re-creating `trainer` in a new session)
+# trainer.train(resume_from_checkpoint=True)
+"""
+    )
+)
+
+cells.append(
+    cell_code(
+        """
+SAVE_DIR = "/content/compilertetris_lora"
+trainer.model.save_pretrained(SAVE_DIR)
+tokenizer.save_pretrained(SAVE_DIR)
+print("Saved to", SAVE_DIR)
+"""
+    )
+)
+
+cells.append(
+    cell_code(
+        """
+from huggingface_hub import login, HfApi
+
+login()
+HfApi().create_repo(HF_ADAPTER_REPO, exist_ok=True, repo_type="model")
+trainer.model.push_to_hub(HF_ADAPTER_REPO, private=True)
+tokenizer.push_to_hub(HF_ADAPTER_REPO, private=True)
+print("Pushed to https://huggingface.co/" + HF_ADAPTER_REPO)
+"""
+    )
+)
+
+if __name__ == "__main__":
+    nb = {
+        "nbformat": 4,
+        "nbformat_minor": 5,
+        "metadata": {
+            "colab": {"provenance": [], "gpuType": "T4"},
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3",
+            },
+            "language_info": {"name": "python"},
+        },
+        "cells": cells,
+    }
+    OUT.write_text(json.dumps(nb, indent=2), encoding="utf-8")
+    print("Wrote", OUT)

space/space/space/space/write_colab_synth_notebook.py

@@ -26,19 +26,21 @@ def cell_code(s: str) -> dict:
 cells: list = []
 
 cells.append(cell_md("""
-# Compiler Tetris — GRPO with **synthetic Toy-IR corpus** (Colab)
+# Compiler Tetris — GRPO, **synthetic Toy-IR** (Colab **T4 ~8–10 h** preset)
 
-This variant builds **50–200+** training prompts from `program_generator.py` (random const/add/mul chains in the same JSON shape as `runtime_core` / Deliverable 2). The original `metahack1` notebook uses a *different* IR schema; that generator is not mixed in here unless you add a converter.
+**Runtime (estimate):** after ~20 steps, read `it/s` in the progress bar. Approximate hours ≈ `TRAIN_STEPS / (it/s * 3600)`.
+With **`TRAIN_STEPS = 3600`**: at **0.10 it/s** that is about **10 h**; at **0.12 it/s** about **8.3 h**. If you finish much faster/slower, change **`TRAIN_STEPS`** (or the model size / `NUM_GENERATIONS`).
+
+**T4 tips:** 3B + 4bit + below settings fit T4; **7B** or very long `MAX_COMPLETION` can OOM. **High-RAM** runtime in Colab helps. Use **checkpoints** (next cells) in case the session dies.
 
 | What | Value |
 |------|--------|
 | Space | `greedybeserk95/Compilertetris` |
-| Code root after clone | `/content/Compilertetris` |
-| LoRA output | `/content/compilertetris_lora` |
-| This run uses **shorter** `TRAIN_STEPS` than the 3-example notebook (see config cell). |
-| **Checkpoints** | Under `output_dir` as `checkpoint-*` (see `CHECKPOINT_EVERY` / `KEEP_LAST_N_CHECKPOINTS` in the GRPO cell) |
+| Code root | `/content/Compilertetris` |
+| LoRA out | `/content/compilertetris_lora` |
+| **Checkpoints** | `/content/grpo_compilertetris/checkpoint-*` |
 
-**Tuning:** In the dataset cell, set `N_TRAIN_PROGRAMS` (e.g. 80, 120, 200).
+`program_generator` — Toy-IR in `runtime_core` shape; `metahack1` uses a different schema.
 """))
 
 cells.append(
@@ -105,8 +107,8 @@ IM_END = \"<\" + \"|im_end|>\"
     p1 = r'''
 import re
 
-# How many training programs (50–200 recommended; includes 3 builtins if include_builtins=True)
-N_TRAIN_PROGRAMS = 120
+# T4: 200–300 is a good tradeoff (RAM + diversity)
+N_TRAIN_PROGRAMS = 250
 RANDOM_SEED = 42
 
 SYSTEM_PROMPT = f"""You are a compiler optimization agent.
@@ -183,7 +185,7 @@ def env_reward_for_completion(prompt: str, completion: str, max_env_steps: int =
     return float(total)
 
 
-def make_reward_function(max_env_steps: int = 8):
+def make_reward_function(max_env_steps: int = 10):
     def reward_func(prompts: list, completions: list, **kwargs) -> list:
         return [env_reward_for_completion(p, c, max_env_steps) for p, c in zip(prompts, completions)]
     return reward_func
@@ -192,6 +194,7 @@ def make_reward_function(max_env_steps: int = 8):
 cells.append(cell_code("""
 from trl import GRPOConfig, GRPOTrainer
 
+# --- T4-friendly (3B 4-bit). For A100+ you can try Qwen2.5-7B and MAX_COMPLETION=384. ---
 MODEL_NAME = "unsloth/Qwen2.5-3B-Instruct"
 MAX_SEQ_LEN = 1024
 MAX_COMPLETION = 256
@@ -204,37 +207,40 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 )
 model = FastLanguageModel.get_peft_model(
     model,
-    r=16, lora_alpha=16, lora_dropout=0.0, bias="none",
+    r=32, lora_alpha=32, lora_dropout=0.0, bias="none",
     use_gradient_checkpointing="unsloth", random_state=0,
 )
+print("Model:", MODEL_NAME, "| max_seq", MAX_SEQ_LEN, "| completion cap", MAX_COMPLETION)
 """))
 
-# Shorter run than 200 steps — scale up after smoke test
+# ---- T4 / ~8–10 h wall time: tune TRAIN_STEPS after you see it/s in the first minutes ----
 cells.append(
     cell_code(
         """
-NUM_GENERATIONS = 2
-LEARNING_RATE = 2e-5
-# Reduced from 200: increase after you verify loss/reward is stable
-TRAIN_STEPS = 80
+# 4 rollouts per prompt: good for GRPO on T4; 6–8 is heavier (slower, more VRAM)
+NUM_GENERATIONS = 4
+LEARNING_RATE = 1.5e-5
+# Target ~8–10 h on T4 when it/s is ~0.10–0.12 (typical for this stack). Re-tune if your it/s differs.
+#  hours ≈ TRAIN_STEPS / (it/s * 3600)
+TRAIN_STEPS = 3600
+GRAD_ACCUM = 2
 
-# --- Checkpoints (TrainingArguments / GRPO) ---
 GRPO_OUTPUT_DIR = "/content/grpo_compilertetris"
-CHECKPOINT_EVERY = 20   # save a checkpoint every N global steps
-KEEP_LAST_N_CHECKPOINTS = 5  # on disk; older folders are deleted
+CHECKPOINT_EVERY = 200
+KEEP_LAST_N_CHECKPOINTS = 3
 
 grpo_config = GRPOConfig(
     output_dir=GRPO_OUTPUT_DIR,
     learning_rate=LEARNING_RATE,
     per_device_train_batch_size=NUM_GENERATIONS,
-    gradient_accumulation_steps=1,
+    gradient_accumulation_steps=GRAD_ACCUM,
     num_generations=NUM_GENERATIONS,
     max_completion_length=MAX_COMPLETION,
     max_prompt_length=MAX_SEQ_LEN,
     remove_unused_columns=False,
     temperature=0.7,
     max_steps=TRAIN_STEPS,
-    logging_steps=5,
+    logging_steps=20,
     save_strategy="steps",
     save_steps=CHECKPOINT_EVERY,
     save_total_limit=KEEP_LAST_N_CHECKPOINTS,
@@ -287,7 +293,7 @@ cells.append(
 **Resume after disconnect / crash** — re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of:
 
 - `trainer.train(resume_from_checkpoint=True)` — continues from the latest `checkpoint-*` in `output_dir`
-- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-40")` — specific step
+- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-300")` — example path (use an existing `checkpoint-*` folder)
 """
     )
 )

write_colab_synth_notebook.py

@@ -36,11 +36,13 @@ With **`TRAIN_STEPS = 3600`**: at **0.10 it/s** that is about **10 h**; at **0.1
 | What | Value |
 |------|--------|
 | Space | `greedybeserk95/Compilertetris` |
-| Code root | `/content/Compilertetris` |
-| LoRA out | `/content/compilertetris_lora` |
-| **Checkpoints** | `/content/grpo_compilertetris/checkpoint-*` |
+| Code root | `REPO_DIR` (printed in the “paths” cell) |
+| LoRA out | `LORA_DIR` (printed in the “paths” cell) |
+| **Checkpoints** | `GRPO_OUTPUT_DIR + "/checkpoint-*"` |
 
 `program_generator` — Toy-IR in `runtime_core` shape; `metahack1` uses a different schema.
+
+**Training evidence (loss + reward plots):** the cell *after* `trainer.train()` saves `training_loss_and_reward.png` under `output_dir` and shows it in the notebook — use it in your README / writeup.
 """))
 
 cells.append(
@@ -48,21 +50,49 @@ cells.append(
         """
 ## Checkpoints
 
-The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `/content/grpo_compilertetris/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training).
+The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `$GRPO_OUTPUT_DIR/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training).
 """
     )
 )
 
 cells.append(cell_code("""
-# --- Central path config ---
+# --- Central path config (works on Colab *and* local / HF runtimes) ---
+
+import os
+import tempfile
+from pathlib import Path
+
+def _default_workspace_base() -> str:
+    # Colab: /content is writable
+    c = "/content"
+    if os.path.isdir(c) and os.access(c, os.W_OK):
+        return str(Path(c) / "work")
+
+    # Otherwise: a guaranteed-writable temp dir
+    return str(Path(tempfile.gettempdir()) / "compilertetris_work")
+
+
+BASE = os.environ.get("COMPILERTETRIS_BASE", _default_workspace_base())
+os.makedirs(BASE, exist_ok=True)
 
 HF_SPACE_REPO = "https://huggingface.co/spaces/greedybeserk95/Compilertetris"
 HF_ADAPTER_REPO = "greedybeserk95/Compilertetris-grpo-adapter"
-REPO_DIR = "/content/Compilertetris"
+REPO_DIR = str(Path(BASE) / "Compilertetris")
+GRPO_OUTPUT_DIR = str(Path(BASE) / "grpo_compilertetris")
+LORA_DIR = str(Path(BASE) / "compilertetris_lora")
+
+print("BASE        :", BASE)
+print("REPO_DIR    :", REPO_DIR)
+print("GRPO_OUTPUT :", GRPO_OUTPUT_DIR)
+print("LORA_DIR    :", LORA_DIR)
 """))
 
 cells.append(cell_code("""
 import os, subprocess
+from pathlib import Path
+
+# Ensure parent is writable/created
+Path(REPO_DIR).parent.mkdir(parents=True, exist_ok=True)
 
 if os.path.isdir(REPO_DIR + "/.git"):
     subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"], timeout=300)
@@ -74,16 +104,22 @@ print("Repo:", REPO_DIR)
 cells.append(cell_code("""!nvidia-smi"""))
 
 cells.append(cell_code("""
-!pip install -q "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
-!pip install -q trl datasets transformers accelerate peft bitsandbytes
+# One-shot install (A100 + CUDA 12.x):
+# - Unsloth's `cu124-ampere-torch260` extra pulls a **Torch 2.6** CUDA stack (fixes `torch.int1` / torchao mismatches)
+# - Pin TRL+Transformers to a Unsloth-tested pair (per Unsloth issue threads around torch2.6 / transformers4.55.x)
+!python -m pip install -U --no-cache-dir \
+  "transformers==4.55.4" "trl==0.20.0" \
+  matplotlib datasets accelerate peft bitsandbytes safetensors huggingface_hub sentencepiece \
+  "unsloth[cu124-ampere-torch260] @ git+https://github.com/unslothai/unsloth.git"
 """))
 
 cells.append(cell_code("""
+# `import unsloth` should come before `transformers` is imported (happens via other imports).
+import unsloth
 import sys, json, os
 sys.path.insert(0, REPO_DIR)
 
 import torch
-import unsloth
 from unsloth import FastLanguageModel
 from datasets import Dataset
 
@@ -225,7 +261,7 @@ LEARNING_RATE = 1.5e-5
 TRAIN_STEPS = 3600
 GRAD_ACCUM = 2
 
-GRPO_OUTPUT_DIR = "/content/grpo_compilertetris"
+# GRPO_OUTPUT_DIR is set in the "Central path config" cell
 CHECKPOINT_EVERY = 200
 KEEP_LAST_N_CHECKPOINTS = 3
 
@@ -268,6 +304,101 @@ trainer.train()
 print("Done.")
 """))
 
+cells.append(
+    cell_code(
+        r"""
+# --- Loss + reward plots (submission): from trainer.state.log_history after train() ---
+
+!pip install -q matplotlib
+
+import os
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+PLOT_DIR = GRPO_OUTPUT_DIR
+os.makedirs(PLOT_DIR, exist_ok=True)
+PNG_PATH = os.path.join(PLOT_DIR, "training_loss_and_reward.png")
+
+def _reward_from_log_row(h):
+    v = h.get("reward")
+    if isinstance(v, (int, float)):
+        return float(v), "reward"
+    best_k, best_v = None, None
+    for k, v in h.items():
+        if not isinstance(v, (int, float)):
+            continue
+        klow = k.replace("-", "/").lower()
+        if "reward" not in klow or "std" in klow or "clip" in klow or "per_token" in klow:
+            continue
+        if "mean" in klow or k == "reward":
+            return float(v), k
+        if "mean" not in klow and best_k is None:
+            best_k, best_v = k, float(v)
+    if best_k is not None:
+        return best_v, best_k
+    return None, None
+
+def extract_series(history):
+    sl, vl, sr, vr = [], [], [], []
+    rlabel = None
+    for h in history:
+        s = h.get("step")
+        if s is None:
+            continue
+        lo = h.get("loss")
+        if isinstance(lo, (int, float)):
+            sl.append(s)
+            vl.append(float(lo))
+        r_val, rk = _reward_from_log_row(h)
+        if r_val is not None and rk:
+            if rlabel is None:
+                rlabel = rk
+            if rk == rlabel:
+                sr.append(s)
+                vr.append(r_val)
+    return (sl, vl, "loss"), (sr, vr, rlabel or "reward")
+
+(loss_s, loss_v, _lk), (rew_s, rew_v, rew_lab) = extract_series(trainer.state.log_history)
+print("Points — loss:", len(loss_v), "| reward:", len(rew_v))
+if trainer.state.log_history:
+    print("Last log row keys (sample):", list(trainer.state.log_history[-1].keys())[:25])
+
+fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
+if loss_v:
+    ax0.plot(loss_s, loss_v, "b.-", label="loss", linewidth=1, markersize=2)
+    ax0.set_ylabel("training loss")
+    ax0.set_title("GRPO (this Colab run)")
+    ax0.grid(True, alpha=0.3)
+    ax0.legend()
+else:
+    ax0.text(0.5, 0.5, "No 'loss' in log_history", ha="center", transform=ax0.transAxes)
+
+if rew_v:
+    ax1.plot(rew_s, rew_v, "g.-", label=rew_lab, linewidth=1, markersize=2)
+    ax1.set_ylabel("mean reward" if "mean" in (rew_lab or "") else "reward")
+    ax1.set_xlabel("global step")
+    ax1.grid(True, alpha=0.3)
+    ax1.legend()
+else:
+    ax1.text(0.5, 0.5, "No reward column found — see keys above", ha="center", transform=ax1.transAxes)
+    ax1.set_xlabel("global step")
+
+plt.tight_layout()
+plt.savefig(PNG_PATH, dpi=150, bbox_inches="tight")
+print("Saved:", PNG_PATH)
+try:
+    from IPython.display import Image, display
+    display(Image(PNG_PATH))
+except Exception as e:
+    print("Display:", e)
+finally:
+    plt.close("all")
+print("If keys differ:  print(trainer.state.log_history[-1])")
+"""
+    )
+)
+
 cells.append(
     cell_code(
         """
@@ -293,7 +424,7 @@ cells.append(
 **Resume after disconnect / crash** — re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of:
 
 - `trainer.train(resume_from_checkpoint=True)` — continues from the latest `checkpoint-*` in `output_dir`
-- `trainer.train(resume_from_checkpoint="/content/grpo_compilertetris/checkpoint-300")` — example path (use an existing `checkpoint-*` folder)
+- `trainer.train(resume_from_checkpoint="<your GRPO_OUTPUT_DIR>/checkpoint-300")` — example (use an existing `checkpoint-*` folder; print `GRPO_OUTPUT_DIR` from the path cell)
 """
     )
 )
@@ -310,7 +441,7 @@ cells.append(
 cells.append(
     cell_code(
         """
-SAVE_DIR = "/content/compilertetris_lora"
+SAVE_DIR = LORA_DIR
 trainer.model.save_pretrained(SAVE_DIR)
 tokenizer.save_pretrained(SAVE_DIR)
 print("Saved to", SAVE_DIR)