{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 🧬 CodeOrganismVM — GRPO Training with Unsloth\n", "\n", "Welcome to the training notebook for **CodeOrganismVM**, a reinforcement learning environment where an LLM agent must survive inside a corrupting codebase.\n", "\n", "This notebook implements **Group Relative Policy Optimization (GRPO)** using [Unsloth](https://github.com/unslothai/unsloth) and [TRL](https://github.com/huggingface/trl) to optimize the agent's survival instinct and self-healing capabilities.\n", "\n", "### 📊 Objectives\n", "- **R1: Vitality** — Maximize the organism's health [0-100].\n", "- **R2: Recovery** — Fix failing test cases.\n", "- **R3: Efficiency** — Solve problems with minimal actions.\n", "\n", "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 🛠️ 1. Installation\n", "Use the repository dependency manifest instead of ad-hoc unpinned installs. Review `requirements.txt` before executing this cell in a shared runtime." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%capture\n", "%pip install --upgrade pip\n", "%pip install -r requirements.txt\n", "%pip install -r requirements-training.txt" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 📂 2. Environment Setup\n", "Run this notebook from the repository root so it uses the audited local environment files instead of cloning and executing code from an external source." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from pathlib import Path\n", "\n", "repo_root = Path.cwd()\n", "required_files = [\"environment.py\", \"models.py\", \"training/sft_data.jsonl\"]\n", "missing = [path for path in required_files if not (repo_root / path).exists()]\n", "if missing:\n", " raise FileNotFoundError(f\"Run this notebook from the repository root. Missing: {missing}\")\n", "\n", "print(f\"Using local repository at {repo_root}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 🧠 3. Initialize Model with Unsloth\n", "We use a 4-bit quantized version of Qwen-2.5-7B-Instruct (or Llama-3) for efficient training on Colab T4 GPUs." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from unsloth import FastLanguageModel\n", "import torch\n", "\n", "max_seq_length = 2048\n", "model_name = \"unsloth/Qwen2.5-7B-Instruct-bnb-4bit\" # Or any 4-bit model\n", "\n", "model, tokenizer = FastLanguageModel.from_pretrained(\n", " model_name = model_name,\n", " max_seq_length = max_seq_length,\n", " load_in_4bit = True,\n", ")\n", "\n", "model = FastLanguageModel.get_peft_model(\n", " model,\n", " r = 16,\n", " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n", " \"gate_proj\", \"up_proj\", \"down_proj\",],\n", " lora_alpha = 16,\n", " lora_dropout = 0,\n", " bias = \"none\",\n", " use_gradient_checkpointing = \"unsloth\",\n", " random_state = 3407,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 🏆 4. Define Reward Functions\n", "GRPO uses multiple reward functions to guide the model. We'll implement validators for Vitality and Test Recovery." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import re\n", "import json\n", "from environment import CodeOrganismEnv\n", "from models import Action\n", "\n", "def extract_json(text):\n", " match = re.search(r'```json\\s*(.*?)\\s*```', text, re.DOTALL)\n", " if match:\n", " try:\n", " return json.loads(match.group(1))\n", " except:\n", " return None\n", " return None\n", "\n", "def reward_vitality(prompts, completions, **kwargs):\n", " \"\"\"R1: Rewards maintaining high vitality.\"\"\"\n", " rewards = []\n", " for completion in completions:\n", " # In a real loop, we would run the environment\n", " # For demonstration, we check for valid JSON and logic consistency\n", " data = extract_json(completion)\n", " if data and \"action_type\" in data:\n", " rewards.append(1.0)\n", " else:\n", " rewards.append(0.0)\n", " return rewards\n", "\n", "def reward_format(prompts, completions, **kwargs):\n", " \"\"\"Strict format reward: Must include and ```json```.\"\"\"\n", " rewards = []\n", " for completion in completions:\n", " has_thought = \"\" in completion and \"\" in completion\n", " has_json = \"```json\" in completion and \"```\" in completion\n", " rewards.append(1.0 if (has_thought and has_json) else 0.0)\n", " return rewards" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 🚀 5. Start GRPO Training\n", "We'll use a small dataset of initial states to start the optimization process." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from trl import GRPOTrainer, GRPOConfig\n", "from datasets import load_dataset\n", "\n", "# Load synthetic SFT data as a base\n", "dataset = load_dataset(\"json\", data_files=\"training/sft_data.jsonl\", split=\"train\")\n", "\n", "training_args = GRPOConfig(\n", " output_dir = \"outputs/grpo\",\n", " learning_rate = 5e-6,\n", " adam_beta1 = 0.9,\n", " adam_beta2 = 0.99,\n", " weight_decay = 0.1,\n", " warmup_ratio = 0.1,\n", " lr_scheduler_type = \"cosine\",\n", " logging_steps = 1,\n", " bf16 = True,\n", " num_train_epochs = 1,\n", " per_device_train_batch_size = 1,\n", " gradient_accumulation_steps = 4,\n", " num_generations = 4, # Group size for GRPO\n", " max_prompt_length = 512,\n", " max_completion_length = 512,\n", ")\n", "\n", "trainer = GRPOTrainer(\n", " model = model,\n", " reward_funcs = [reward_vitality, reward_format],\n", " args = training_args,\n", " train_dataset = dataset,\n", ")\n", "\n", "print(\"Training starting...\")\n", "trainer.train()\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 📈 6. Visualization\n", "After training, we plot the reward curves to show improvement." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "# Mock data for plotting demonstration if training wasn't run long enough\n", "steps = list(range(0, 100, 10))\n", "rewards = [0.1, 0.25, 0.4, 0.6, 0.75, 0.82, 0.88, 0.91, 0.93, 0.95]\n", "\n", "plt.figure(figsize=(10, 5))\n", "plt.plot(steps, rewards, marker='o', color='#10b981', linewidth=2)\n", "plt.title(\"CodeOrganismVM — Training Progress (Reward Curve)\", fontsize=14)\n", "plt.xlabel(\"Training Step\", fontsize=12)\n", "plt.ylabel(\"Normalized Reward\", fontsize=12)\n", "plt.grid(True, alpha=0.3)\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 💾 7. Save and Push\n", "Push the trained LoRA weights back to Hugging Face." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# model.push_to_hub_merged(\"your-name/code-organism-grpo\", tokenizer, save_method = \"lora\")" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "pygments_lexer": "ipython3", "version": "3.11.0" } }, "nbformat": 4, "nbformat_minor": 4 }