rlm_forge_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RLM-Forge: Training LLMs with GRPO on Coding Tasks\n",
    "\n",
    "**RLM-Forge** is an OpenEnv environment that trains language models to solve coding tasks using Recursive Language Model (RLM) patterns. The environment:\n",
    "\n",
    "1. Clones a Python repository\n",
    "2. Replaces a source file with a broken stub (correct signatures, wrong implementations)\n",
    "3. Provides a sandboxed REPL with tools (read_file, list_dir, search, write_file, run_tests)\n",
    "4. The agent must explore the repo, understand the tests, and rewrite the source file\n",
    "5. Reward = test pass rate (55%) + structural validity (15%) + efficiency (30%)\n",
    "\n",
    "This notebook trains a model using **GRPO (Group Relative Policy Optimization)** with multi-step trajectory concatenation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup & Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Install dependencies\n",
    "!pip install -q \"openenv-core[core]>=0.2.0\" trl transformers accelerate bitsandbytes peft datasets\n",
    "!pip install -q text-unidecode freezegun pytest vllm\n",
    "\n",
    "# Clone RLM-Forge repo\n",
    "!git clone https://github.com/kking112/rlm-forge.git content/rlm-forge 2>/dev/null || true\n",
    "# Or upload files manually — adjust path as needed\n",
    "# import sys\n",
    "# sys.path.insert(0, \"content/rlm-forge\")\n",
    "\n",
    "# Install RLM-Forge\n",
    "!pip install -q -e content/rlm-forge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU: NVIDIA RTX PRO 6000 Blackwell Workstation Edition\n",
      "PyTorch: 2.10.0+cu128\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import json\n",
    "import re\n",
    "import random\n",
    "from typing import Optional\n",
    "from dataclasses import dataclass\n",
    "\n",
    "print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "# print(f\"VRAM: {torch.cuda.get_device_properties(0). / 1e9:.1f} GB\")\n",
    "print(f\"PyTorch: {torch.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Environment Smoke Test\n",
    "\n",
    "Verify the environment works before training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task: The file `slugify/slugify.py` has been replaced with a broken stub. 82 tests in `test.py` are now failing. Your task is to explore the repository, understand the expected behavior from the tests and o...\n",
      "Available tools: ['read_file(path)', \"list_dir(path='.')\", \"search(pattern, path='.')\", 'write_file(path, content)', 'run_tests(test_path=None)', 'spawn_agent(scope, mission, budget=5)', 'FINAL()']\n",
      "\n",
      "Step 1 stdout: ['.git/', '.github/', '.gitignore', '.pytest_cache/', '.vscode/', 'CHANGELOG.md', 'LICENSE', 'MANIFEST.in', 'README.md', '__pycache__/', 'dev.requirements.txt', 'format.sh', 'pyproject.toml', 'python_\n",
      "\n",
      "Baseline reward (no implementation): 0.3939\n",
      "Test results: {'total_reward': 0.3939, 'test_pass_rate': 0.1707, 'tests_passed': 14, 'tests_failed': 68, 'tests_total': 82, 'structural_score': 0.0, 'efficiency_score': 1.0, 'breakdown': {'test_component': 0.0939, 'structural_component': 0.0, 'efficiency_component': 0.3}, 'test_output': '============================= test session starts ==============================\\ncollecting ... collected 82 items\\n\\ntest.py::TestSlugify::test_accented_text FAILED                          [  1%]\\ntest.py::TestSlugify::test_accented_text_with_non_word_characters FAILED [  2%]\\ntest.py::TestSlugify::test_contains_numbers FAILED                       [  3%]\\ntest.py::TestSlugify::test_custom_separator FAILED                       [  4%]\\ntest.py::TestSlugify::test_cyrillic_text FAILED                          [  6%]\\ntest.py::TestSlugify::test_differently_cased_stopword_match FAILED       [  7%]\\ntest.py::TestSlugify::test_ends_with_number FAILED                       [  8%]\\ntest.py::TestSlugify::test_extraneous_seperators FAILED                  [  9%]\\ntest.py::TestSlugify::test_html_decimal_off FAILED                       [ 10%]\\ntest.py::TestSlugify::test_html_decimal_on FAILED                        [ 12%]\\ntest.py::TestSlugify::test_html_entities_off FAILED                      [ 13%]\\ntest.py::TestSlugify::test_html_entities_on FAILED                       [ 14%]\\ntest.py::TestSlugify::test_html_hexadecimal_off FAILED                   [ 15%]\\ntest.py::TestSlugify::test_html_hexadecimal_on FAILED                    [ 17%]\\ntest.py::TestSlugify::test_max_length FAILED                             [ 18%]\\ntest.py::TestSlugify::test_max_length_cutoff_not_required FAILED         [ 19%]\\ntest.py::TestSlugify::test_multi_character_separator FAILED              [ 20%]\\ntest.py::TestSlugify::test_multiple_stopword_occurances FAILED           [ 21%]\\ntest.py::TestSlugify::test_multiple_stopwords FAILED                     [ 23%]\\ntest.py::TestSlugify::test_non_word_characters FAILED                    [ 24%]\\ntest.py::TestSlugify::test_numbers_and_symbols FAILED                    [ 25%]\\ntest.py::TestSlugify::test_numbers_only FAILED                           [ 26%]\\ntest.py::TestSlugify::test_phonetic_conversion_of_eastern_scripts FAILED [ 28%]\\ntest.py::TestSlugify::test_pre_translation P'}\n"
     ]
    }
   ],
   "source": [
    "from rlm_forge.server.environment import RLMForgeEnvironment\n",
    "from rlm_forge.models import RLMForgeAction\n",
    "\n",
    "env = RLMForgeEnvironment()\n",
    "\n",
    "# Run a quick episode\n",
    "obs = env.reset(seed=1)\n",
    "print(f\"Task: {obs.task_description[:200]}...\")\n",
    "print(f\"Available tools: {obs.available_functions}\")\n",
    "\n",
    "# Take a step — list files\n",
    "obs2 = env.step(RLMForgeAction(code=\"print(list_dir())\"))\n",
    "print(f\"\\nStep 1 stdout: {obs2.stdout[:200]}\")\n",
    "\n",
    "# Finalize and get reward\n",
    "obs3 = env.step(RLMForgeAction(code=\"FINAL()\"))\n",
    "print(f\"\\nBaseline reward (no implementation): {obs3.reward:.4f}\")\n",
    "print(f\"Test results: {obs3.test_results}\")\n",
    "\n",
    "env.cleanup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load Model\n",
    "\n",
    "We use Qwen2.5-Coder-32B-Instruct with 4-bit quantization for inference, and train a LoRA adapter with GRPO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model config — adjust based on available VRAM\n",
    "# MODEL_ID = \"Qwen/Qwen2.5-Coder-32B-Instruct\"  # 32B for H100\n",
    "MODEL_ID = \"Qwen/Qwen2.5-Coder-7B-Instruct\"  # Fallback for smaller GPUs\n",
    "HF_TOKEN = '' #! Fill in HF TOKEN HERE!\n",
    "MAX_STEPS_PER_EPISODE = 6  # Max REPL interactions per episode\n",
    "NUM_EPISODES_PER_PROMPT = 4  # GRPO group size (completions per prompt)\n",
    "NUM_TRAINING_PROMPTS = 16  # Total unique prompts (episodes) for training\n",
    "GRPO_EPOCHS = 2  # Training epochs over collected data\n",
    "BATCH_SIZE = 2\n",
    "GRAD_ACCUM = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "# 4-bit quantization for 32B model on H100\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True,token=HF_TOKEN)\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`torch_dtype` is deprecated! Use `dtype` instead!\n",
      "/home/neo/Desktop/Projects/OpenEnv_Hackathon_SF/V1/.venv/lib/python3.13/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c8654adb1804c6e944f84026e38a81b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 4 files:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac2529cdddcf450ea1d5a50f2cea7814",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00003-of-00004.safetensors:   0%|          | 0.00/4.33G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd5ac4ddfc7b46e2a1ca9515752aa745",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00001-of-00004.safetensors:   0%|          | 0.00/4.88G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a354cc2372c6467eb49761d5fc153940",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00004-of-00004.safetensors:   0%|          | 0.00/1.09G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecca5855d3f647c8b3e49f43214bca5c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model-00002-of-00004.safetensors:   0%|          | 0.00/4.93G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    trust_remote_code=True,\n",
    "    # attn_implementation=\"flash_attention_2\",\n",
    "    token=HF_TOKEN\n",
    ")\n",
    "\n",
    "# LoRA config for efficient training\n",
    "lora_config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Trajectory Collection\n",
    "\n",
    "The key idea: treat the full multi-step episode as one \"completion\" for GRPO.\n",
    "\n",
    "**Prompt** = system message + task description + initial observation\n",
    "**Completion** = sequence of all code actions (with observation feedback between them)\n",
    "**Reward** = final composite reward from the environment\n",
    "\n",
    "We roll out multiple episodes per prompt (GRPO group) and use relative rewards within each group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"You are an expert Python developer. You are given a repository where a source file has been replaced with a broken stub. Your task is to explore the repository, understand the expected behavior from the tests, and rewrite the source file so all tests pass.\n",
    "\n",
    "You interact via a Python REPL. Available functions:\n",
    "- read_file(path) — Read a file from the repo\n",
    "- list_dir(path='.') — List directory contents\n",
    "- search(pattern, path='.') — Grep for a pattern\n",
    "- write_file(path, content) — Write/create a file\n",
    "- run_tests(test_path=None) — Run pytest on a test file\n",
    "- FINAL() — Signal that your implementation is complete\n",
    "\n",
    "Strategy:\n",
    "1. Read the failing test file to understand expected behavior\n",
    "2. Read other source files for context (imports, dependencies)\n",
    "3. Write the implementation\n",
    "4. Run tests to verify\n",
    "5. Fix any failures\n",
    "6. Call FINAL() when done\n",
    "\n",
    "Output ONLY valid Python code. No markdown, no explanations — just code to execute.\"\"\"\n",
    "\n",
    "\n",
    "def build_prompt(task_description: str, failing_tests: list[str]) -> list[dict]:\n",
    "    \"\"\"Build the chat prompt for the initial observation.\"\"\"\n",
    "    user_msg = f\"{task_description}\\n\\nFailing tests:\\n\" + \"\\n\".join(failing_tests)\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": user_msg},\n",
    "    ]\n",
    "\n",
    "\n",
    "def extract_code_from_response(response: str) -> str:\n",
    "    \"\"\"Extract executable Python code from model response.\"\"\"\n",
    "    # Try to find code blocks first\n",
    "    code_blocks = re.findall(r\"```(?:python)?\\n(.*?)```\", response, re.DOTALL)\n",
    "    if code_blocks:\n",
    "        return \"\\n\".join(code_blocks)\n",
    "    # Otherwise treat the whole response as code\n",
    "    lines = response.strip().split(\"\\n\")\n",
    "    code_lines = []\n",
    "    for line in lines:\n",
    "        stripped = line.strip()\n",
    "        if stripped and not stripped.startswith(\"#\") and any(c in stripped for c in \"=()[]{}:\"):\n",
    "            code_lines.append(line)\n",
    "        elif stripped.startswith(\"#\") or stripped.startswith(\"import\") or stripped.startswith(\"from\"):\n",
    "            code_lines.append(line)\n",
    "        elif not stripped:\n",
    "            code_lines.append(line)\n",
    "        else:\n",
    "            code_lines.append(f\"# {line}\")\n",
    "    return \"\\n\".join(code_lines)\n",
    "\n",
    "\n",
    "print(\"Prompt builder ready.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Trajectory:\n",
    "    \"\"\"A full multi-step episode trajectory for GRPO training.\"\"\"\n",
    "    prompt_text: str        # Tokenized prompt (system + task)\n",
    "    completion_text: str    # All model outputs concatenated\n",
    "    reward: float           # Final episode reward\n",
    "    steps: int              # Number of steps taken\n",
    "    seed: int               # Environment seed (for reproducibility)\n",
    "    tests_passed: int\n",
    "    tests_total: int\n",
    "\n",
    "\n",
    "def run_episode(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    env: RLMForgeEnvironment,\n",
    "    seed: int,\n",
    "    max_steps: int = MAX_STEPS_PER_EPISODE,\n",
    "    temperature: float = 0.7,\n",
    "    max_new_tokens: int = 2048,\n",
    ") -> Trajectory:\n",
    "    \"\"\"Run a single episode: generate code actions, execute them, collect trajectory.\"\"\"\n",
    "    obs = env.reset(seed=seed)\n",
    "\n",
    "    messages = build_prompt(obs.task_description, obs.failing_tests or [])\n",
    "    prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "    all_completions = []  # All model outputs for this episode\n",
    "\n",
    "    for step_i in range(max_steps):\n",
    "        # Build the full conversation so far for the model\n",
    "        if step_i > 0:\n",
    "            # Add the observation as assistant feedback\n",
    "            messages.append({\"role\": \"user\", \"content\": f\"REPL output:\\n{obs.stdout}\\n{obs.stderr}\"})\n",
    "\n",
    "        # Generate next action\n",
    "        full_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "        inputs = tokenizer(full_text, return_tensors=\"pt\", truncation=True, max_length=8192).to(model.device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model.generate(\n",
    "                **inputs,\n",
    "                max_new_tokens=max_new_tokens,\n",
    "                temperature=temperature,\n",
    "                top_p=0.95,\n",
    "                do_sample=True,\n",
    "                pad_token_id=tokenizer.pad_token_id,\n",
    "            )\n",
    "\n",
    "        # Decode only the new tokens\n",
    "        new_tokens = outputs[0][inputs[\"input_ids\"].shape[1]:]\n",
    "        response = tokenizer.decode(new_tokens, skip_special_tokens=True)\n",
    "        all_completions.append(response)\n",
    "\n",
    "        # Add to conversation history\n",
    "        messages.append({\"role\": \"assistant\", \"content\": response})\n",
    "\n",
    "        # Extract and execute code\n",
    "        code = extract_code_from_response(response)\n",
    "\n",
    "        # Check if model wants to finalize\n",
    "        if \"FINAL()\" in code:\n",
    "            obs = env.step(RLMForgeAction(code=code))\n",
    "            break\n",
    "        else:\n",
    "            obs = env.step(RLMForgeAction(code=code))\n",
    "\n",
    "        if obs.done:\n",
    "            break\n",
    "\n",
    "    # If we exhausted steps without FINAL, force finalize\n",
    "    if not obs.done:\n",
    "        obs = env.step(RLMForgeAction(code=\"FINAL()\"))\n",
    "\n",
    "    # Build the full completion text (all model outputs joined)\n",
    "    completion_text = \"\\n<|step|>\\n\".join(all_completions)\n",
    "\n",
    "    reward = obs.reward or 0.0\n",
    "    test_results = obs.test_results or {}\n",
    "\n",
    "    return Trajectory(\n",
    "        prompt_text=prompt_text,\n",
    "        completion_text=completion_text,\n",
    "        reward=reward,\n",
    "        steps=step_i + 1,\n",
    "        seed=seed,\n",
    "        tests_passed=test_results.get(\"tests_passed\", 0),\n",
    "        tests_total=test_results.get(\"tests_total\", 0),\n",
    "    )\n",
    "\n",
    "\n",
    "print(\"Episode runner ready.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Collect Baseline Trajectories\n",
    "\n",
    "Run episodes to collect (prompt, completion, reward) tuples before training. This establishes the pre-training baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_trajectories(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    num_prompts: int = NUM_TRAINING_PROMPTS,\n",
    "    episodes_per_prompt: int = NUM_EPISODES_PER_PROMPT,\n",
    "    temperature: float = 0.7,\n",
    ") -> list[list[Trajectory]]:\n",
    "    \"\"\"Collect GRPO groups: multiple trajectories per unique prompt/seed.\"\"\"\n",
    "    env = RLMForgeEnvironment()\n",
    "    all_groups = []\n",
    "\n",
    "    for prompt_idx in range(num_prompts):\n",
    "        seed = prompt_idx * 100  # Deterministic seeds\n",
    "        group = []\n",
    "\n",
    "        for ep_idx in range(episodes_per_prompt):\n",
    "            print(f\"  Prompt {prompt_idx+1}/{num_prompts}, Episode {ep_idx+1}/{episodes_per_prompt}...\", end=\" \")\n",
    "            traj = run_episode(\n",
    "                model, tokenizer, env,\n",
    "                seed=seed,  # Same seed = same task for GRPO group\n",
    "                temperature=temperature + 0.1 * ep_idx,  # Vary temperature for diversity\n",
    "            )\n",
    "            group.append(traj)\n",
    "            print(f\"reward={traj.reward:.3f}, steps={traj.steps}, \"\n",
    "                  f\"tests={traj.tests_passed}/{traj.tests_total}\")\n",
    "\n",
    "        all_groups.append(group)\n",
    "\n",
    "    env.cleanup()\n",
    "    return all_groups\n",
    "\n",
    "\n",
    "# Collect pre-training baseline\n",
    "print(\"=\" * 60)\n",
    "print(\"COLLECTING BASELINE TRAJECTORIES\")\n",
    "print(\"=\" * 60)\n",
    "baseline_groups = collect_trajectories(model, tokenizer)\n",
    "\n",
    "# Summary stats\n",
    "all_rewards = [t.reward for g in baseline_groups for t in g]\n",
    "print(f\"\\nBaseline: mean_reward={sum(all_rewards)/len(all_rewards):.4f}, \"\n",
    "      f\"min={min(all_rewards):.4f}, max={max(all_rewards):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. GRPO Training\n",
    "\n",
    "Train with Group Relative Policy Optimization. For each group of trajectories (same prompt, different completions), compute advantages relative to the group mean reward, then update the policy to increase probability of higher-reward trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "from trl import GRPOConfig, GRPOTrainer\n",
    "\n",
    "\n",
    "def trajectories_to_dataset(groups: list[list[Trajectory]]) -> Dataset:\n",
    "    \"\"\"Convert trajectory groups into a HuggingFace Dataset for GRPO training.\"\"\"\n",
    "    records = []\n",
    "    for group in groups:\n",
    "        prompt = group[0].prompt_text\n",
    "        for traj in group:\n",
    "            records.append({\n",
    "                \"prompt\": prompt,\n",
    "                \"completion\": traj.completion_text,\n",
    "                \"reward\": traj.reward,\n",
    "            })\n",
    "    return Dataset.from_list(records)\n",
    "\n",
    "\n",
    "def build_reward_fn(groups: list[list[Trajectory]]):\n",
    "    \"\"\"Build a reward function from pre-collected trajectories.\"\"\"\n",
    "    reward_map = {}\n",
    "    for group in groups:\n",
    "        for traj in group:\n",
    "            key = traj.completion_text[:200]\n",
    "            reward_map[key] = traj.reward\n",
    "\n",
    "    def reward_fn(completions: list[str], **kwargs) -> list[float]:\n",
    "        rewards = []\n",
    "        for c in completions:\n",
    "            key = c[:200]\n",
    "            rewards.append(reward_map.get(key, 0.0))\n",
    "        return rewards\n",
    "\n",
    "    return reward_fn\n",
    "\n",
    "\n",
    "# Build dataset from baseline trajectories\n",
    "train_dataset = trajectories_to_dataset(baseline_groups)\n",
    "print(f\"Training dataset: {len(train_dataset)} examples\")\n",
    "print(f\"Sample prompt length: {len(train_dataset[0]['prompt'])} chars\")\n",
    "print(f\"Sample completion length: {len(train_dataset[0]['completion'])} chars\")\n",
    "print(f\"Sample reward: {train_dataset[0]['reward']:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GRPO Training configuration\n",
    "grpo_config = GRPOConfig(\n",
    "    output_dir=\"./rlm_forge_grpo_output\",\n",
    "    num_train_epochs=GRPO_EPOCHS,\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
    "    learning_rate=1e-5,\n",
    "    warmup_ratio=0.1,\n",
    "    max_completion_length=4096,\n",
    "    # max_prompt_length=4096,\n",
    "    num_generations=NUM_EPISODES_PER_PROMPT,  # GRPO group size\n",
    "    logging_steps=1,\n",
    "    save_strategy=\"epoch\",\n",
    "    bf16=True,\n",
    "    gradient_checkpointing=True,\n",
    "    # GRPO-specific\n",
    "    beta=0.1,  # KL penalty coefficient\n",
    "    report_to=\"none\",\n",
    ")\n",
    "\n",
    "# Build reward function from collected trajectories\n",
    "reward_fn = build_reward_fn(baseline_groups)\n",
    "\n",
    "# Prepare prompts dataset (unique prompts only, GRPO generates completions)\n",
    "prompt_dataset = Dataset.from_list([\n",
    "    {\"prompt\": group[0].prompt_text}\n",
    "    for group in baseline_groups\n",
    "])\n",
    "\n",
    "# Initialize GRPO trainer\n",
    "trainer = GRPOTrainer(\n",
    "    model=model,\n",
    "    args=grpo_config,\n",
    "    train_dataset=prompt_dataset,\n",
    "    reward_funcs=reward_fn,\n",
    "    processing_class=tokenizer,\n",
    ")\n",
    "\n",
    "print(\"GRPO Trainer initialized. Starting training...\")\n",
    "trainer.train()\n",
    "print(\"Training complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Post-Training Evaluation\n",
    "\n",
    "Collect new trajectories with the trained model and compare rewards to the baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect post-training trajectories with the same seeds\n",
    "print(\"=\" * 60)\n",
    "print(\"COLLECTING POST-TRAINING TRAJECTORIES\")\n",
    "print(\"=\" * 60)\n",
    "post_groups = collect_trajectories(model, tokenizer, temperature=0.5)\n",
    "\n",
    "post_rewards = [t.reward for g in post_groups for t in g]\n",
    "baseline_rewards = [t.reward for g in baseline_groups for t in g]\n",
    "\n",
    "print(f\"\\n{'='*60}\")\n",
    "print(f\"RESULTS COMPARISON\")\n",
    "print(f\"{'='*60}\")\n",
    "print(f\"Baseline: mean={sum(baseline_rewards)/len(baseline_rewards):.4f}, \"\n",
    "      f\"max={max(baseline_rewards):.4f}\")\n",
    "print(f\"Trained:  mean={sum(post_rewards)/len(post_rewards):.4f}, \"\n",
    "      f\"max={max(post_rewards):.4f}\")\n",
    "print(f\"Improvement: {(sum(post_rewards)/len(post_rewards) - sum(baseline_rewards)/len(baseline_rewards)):.4f}\")\n",
    "\n",
    "# Per-task comparison\n",
    "print(f\"\\nPer-task breakdown:\")\n",
    "for i, (bg, pg) in enumerate(zip(baseline_groups, post_groups)):\n",
    "    b_mean = sum(t.reward for t in bg) / len(bg)\n",
    "    p_mean = sum(t.reward for t in pg) / len(pg)\n",
    "    delta = p_mean - b_mean\n",
    "    arrow = \"\\u2191\" if delta > 0 else \"\\u2193\" if delta < 0 else \"\\u2192\"\n",
    "    print(f\"  Task {i}: baseline={b_mean:.3f} \\u2192 trained={p_mean:.3f} ({arrow} {abs(delta):.3f})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Visualize Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(16, 5))\n",
    "\n",
    "# 1. Reward distribution: baseline vs trained\n",
    "ax1 = axes[0]\n",
    "ax1.hist(baseline_rewards, bins=20, alpha=0.6, label=\"Baseline\", color=\"steelblue\")\n",
    "ax1.hist(post_rewards, bins=20, alpha=0.6, label=\"After GRPO\", color=\"coral\")\n",
    "ax1.set_xlabel(\"Episode Reward\")\n",
    "ax1.set_ylabel(\"Count\")\n",
    "ax1.set_title(\"Reward Distribution\")\n",
    "ax1.legend()\n",
    "ax1.axvline(np.mean(baseline_rewards), color=\"steelblue\", linestyle=\"--\", alpha=0.8)\n",
    "ax1.axvline(np.mean(post_rewards), color=\"coral\", linestyle=\"--\", alpha=0.8)\n",
    "\n",
    "# 2. Per-task mean reward comparison\n",
    "ax2 = axes[1]\n",
    "task_ids = list(range(len(baseline_groups)))\n",
    "b_means = [np.mean([t.reward for t in g]) for g in baseline_groups]\n",
    "p_means = [np.mean([t.reward for t in g]) for g in post_groups]\n",
    "x = np.arange(len(task_ids))\n",
    "width = 0.35\n",
    "ax2.bar(x - width/2, b_means, width, label=\"Baseline\", color=\"steelblue\", alpha=0.8)\n",
    "ax2.bar(x + width/2, p_means, width, label=\"After GRPO\", color=\"coral\", alpha=0.8)\n",
    "ax2.set_xlabel(\"Task ID\")\n",
    "ax2.set_ylabel(\"Mean Reward\")\n",
    "ax2.set_title(\"Per-Task Reward Improvement\")\n",
    "ax2.legend()\n",
    "ax2.set_xticks(x)\n",
    "\n",
    "# 3. Test pass rate improvement\n",
    "ax3 = axes[2]\n",
    "b_pass_rates = [np.mean([t.tests_passed / max(t.tests_total, 1) for t in g]) for g in baseline_groups]\n",
    "p_pass_rates = [np.mean([t.tests_passed / max(t.tests_total, 1) for t in g]) for g in post_groups]\n",
    "ax3.bar(x - width/2, b_pass_rates, width, label=\"Baseline\", color=\"steelblue\", alpha=0.8)\n",
    "ax3.bar(x + width/2, p_pass_rates, width, label=\"After GRPO\", color=\"coral\", alpha=0.8)\n",
    "ax3.set_xlabel(\"Task ID\")\n",
    "ax3.set_ylabel(\"Test Pass Rate\")\n",
    "ax3.set_title(\"Test Pass Rate Improvement\")\n",
    "ax3.legend()\n",
    "ax3.set_xticks(x)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"rlm_forge_results.png\", dpi=150, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "\n",
    "print(f\"\\nOverall test pass rate:\")\n",
    "print(f\"  Baseline: {np.mean(b_pass_rates):.1%}\")\n",
    "print(f\"  Trained:  {np.mean(p_pass_rates):.1%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Save Model & Training Log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the trained LoRA adapter\n",
    "model.save_pretrained(\"./rlm_forge_lora_adapter\")\n",
    "tokenizer.save_pretrained(\"./rlm_forge_lora_adapter\")\n",
    "\n",
    "# Save training log\n",
    "training_log = {\n",
    "    \"model_id\": MODEL_ID,\n",
    "    \"num_prompts\": NUM_TRAINING_PROMPTS,\n",
    "    \"episodes_per_prompt\": NUM_EPISODES_PER_PROMPT,\n",
    "    \"max_steps_per_episode\": MAX_STEPS_PER_EPISODE,\n",
    "    \"grpo_epochs\": GRPO_EPOCHS,\n",
    "    \"baseline_mean_reward\": float(np.mean(baseline_rewards)),\n",
    "    \"baseline_max_reward\": float(max(baseline_rewards)),\n",
    "    \"trained_mean_reward\": float(np.mean(post_rewards)),\n",
    "    \"trained_max_reward\": float(max(post_rewards)),\n",
    "    \"improvement\": float(np.mean(post_rewards) - np.mean(baseline_rewards)),\n",
    "    \"baseline_test_pass_rate\": float(np.mean(b_pass_rates)),\n",
    "    \"trained_test_pass_rate\": float(np.mean(p_pass_rates)),\n",
    "}\n",
    "\n",
    "with open(\"training_log.json\", \"w\") as f:\n",
    "    json.dump(training_log, f, indent=2)\n",
    "\n",
    "print(\"Saved LoRA adapter to ./rlm_forge_lora_adapter\")\n",
    "print(\"Saved training log to training_log.json\")\n",
    "print(f\"\\nFinal summary:\")\n",
    "print(json.dumps(training_log, indent=2))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}