Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* 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
+assets/Logs.png filter=lfs diff=lfs merge=lfs -text
+assets/molforge_architecture.png filter=lfs diff=lfs merge=lfs -text
+assets/reward_curve.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,191 @@
+---
+title: "MolForge: Verifier-Driven RL for Drug Discovery"
+emoji: 🧬
+colorFrom: indigo
+colorTo: blue
+sdk: static
+pinned: false
+license: mit
+tags:
+- reinforcement-learning
+- drug-discovery
+- chemistry
+- multi-agent
+- oncology
+- molecular-simulation
+- openenv
+---
+
+# MolForge: Verifier-Driven RL for Drug Discovery
+
+MolForge is a reinforcement learning environment that simulates a **medical oncology discovery lab**. Unlike traditional LLM tasks where the model generates a final answer in one shot, MolForge forces the model to execute the **scientific method** under real-world constraints: budget, toxicity, and synthesis complexity.
+
+### The Scientific Method as a Workflow
+
+Imagine a biotech team tasked with optimizing a lead candidate for **KRAS G12C** (including a high-difficulty resistance panel). The model doesn't just "write" a molecule; it controls a specialist team that must navigate a resource-constrained laboratory:
+
+- **Lead Chemist**: Proposes molecular edits and decides when to submit.
+- **Assay Planner**: Allocates limited budget to run empirical tests.
+- **Toxicologist**: Reviews safety risks and can object to unsafe designs.
+- **Process Chemist**: Evaluates whether the molecule is practical to synthesize.
+
+Every action—editing a fragment, running a docking simulation, or ordering a toxicity assay—is a decision that impacts the final outcome. The model must learn to gather enough evidence to justify a submission while keeping the project within budget.
+
+> **Core Philosophy:** The LLM is not the judge. The LLM is the scientist being judged by external, verifiable reality.
+
+
+## What Makes MolForge Special?
+
+MolForge is built to move beyond simple "molecule generation" into "scientific workflow optimization." Here are the seven core pillars that make it unique:
+
+1. [**Verifier-Based Evaluation**](#1-verifier-based-evaluation): The LLM is the scientist, not the judge. It is held accountable by real-world verifiers like **RDKit** and **TDC**.
+2. [**Chemical & Molecular Simulations**](#2-chemical--molecular-simulations): Realistic simulation of potency and existence using heuristic docking, **RDKit**, and **TDC**.
+3. [**Self-Correction & Improvement Loop**](#3-self-correction--improvement-loop): After each edit, agents receive structured feedback from verifiers, allowing them to self-correct.
+4. [**Decomposed Reward Architecture**](#4-decomposed-reward-architecture): Multi-step rewards for every action (research, edits, coordination) provide high observability.
+5. [**Scientific Model Improvement**](#5-scientific-model-improvement): Real verifier feedback (Reviews) guides the model toward scientifically sound designs.
+6. [**Strategic Training Modes**](#6-strategic-training-modes): A dual-mode system using **Curriculum mode** (partial credit) and **Assay-Gated mode** (strict).
+7. [**Multi-Agent Governance**](#7-multi-agent-governance): A specialized team that plans, executes, and shares information to coordinate the next plan of action.
+
+---
+
+
+## Architecture
+
+The architecture is a closed scientific feedback loop:
+
+![MolForge Architecture](assets/molforge_architecture.png)
+
+## Scientific Verifier Layers
+
+### RDKit: chemical plausibility
+
+RDKit checks molecule-like behavior and chemistry descriptors. In MolForge, this layer is used to keep the molecule edits grounded in chemical reality. It supports descriptor-style reasoning such as lipophilicity, polarity, tractability, and drug-likeness.
+
+### TDC: biomedical outcome signals
+
+Therapeutics Data Commons represents the medical outcome side of the environment. It provides the project with a path toward realistic prediction tasks such as toxicity, synthesis difficulty, and drug-likeness. In the default Docker deployment, RDKit remains active and TDC is optional because it can pull a heavier platform-sensitive ML stack.
+
+### Heuristic docking: receptor fit
+
+MolForge includes a docking-style surrogate that answers three fast questions:
+
+| Check | Question | Why it matters |
+| --- | --- | --- |
+| Pocket matching | Does the hinge fragment fit the receptor pocket? | Better pocket complementarity improves potency. |
+| Lipophilic match | Is LogP near the target pocket's hydrophobic comfort zone around `3.0`? | Too much lipophilicity can increase toxicity; too little can weaken binding. |
+| Polarity match | Is TPSA near a useful range around `85.0`? | Polarity affects binding, permeability, and clash risk. |
+
+This gives the environment fast receptor-aware feedback in milliseconds, which is important for RL.
+
+## Training Story
+
+The training pipeline has two stages:
+
+1. **SFT warm start**
+2. **RL with verifier rewards**
+
+### Base model
+
+The model used for the main run is:
+
+```text
+unsloth/Qwen3.5-2B
+```
+
+The raw base model was not reliable enough for the environment at first. It often failed to produce the exact structured JSON actions that MolForge expects, and it did not consistently respect the specialist-agent interaction format.
+
+So the first step was a small SFT warm start. This stage is not meant to teach the model the optimal chemistry. It teaches the model how to speak the environment's action language:
+
+- valid JSON actions
+- correct role/action pairing
+- correct molecule slots and fragments
+- concise rationales
+- evidence fields based only on visible observations
+- expected-effect fields such as potency up/down or toxicity up/down
+- valid specialist messages where needed
+
+### Training Results
+After SFT, the policy is trained with GRPO-style RL against MolForge itself. During training, the model explores the 256-combination molecular edit space, receiving rewards for molecule quality, evidence coverage, and budget discipline.
+
+![Reward Curve](assets/reward_curve.png)
+![Training Logs](assets/Logs.png)
+
+As shown in the reward curve and logs, the model successfully learns to navigate the scientific constraints, moving from early exploration to consistent, verifier-backed molecule submissions. For strict evaluation, the environment switches back to `assay_gated` mode.
+
+
+
+
+
+
+
+
+
+
+## Example Successful Strategy
+
+A strong MolForge policy usually learns to:
+
+1. Start with a cheap property assay.
+2. Edit one high-impact molecular slot.
+3. Re-assay after the edit.
+4. Repair toxicity or synthesis issues before chasing more potency.
+5. Save expensive toxicity or MD-style assays for candidates that are worth confirming.
+6. Use specialist review messages before risky submission.
+7. Restart in the hard scenario if the scaffold appears trapped.
+8. Submit only when potency, toxicity, and synthesis evidence are present for the current molecule.
+
+
+
+## Why This Project Matters
+
+MolForge is designed to test a deeper kind of AI capability than simple answer generation. The model must work inside a scientific feedback loop where actions are checked, evidence costs money, unsafe decisions can be blocked, and the final answer only matters if the path to that answer is experimentally justified.
+
+The strongest part of the project is that the LLM is not trusted by default. It has to earn trust through verifier-backed decisions.
+
+## Deep Dive: What Makes MolForge Special?
+
+### 1. Verifier-Based Evaluation
+In many LLM systems, the model itself is used as a judge, reviewer, or evaluator. MolForge flips that pattern. The LLM is the scientist being judged, not the judge. It is held accountable by real-world verifiers like **RDKit**, **TDC**, and molecular simulation engines. This ensures that the model's progress is grounded in chemical and biological reality, not just persuasive language.
+
+### 2. Chemical & Molecular Simulations
+MolForge doesn't just predict outcomes; it utilizes multiple simulation layers to ground the model's decisions:
+*   **Chemical Plausibility (RDKit):** Decides if the molecule generated by the LLM (via edits) can actually exist in chemical reality. [Visit RDKit](https://www.rdkit.org)
+*   **Medical Outcomes (TDC):** Predicts the most probable medical outcomes and properties using the [Therapeutics Data Commons](https://tdcommons.ai).
+*   **Heuristic Docking Score:** A fast, physics-inspired simulation that updates **potency** in milliseconds based on three rules:
+    1.  **Pocket Matching:** Does the fragment structurally fit the target receptor pocket (e.g., KRAS G12C)?
+    2.  **Lipophilic Match:** Is the LogP near the ideal **3.0** to sit comfortably in the hydrophobic pocket?
+    3.  **Polarity Match:** Is the TPSA near the ideal **85.0** to avoid repulsive polar clashes?
+
+### 3. Self-Correction & Improvement Loop
+MolForge is an iterative environment. After each proposed molecular modification, the model receives a structured review from the verifiers. This feedback allows the agent to recognize liabilities (like toxicity or low potency) and correct them in the next step. This creates a genuine **self-improvement loop** within each episode.
+
+### 4. Decomposed Reward Architecture
+The reward function is not a single "black box" scalar. We use a multi-step reward system where small-scale rewards are designed for every individual action—research, molecular edits, and inter-agent coordination. While we may output a single total reward for training simplicity (especially for the hackathon), the decomposed components allow for massive observability into which sections of the workflow are lacking.
+
+### 5. Scientific Model Improvement
+We use real verifier feedback to drive model improvement. By providing constant, verifiable reviews, we train the model to improve its designs based on evidence. This moves the model away from simple pattern matching and toward a more rigorous, evidence-based design process.
+
+### 6. Strategic Training Modes: Curriculum vs. Assay-Gated
+To solve the "sparse reward" problem common in RL, MolForge uses two distinct modes:
+*   **Curriculum Mode (Training):** If a model fails to submit but showed good scientific behavior, it receives "Partial Credit" (up to +0.75). It gets points for gathering evidence and designing promising molecules. These "breadcrumbs" teach the model how to explore before it discovers the terminal submission bonus.
+*   **Assay-Gated Mode (Evaluation):** This is the strict, official mode used for hackathon grading. There is **zero partial credit**. If the model fails to explicitly `submit` a high-potency, safe molecule before the budget runs out, its score is exactly `0.0`.
+
+### 7. Multi-Agent Governance
+Drug discovery is a team effort. MolForge implements a multi-agent system where specialized roles (Lead Chemist, Toxicologist, Assay Planner) review each other's moves, plans, and executions. Crucially, these agents **share information and coordinate** between themselves to decide the next plan of action, ensuring that every decision undergoes a rigorous "peer review" process before execution.
+
+---
+
+## Final Takeaway
+
+MolForge is special because it treats the LLM as a trainable research agent inside a controlled scientific environment, not as an oracle. The model is judged by chemistry and biomedical verifiers, corrected by specialist feedback, constrained by assay budget, and scored by a reward system that can explain where the policy succeeded or failed.
+
+The important pieces work together:
+
+- **Verifier-first evaluation:** RDKit, TDC-style signals, and docking-style simulation judge the model's actions.
+- **Multi-agent review:** specialist roles create checks and balances around each decision.
+- **Self-improvement loop:** every action produces feedback that the next action can respond to.
+- **Decomposed rewards:** the environment tracks molecule quality, evidence, budget, coordination, and safety separately.
+- **Curriculum to strict evaluation:** training can use partial-credit breadcrumbs, while final evaluation remains unforgiving.
+- **Dynamic molecular search:** the model explores 256 fragment combinations across three starting scientific scenarios instead of memorizing one answer.
+
+That is the project thesis: useful scientific agents should not merely generate plausible ideas. They should operate in a loop where the world pushes back.

molforge_grpo_official_submission.ipynb

@@ -0,0 +1,368 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# MolForge GRPO Training Pipeline\n",
+        "This notebook implements the Reinforcement Learning (GRPO) training pipeline for the MolForge environment.\n",
+        "We train the model using a **Proposer-Critic-Selector** architecture and targeted **reward shaping** to overcome local minima."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install -U \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
+        "!pip install -U \"trl>=0.21.0\" peft accelerate bitsandbytes datasets matplotlib pandas huggingface_hub \"openenv-core[core]>=0.2.3\" rdkit jmespath xformers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "import sys\n",
+        "from pathlib import Path\n",
+        "\n",
+        "# Clone the repository\n",
+        "if not Path(\"/content/molt_lab\").exists():\n",
+        "    !git clone https://github.com/Adhitya-Vardhan/molt_lab.git /content/molt_lab\n",
+        "\n",
+        "# Add project root to path\n",
+        "if \"/content/molt_lab\" not in sys.path:\n",
+        "    sys.path.insert(0, \"/content/molt_lab\")\n",
+        "    \n",
+        "# Change working directory\n",
+        "os.chdir(\"/content/molt_lab\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import time\n",
+        "import os\n",
+        "\n",
+        "# Training Configuration\n",
+        "os.environ[\"MOLFORGE_REWARD_MODE\"] = \"curriculum\"\n",
+        "os.environ[\"MOLFORGE_TRAINING_RANDOMIZATION\"] = \"1\"\n",
+        "\n",
+        "RL_MAX_STEPS = 300\n",
+        "NUM_GENERATIONS = 2\n",
+        "PER_DEVICE_BATCH = 2\n",
+        "GRAD_ACCUM = 4\n",
+        "LEARNING_RATE = 2e-6\n",
+        "MAX_SEQ_LENGTH = 2048\n",
+        "MAX_PROMPT_LENGTH = 1536\n",
+        "MAX_COMPLETION_LENGTH = 384\n",
+        "\n",
+        "RUN_NAME = time.strftime(\"molforge_grpo_%Y%m%d_%H%M%S\")\n",
+        "OUTPUT_DIR = Path(f\"/content/molforge_rl_runs/{RUN_NAME}\")\n",
+        "ADAPTER_SAVE_DIR = OUTPUT_DIR / \"adapters\"\n",
+        "PLOT_DIR = OUTPUT_DIR / \"plots\"\n",
+        "\n",
+        "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
+        "PLOT_DIR.mkdir(parents=True, exist_ok=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Reward Function & OpenEnv Integration\n",
+        "We implement a custom reward function that wraps the native `MolForgeEnvironment`. \n",
+        "To prevent \"reward hacking\" (where the model endlessly farms `run_assay` for safe points), we apply targeted reward shaping."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "from typing import Any, Dict, Tuple\n",
+        "from inference_common import (\n",
+        "    MolForgeAction,\n",
+        "    attach_reasoning_fields,\n",
+        "    attach_team_messages,\n",
+        "    extract_json,\n",
+        ")\n",
+        "from server.molforge_environment import MolForgeEnvironment\n",
+        "from models import MolForgeState\n",
+        "\n",
+        "def replay_to_state(record: dict[str, Any]) -> MolForgeEnvironment:\n",
+        "    \"\"\"Replays actions to reach a specific state.\"\"\"\n",
+        "    env = MolForgeEnvironment()\n",
+        "    # Set randomization and seed if provided\n",
+        "    if record.get(\"randomized\"):\n",
+        "        os.environ[\"MOLFORGE_TRAINING_RANDOMIZATION\"] = \"1\"\n",
+        "    os.environ[\"MOLFORGE_RAND_SEED\"] = str(record.get(\"random_seed\", \"rl\"))\n",
+        "    \n",
+        "    observation = env.reset()\n",
+        "    for action_payload in record.get(\"pre_actions\", []):\n",
+        "        action = MolForgeAction(**action_payload)\n",
+        "        observation = env.step(attach_team_messages(observation, attach_reasoning_fields(observation, action)))\n",
+        "        if observation.done:\n",
+        "            break\n",
+        "    return env\n",
+        "\n",
+        "def evaluate_completion(prompt_str: str, completion_str: str, record: dict[str, Any]) -> Tuple[float, dict]:\n",
+        "    diagnostics = {\"valid_json\": False}\n",
+        "    try:\n",
+        "        action_dict = extract_json(completion_str)\n",
+        "        action = MolForgeAction(**action_dict)\n",
+        "    except Exception:\n",
+        "        return -1.2, diagnostics\n",
+        "\n",
+        "    diagnostics[\"valid_json\"] = True\n",
+        "    env = replay_to_state(record)\n",
+        "    \n",
+        "    # Step the OpenEnv environment\n",
+        "    observation = env._build_observation(reward=0.0, done=False, reward_components=[])\n",
+        "    action = attach_team_messages(observation, attach_reasoning_fields(observation, action))\n",
+        "    next_observation = env.step(action)\n",
+        "    \n",
+        "    reward = float(next_observation.reward)\n",
+        "    grader_scores = next_observation.metadata.get(\"terminal_grader_scores\", {})\n",
+        "    \n",
+        "    # Reward Shaping\n",
+        "    if action.action_type == \"run_assay\" and reward > 0:\n",
+        "        reward *= 0.25\n",
+        "    elif action.action_type == \"submit\":\n",
+        "        sub_score = float(grader_scores.get(\"submission_score\", 0.0))\n",
+        "        if sub_score > 0.0: reward += sub_score * 3.0\n",
+        "    elif action.action_type == \"edit\" and reward > 0:\n",
+        "        reward *= 1.5\n",
+        "\n",
+        "    diagnostics.update({\"action_type\": action.action_type, \"reward\": reward, \"done\": next_observation.done})\n",
+        "    return reward, diagnostics\n",
+        "\n",
+        "def molforge_reward_func(prompts, completions, **kwargs) -> list[float]:\n",
+        "    rewards = []\n",
+        "    # When using dynamic dataset, columns like scenario_id, pre_actions etc are in kwargs\n",
+        "    for i in range(len(completions)):\n",
+        "        # Reconstruct the record from kwargs\n",
+        "        record = {\n",
+        "            \"pre_actions\": kwargs[\"record\"][i][\"pre_actions\"] if \"record\" in kwargs else [],\n",
+        "            \"randomized\": True,\n",
+        "            \"random_seed\": \"dynamic-rl\"\n",
+        "        }\n",
+        "        prompt_str = prompts[i][-1][\"content\"] if isinstance(prompts[i], list) else str(prompts[i])\n",
+        "        completion_str = completions[i][0][\"content\"] if isinstance(completions[i], list) else str(completions[i])\n",
+        "        reward, _ = evaluate_completion(prompt_str, completion_str, record)\n",
+        "        rewards.append(reward)\n",
+        "    return rewards\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Model & Tokenizer Loading"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from unsloth import FastLanguageModel\n",
+        "\n",
+        "# Set this to your SFT checkpoint Deployed to hugging face space \n",
+        "# SFT trained on only to mimic the response behavioiur of the model (structured responses visit the hf blog for more detailed explanation )\n",
+        "SFT_ADAPTER_PATH = \"Adhitya122/qwen3_5_2b_molforge_sft_v4\"\n",
+        "\n",
+        "print(\"Loading model and applying Unsloth optimizations...\")\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name=SFT_ADAPTER_PATH,\n",
+        "    max_seq_length=MAX_SEQ_LENGTH,\n",
+        "    dtype=None,\n",
+        "    load_in_4bit=True,\n",
+        ")\n",
+        "\n",
+        "# Enable fast training paths\n",
+        "FastLanguageModel.for_training(model)\n",
+        "\n",
+        "# Extract underlying tokenizer if it is wrapped in a vision processor\n",
+        "if hasattr(tokenizer, \"tokenizer\"):\n",
+        "    tokenizer = tokenizer.tokenizer"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### GRPO Training Loop"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from datasets import Dataset\n",
+        "from scripts.generate_sft_compact_policy_v4_dataset import compact_action_payload, COMPACT_ACTION_SYSTEM_PROMPT\n",
+        "from inference_common import heuristic_team_action\n",
+        "import random\n",
+        "\n",
+        "def build_dynamic_prompts(episodes=50, max_turns=5) -> Dataset:\n",
+        "    \"\"\"Generates training prompts by playing the environment with a heuristic expert.\"\"\"\n",
+        "    print(f\"Generating {episodes} episodes of dynamic prompts...\")\n",
+        "    records = []\n",
+        "    env = MolForgeEnvironment()\n",
+        "    \n",
+        "    for _ in range(episodes):\n",
+        "        observation = env.reset()\n",
+        "        pre_actions = []\n",
+        "        \n",
+        "        for _ in range(max_turns):\n",
+        "            if observation.done:\n",
+        "                break\n",
+        "            \n",
+        "            # Capture the current state as a prompt\n",
+        "            prompt_payload = compact_action_payload(observation)\n",
+        "            records.append({\n",
+        "                \"prompt\": [\n",
+        "                    {\"role\": \"system\", \"content\": COMPACT_ACTION_SYSTEM_PROMPT},\n",
+        "                    {\"role\": \"user\", \"content\": json.dumps(prompt_payload)}\n",
+        "                ],\n",
+        "                \"record\": {\n",
+        "                    \"scenario_id\": observation.scenario_id,\n",
+        "                    \"difficulty\": observation.difficulty,\n",
+        "                    \"step_index\": observation.step_index,\n",
+        "                    \"pre_actions\": list(pre_actions),\n",
+        "                    \"randomized\": True,\n",
+        "                    \"random_seed\": \"dynamic-rl\"\n",
+        "                }\n",
+        "            })\n",
+        "            \n",
+        "            # Use expert to move to the next state\n",
+        "            action = heuristic_team_action(observation)\n",
+        "            observation = env.step(action)\n",
+        "            pre_actions.append({\"action_type\": action.action_type, \"acting_role\": action.acting_role})\n",
+        "            \n",
+        "    random.shuffle(records)\n",
+        "    return Dataset.from_list(records)\n",
+        "\n",
+        "# Generate the dataset dynamically (no .jsonl needed!)\n",
+        "dataset = build_dynamic_prompts(episodes=20, max_turns=6)\n",
+        "print(f\"Dynamic dataset created with {len(dataset)} prompt states.\")\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from trl import GRPOConfig, GRPOTrainer\n",
+        "import inspect\n",
+        "\n",
+        "# Safe GRPO Configuration (detects supported arguments automatically)\n",
+        "config_kwargs = {\n",
+        "    \"output_dir\": str(OUTPUT_DIR),\n",
+        "    \"learning_rate\": LEARNING_RATE,\n",
+        "    \"per_device_train_batch_size\": PER_DEVICE_BATCH,\n",
+        "    \"gradient_accumulation_steps\": GRAD_ACCUM,\n",
+        "    \"max_prompt_length\": MAX_PROMPT_LENGTH,\n",
+        "    \"max_completion_length\": MAX_COMPLETION_LENGTH,\n",
+        "    \"num_generations\": NUM_GENERATIONS,\n",
+        "    \"max_steps\": RL_MAX_STEPS,\n",
+        "    \"logging_steps\": 1,\n",
+        "    \"save_steps\": 25,\n",
+        "    \"bf16\": True,\n",
+        "    \"report_to\": \"none\",\n",
+        "    \"log_completions\": True,\n",
+        "}\n",
+        "\n",
+        "# Filter arguments to only pass what the current library version supports\n",
+        "supported_params = inspect.signature(GRPOConfig.__init__).parameters\n",
+        "filtered_kwargs = {k: v for k, v in config_kwargs.items() if k in supported_params}\n",
+        "\n",
+        "training_args = GRPOConfig(**filtered_kwargs)\n",
+        "\n",
+        "# Initialize Trainer\n",
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    reward_funcs=molforge_reward_func,\n",
+        "    args=training_args,\n",
+        "    train_dataset=dataset,\n",
+        "    processing_class=tokenizer,\n",
+        ")\n",
+        "\n",
+        "print(\"Starting GRPO Training...\")\n",
+        "trainer.train()\n",
+        "\n",
+        "print(f\"Training complete. Saving adapters to {ADAPTER_SAVE_DIR}\")\n",
+        "trainer.save_model(str(ADAPTER_SAVE_DIR))\n",
+        "tokenizer.save_pretrained(str(ADAPTER_SAVE_DIR))\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "import pandas as pd\n",
+        "\n",
+        "# Extract metrics from trainer log history\n",
+        "log_history = trainer.state.log_history\n",
+        "df = pd.DataFrame(log_history)\n",
+        "\n",
+        "print(\"Generating plots...\")\n",
+        "\n",
+        "if \"loss\" in df.columns:\n",
+        "    plt.figure(figsize=(10, 5))\n",
+        "    df_loss = df.dropna(subset=[\"loss\"])\n",
+        "    plt.plot(df_loss[\"step\"], df_loss[\"loss\"], label=\"Loss\", color=\"blue\")\n",
+        "    plt.title(\"Training Loss\")\n",
+        "    plt.grid(True)\n",
+        "    plt.show()\n",
+        "\n",
+        "if \"reward\" in df.columns:\n",
+        "    plt.figure(figsize=(10, 5))\n",
+        "    df_reward = df.dropna(subset=[\"reward\"])\n",
+        "    plt.plot(df_reward[\"step\"], df_reward[\"reward\"], label=\"Reward\", color=\"green\")\n",
+        "    plt.title(\"Reward Curve\")\n",
+        "    plt.grid(True)\n",
+        "    plt.show()\n",
+        "\n",
+        "import shutil\n",
+        "from google.colab import files\n",
+        "\n",
+        "print(f\"Zipping results from {OUTPUT_DIR}...\")\n",
+        "zip_filename = f\"{RUN_NAME}_results\"\n",
+        "shutil.make_archive(zip_filename, \"zip\", OUTPUT_DIR)\n",
+        "\n",
+        "print(f\"Downloading {zip_filename}.zip...\")\n",
+        "files.download(f\"{zip_filename}.zip\")\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}