diff --git a/README.md b/README.md
index 17995d1f..be2fdbc5 100644
--- a/README.md
+++ b/README.md
@@ -38,26 +38,42 @@ We provide all runs and data from our paper [here](https://drive.google.com/driv
 ## Table of Contents
 
 1. [Introduction](#introduction)
-2. [Requirements](#requirements)
+2. [🚀 Quick Start: Google Colab](#-quick-start-google-colab)
+3. [Requirements](#requirements)
    - [Installation](#installation)
    - [Supported Models and API Keys](#supported-models-and-api-keys)
-3. [Setting Up the Templates](#setting-up-the-templates)
+4. [Setting Up the Templates](#setting-up-the-templates)
    - [NanoGPT Template](#nanogpt-template)
    - [2D Diffusion Template](#2d-diffusion-template)
    - [Grokking Template](#grokking-template)
-4. [Run AI Scientist Paper Generation Experiments](#run-ai-scientist-paper-generation-experiments)
-5. [Getting an LLM-Generated Paper Review](#getting-an-llm-generated-paper-review)
-6. [Making Your Own Template](#making-your-own-template)
+5. [Run AI Scientist Paper Generation Experiments](#run-ai-scientist-paper-generation-experiments)
+6. [Getting an LLM-Generated Paper Review](#getting-an-llm-generated-paper-review)
+7. [Making Your Own Template](#making-your-own-template)
    - [Community-Contributed Templates](#community-contributed-templates)
-7. [Template Resources](#template-resources)
-8. [Citing The AI Scientist](#citing-the-ai-scientist)
-9. [Frequently Asked Questions](#frequently-asked-questions)
-10. [Containerization](#containerization)
+8. [Template Resources](#template-resources)
+9. [Citing The AI Scientist](#citing-the-ai-scientist)
+10. [Frequently Asked Questions](#frequently-asked-questions)
+11. [Containerization](#containerization)
 
 ## Introduction
 
 We provide three templates, which were used in our paper, covering the following domains: **NanoGPT**, **2D Diffusion**, and **Grokking**. These templates enable The AI Scientist to generate ideas and conduct experiments in these areas. We accept contributions of new templates from the community, but please note that they are not maintained by us. All other templates beyond the three provided are community contributions.
 
+## Quick Start: Google Colab
+
+**New!** Try AI Scientist in your browser with zero local setup:
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SakanaAI/AI-Scientist/blob/main/notebooks/ai_scientist_2d_diffusion.ipynb)
+
+- **30-45 minutes** to your first AI-generated paper
+- **$3-10** total cost (with GPT-4o-mini or Claude)
+- **Free GPU** included (Colab T4)
+- **2D Diffusion template** - beautiful visualizations
+
+Perfect for first-time users! See [notebooks/README.md](notebooks/README.md) for detailed instructions, cost breakdown, and troubleshooting.
+
+For full functionality (parallel experiments, all templates, custom datasets), follow the local installation guide below.
+
 ## Requirements
 
 This code is designed to run on Linux with NVIDIA GPUs using CUDA and PyTorch. Support for other GPU architectures may be possible by following the [PyTorch guidelines](https://pytorch.org/get-started/locally/). The current templates would likely take an infeasible amount of time on CPU-only machines. Running on other operating systems may require significant adjustments.
diff --git a/notebooks/ai_scientist_2d_diffusion.ipynb b/notebooks/ai_scientist_2d_diffusion.ipynb
new file mode 100644
index 00000000..b6beb04d
--- /dev/null
+++ b/notebooks/ai_scientist_2d_diffusion.ipynb
@@ -0,0 +1,728 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "copy-warning"
+      },
+      "source": [
+        "> **\u26a0\ufe0f IMPORTANT: Make a Copy First!**\n",
+        ">\n",
+        "> Before running this notebook:\n",
+        "> 1. Click **File \u2192 Save a copy in Drive**\n",
+        "> 2. Work in your copy (not this original)\n",
+        ">\n",
+        "> This ensures your progress is saved and you won't lose work!\n",
+        "\n",
+        "---\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "header"
+      },
+      "source": [
+        "# \ud83e\uddd1\u200d\ud83d\udd2c AI Scientist - Generate Your First Research Paper!\n",
+        "\n",
+        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SakanaAI/AI-Scientist/blob/main/notebooks/ai_scientist_2d_diffusion.ipynb)\n",
+        "\n",
+        "Welcome! This notebook lets you run **The AI Scientist** on Google Colab to automatically:\n",
+        "1. Generate novel research ideas\n",
+        "2. Run experiments\n",
+        "3. Write a complete research paper\n",
+        "4. Create publication-ready figures\n",
+        "\n",
+        "**Template**: 2D Diffusion (studies diffusion models on low-dimensional datasets)\n",
+        "\n",
+        "---\n",
+        "\n",
+        "### LLM Provider Comparison\n",
+        "\n",
+        "| Provider | Model | Cost per Paper | Quality | Speed | Recommended For |\n",
+        "|----------|-------|----------------|---------|-------|-----------------|\n",
+        "| **DeepSeek** | deepseek-chat | **$1-2** | Good | Fast | Budget users |\n",
+        "| **OpenAI** | gpt-4o-mini | **$2-4** | Good | Fast | Most users |\n",
+        "| **OpenAI** | gpt-4o | **$5-8** | Great | Medium | Better quality |\n",
+        "| **Anthropic** | Claude Sonnet 3.5 | **$8-12** | **Best** | Medium | Best results |\n",
+        "\n",
+        "### Cost Breakdown by Stage\n",
+        "\n",
+        "| Stage | Time | API Calls | Cost (GPT-4o-mini) | Cost (Claude Sonnet) |\n",
+        "|-------|------|-----------|-------------------|---------------------|\n",
+        "| Setup | 5 min | 0 | Free | Free |\n",
+        "| Baseline Run | 5-10 min | 0 | Free | Free |\n",
+        "| Idea Generation | 2-5 min | ~5-10 | $0.50-1 | $1-2 |\n",
+        "| Experiments (Aider) | 10-20 min | ~20-50 | $1-3 | $4-8 |\n",
+        "| Paper Writing | 5-10 min | ~10-20 | $0.50-1 | $2-3 |\n",
+        "| **TOTAL** | **30-45 min** | **~35-80** | **$2-5** | **$7-13** |\n",
+        "\n",
+        "*Note: Costs scale linearly with number of ideas. Above estimates are for 1 idea.*\n",
+        "\n",
+        "### Cost Control Tips\n",
+        "\n",
+        "1. **Start with 1 idea**: Default setting keeps costs low\n",
+        "2. **Use cheaper models**: DeepSeek or GPT-4o-mini for testing\n",
+        "3. **Skip novelty check**: Saves ~5-10 API calls per idea\n",
+        "4. **Skip review step**: Saves ~10-20 API calls (can review manually)\n",
+        "5. **Use free GPU**: Colab T4 is sufficient, no need for paid tiers\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "setup-header"
+      },
+      "source": [
+        "---\n",
+        "# Part 1: Setup & Installation \n",
+        "\n",
+        "Installing dependencies and configuring the environment."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "check-gpu"
+      },
+      "outputs": [],
+      "source": [
+        "# Check GPU availability\n",
+        "import subprocess\n",
+        "import sys\n",
+        "\n",
+        "gpu_info = subprocess.run(['nvidia-smi'], capture_output=True, text=True)\n",
+        "if gpu_info.returncode == 0:\n",
+        "    print(\"GPU detected: \")\n",
+        "    print(gpu_info.stdout)\n",
+        "else:\n",
+        "    print(\"No GPU detected!\")\n",
+        "    print(\"Please enable GPU: Runtime \u2192 Change runtime type \u2192 T4 GPU\")\n",
+        "    sys.exit(1)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "install-latex"
+      },
+      "outputs": [],
+      "source": [
+        "%%bash\n",
+        "# Install minimal LaTeX (~500MB vs 5GB texlive-full)\n",
+        "echo \"Installing LaTeX (this takes ~3-5 minutes)...\"\n",
+        "apt-get update -qq\n",
+        "apt-get install -y -qq \\\n",
+        "    texlive-latex-base \\\n",
+        "    texlive-latex-extra \\\n",
+        "    texlive-fonts-recommended \\\n",
+        "    texlive-fonts-extra \\\n",
+        "    cm-super \\\n",
+        "    dvipng \\\n",
+        "    chktex \\\n",
+        "    > /dev/null 2>&1\n",
+        "\n",
+        "echo \"LaTeX installed successfully!\"\n",
+        "pdflatex --version | head -1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "clone-repo"
+      },
+      "outputs": [],
+      "source": [
+        "# Clone AI Scientist repository\n",
+        "import os\n",
+        "\n",
+        "if not os.path.exists('AI-Scientist'):\n",
+        "    !git clone https://github.com/SakanaAI/AI-Scientist.git\n",
+        "    print(\"Repository cloned\")\n",
+        "else:\n",
+        "    print(\"Repository already exists\")\n",
+        "\n",
+        "%cd AI-Scientist"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "install-python-deps"
+      },
+      "outputs": [],
+      "source": [
+        "# Install Python dependencies\n",
+        "print(\"Installing Python dependencies...\")\n",
+        "!pip install -q -r requirements.txt\n",
+        "\n",
+        "# Install 2D Diffusion specific dependencies\n",
+        "!git clone https://github.com/gregversteeg/NPEET.git\n",
+        "%cd NPEET\n",
+        "!pip install -q .\n",
+        "%cd ..\n",
+        "!pip install -q scikit-learn\n",
+        "\n",
+        "print(\"All dependencies installed!\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "api-key-header"
+      },
+      "source": [
+        "## Configure API Key\n",
+        "\n",
+        "Choose **one** of these options:\n",
+        "\n",
+        "### Option A: Use Colab Secrets (Recommended)\n",
+        "1. Click the \ud83d\udd11 key icon in the left sidebar\n",
+        "2. Add a new secret:\n",
+        "   - Name: `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, `DEEPSEEK_API_KEY`)\n",
+        "   - Value: Your API key\n",
+        "3. Enable \"Notebook access\"\n",
+        "4. Run the cell below\n",
+        "\n",
+        "### Option B: Manual Entry\n",
+        "Uncomment and run the manual entry cell instead."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "api-key-secrets"
+      },
+      "outputs": [],
+      "source": [
+        "# Option A: Use Colab Secrets\n",
+        "from google.colab import userdata\n",
+        "import os\n",
+        "\n",
+        "# Try to get key from secrets \n",
+        "try:\n",
+        "    os.environ['OPENAI_API_KEY'] = userdata.get('OPENAI_API_KEY')\n",
+        "    print(\"OpenAI API key loaded from secrets\")\n",
+        "    model_choice = \"gpt-4o-mini\"  # Affordable and good quality \n",
+        "except:\n",
+        "    try:\n",
+        "        os.environ['ANTHROPIC_API_KEY'] = userdata.get('ANTHROPIC_API_KEY')\n",
+        "        print(\"Anthropic API key loaded from secrets\")\n",
+        "        model_choice = \"claude-3-5-sonnet-20241022\"  # Best quality\n",
+        "    except:\n",
+        "        try:\n",
+        "            os.environ['DEEPSEEK_API_KEY'] = userdata.get('DEEPSEEK_API_KEY')\n",
+        "            print(\"DeepSeek API key loaded from secrets\")\n",
+        "            model_choice = \"deepseek-chat\"  # Cheapest \n",
+        "        except:\n",
+        "            print(\"No API key found in secrets. Use Option B below.\")\n",
+        "            model_choice = None"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "api-key-manual"
+      },
+      "outputs": [],
+      "source": [
+        "# Option B: Manual Entry (uncomment to use)\n",
+        "import getpass\n",
+        "import os\n",
+        "\n",
+        "print(\"Choose your LLM provider:\")\n",
+        "print(\"1. OpenAI (gpt-4o-mini) - ~$2-4 per paper\")\n",
+        "print(\"2. Anthropic (claude-3-5-sonnet) - ~$8-12 per paper\")\n",
+        "print(\"3. DeepSeek (deepseek-chat) - ~$1-2 per paper\")\n",
+        "\n",
+        "provider = input(\"Enter choice (1/2/3): \").strip()\n",
+        "\n",
+        "if provider == \"1\":\n",
+        "    os.environ['OPENAI_API_KEY'] = getpass.getpass('Enter your OpenAI API key: ')\n",
+        "    model_choice = \"gpt-4o-mini\"\n",
+        "elif provider == \"2\":\n",
+        "    os.environ['ANTHROPIC_API_KEY'] = getpass.getpass('Enter your Anthropic API key: ')\n",
+        "    model_choice = \"claude-3-5-sonnet-20241022\"\n",
+        "elif provider == \"3\":\n",
+        "    os.environ['DEEPSEEK_API_KEY'] = getpass.getpass('Enter your DeepSeek API key: ')\n",
+        "    model_choice = \"deepseek-chat\"\n",
+        "else:\n",
+        "    raise ValueError(\"Invalid choice\")\n",
+        "\n",
+        "print(f\"\\nUsing model: {model_choice}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "verify-setup"
+      },
+      "outputs": [],
+      "source": [
+        "# Verify setup\n",
+        "import torch\n",
+        "\n",
+        "assert model_choice is not None, \"No API key configured! Use Option A or B above.\"\n",
+        "assert torch.cuda.is_available(), \"No GPU available! Change runtime type to GPU.\"\n",
+        "\n",
+        "print(\"\\n\" + \"=\"*50)\n",
+        "print(\"Setup Complete!\")\n",
+        "print(\"=\"*50)\n",
+        "print(f\"Model: {model_choice}\")\n",
+        "print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
+        "print(f\"Working directory: {os.getcwd()}\")\n",
+        "print(\"=\"*50)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "baseline-header"
+      },
+      "source": [
+        "---\n",
+        "# Part 2: Generate Baseline Run \n",
+        "\n",
+        "Running the baseline experiment to establish performance metrics.\n",
+        "\n",
+        "This trains a simple diffusion model on 2D datasets "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "run-baseline"
+      },
+      "outputs": [],
+      "source": [
+        "import subprocess\n",
+        "import os\n",
+        "\n",
+        "template_dir = \"templates/2d_diffusion\"\n",
+        "os.chdir(template_dir)\n",
+        "\n",
+        "print(\"\ud83d\udd2c Running baseline experiment...\")\n",
+        "print(\"This will train a diffusion model on 2D datasets.\")\n",
+        "print(\"Expected time: 5-10 minutes\\n\")\n",
+        "\n",
+        "# Run baseline\n",
+        "result = subprocess.run(\n",
+        "    [\"python\", \"experiment.py\", \"--out_dir=run_0\"],\n",
+        "    capture_output=False,  # Show output\n",
+        "    text=True\n",
+        ")\n",
+        "\n",
+        "if result.returncode == 0:\n",
+        "    print(\"\\nBaseline experiment completed!\")\n",
+        "else:\n",
+        "    print(\"\\nBaseline experiment failed!\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "show-baseline-results"
+      },
+      "outputs": [],
+      "source": [
+        "# Display baseline results\n",
+        "import json\n",
+        "from IPython.display import Image, display\n",
+        "import os\n",
+        "\n",
+        "# Show final metrics\n",
+        "with open(\"run_0/final_info.json\", \"r\") as f:\n",
+        "    baseline_results = json.load(f)\n",
+        "\n",
+        "print(\"Baseline Results:\")\n",
+        "print(\"=\" * 50)\n",
+        "for dataset, metrics in baseline_results.items():\n",
+        "    print(f\"\\n{dataset}:\")\n",
+        "    for metric, value in metrics.items():\n",
+        "        print(f\"  {metric}: {value}\")\n",
+        "\n",
+        "# Generate and show plots\n",
+        "print(\"\\nGenerating plots...\")\n",
+        "!python plot.py\n",
+        "\n",
+        "# Display generated plots\n",
+        "plot_files = [f for f in os.listdir('.') if f.endswith('.png')]\n",
+        "print(f\"\\nGenerated {len(plot_files)} plots:\\n\")\n",
+        "for plot_file in sorted(plot_files)[:3]:  # Show first 3 plots\n",
+        "    print(f\"Showing: {plot_file}\")\n",
+        "    display(Image(filename=plot_file, width=400))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ideas-header"
+      },
+      "source": [
+        "---\n",
+        "# Part 3: Generate Research Ideas\n",
+        "\n",
+        "The AI will now generate novel research ideas based on the baseline results.\n",
+        "\n",
+        "**Cost control**: Default to generating 1 idea. You can increase this, but it will be costly!"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "config-ideas"
+      },
+      "outputs": [],
+      "source": [
+        "# Configuration\n",
+        "NUM_IDEAS = 1  # \u26a0\ufe0f You can increase this, but it will be costly!\n",
+        "SKIP_NOVELTY_CHECK = True  # Set False to check against existing literature. Slower, more API calls\n",
+        "\n",
+        "print(f\"Configuration:\")\n",
+        "print(f\"  Model: {model_choice}\")\n",
+        "print(f\"  Number of ideas: {NUM_IDEAS}\")\n",
+        "print(f\"  Novelty check: {'Disabled' if SKIP_NOVELTY_CHECK else 'Enabled'}\")\n",
+        "print(f\"\\nEstimated cost for idea generation: ${NUM_IDEAS * 0.5:.2f} - ${NUM_IDEAS * 2:.2f}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "generate-ideas"
+      },
+      "outputs": [],
+      "source": [
+        "# Generate ideas by calling function directly\n",
+        "from ai_scientist.generate_ideas import generate_ideas\n",
+        "from ai_scientist.llm import create_client\n",
+        "\n",
+        "print(\"Generating research ideas...\\n\")\n",
+        "\n",
+        "  # Create client\n",
+        "client, client_model = create_client(model_choice)\n",
+        "\n",
+        "  # Generate ideas with controlled count\n",
+        "ideas = generate_ideas(\n",
+        "  base_dir=\"templates/2d_diffusion\",\n",
+        "  client=client,\n",
+        "  model=client_model,\n",
+        "  skip_generation=False,\n",
+        "  max_num_generations=NUM_IDEAS,  \n",
+        "  num_reflections=3,\n",
+        ")\n",
+        "\n",
+        "print(f\"\\nGenerated {len(ideas)} ideas!\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "show-ideas"
+      },
+      "outputs": [],
+      "source": [
+        "# Display generated ideas\n",
+        "import json\n",
+        "\n",
+        "with open(\"templates/2d_diffusion/ideas.json\", \"r\") as f:\n",
+        "    ideas = json.load(f)\n",
+        "\n",
+        "print(\"\\n\" + \"=\"*70)\n",
+        "print(\"Generated Research Ideas\")\n",
+        "print(\"=\"*70)\n",
+        "\n",
+        "novel_ideas = [idea for idea in ideas if idea.get(\"novel\", True)]\n",
+        "\n",
+        "for i, idea in enumerate(novel_ideas, 1):\n",
+        "    print(f\"\\n{'='*70}\")\n",
+        "    print(f\"Idea {i}: {idea['Title']}\")\n",
+        "    print(f\"{'='*70}\")\n",
+        "    print(f\"Name: {idea['Name']}\")\n",
+        "    print(f\"\\nExperiment Plan:\\n{idea['Experiment']}\")\n",
+        "    print(f\"\\nRatings:\")\n",
+        "    print(f\"  Interestingness: {idea.get('Interestingness', 'N/A')}/10\")\n",
+        "    print(f\"  Feasibility: {idea.get('Feasibility', 'N/A')}/10\")\n",
+        "    print(f\"  Novelty: {idea.get('Novelty', 'N/A')}/10\")\n",
+        "\n",
+        "print(f\"\\n{'='*70}\")\n",
+        "print(f\"Total novel ideas: {len(novel_ideas)}\")\n",
+        "print(f\"{'='*70}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "experiments-header"
+      },
+      "source": [
+        "---\n",
+        "# Part 4: Run Experiments\n",
+        "\n",
+        "The AI will now:\n",
+        "1. Modify the experiment code to implement the idea\n",
+        "2. Run multiple experiments\n",
+        "3. Collect and analyze results\n",
+        "\n",
+        "**Note**: This uses Aider (AI pair programmer) to write code autonomously."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Add \"novel\" field to all ideas\n",
+        "import json\n",
+        "\n",
+        "ideas_file = \"templates/2d_diffusion/ideas.json\"\n",
+        "\n",
+        "# Load ideas\n",
+        "with open(ideas_file, \"r\") as f:\n",
+        "    ideas = json.load(f)\n",
+        "\n",
+        "for idea in ideas:\n",
+        "    if \"novel\" not in idea:\n",
+        "        idea[\"novel\"] = True\n",
+        "\n",
+        "\n",
+        "with open(ideas_file, \"w\") as f:\n",
+        "    json.dump(ideas, f, indent=4)\n",
+        "\n",
+        "print(f\"Fixed {len(ideas)} ideas \u2014 added 'novel' field\")\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "config-experiments"
+      },
+      "outputs": [],
+      "source": [
+        "# Part 4: Run AI Scientist  \n",
+        "\n",
+        "print(\"=\" * 70)\n",
+        "print(\"Run FULL AI Scientist Pipeline\")\n",
+        "print(\"=\" * 70)\n",
+        "\n",
+        "import os\n",
+        "import subprocess\n",
+        "import json\n",
+        "\n",
+        "# Check environment\n",
+        "os.chdir(\"/content/AI-Scientist\")\n",
+        "print(\"\\n\u2713 Current directory:\", os.getcwd())\n",
+        "print(\"\u2713 Ideas file:\", \"EXISTS\" if os.path.exists(\"templates/2d_diffusion/ideas.json\") else \"MISSING\")\n",
+        "print(\"\u2713 Launch script:\", \"EXISTS\" if os.path.exists(\"launch_scientist.py\") else \"MISSING\")\n",
+        "\n",
+        "# Try to see what ideas we have\n",
+        "try:\n",
+        "    with open(\"templates/2d_diffusion/ideas.json\", \"r\") as f:\n",
+        "        ideas = json.load(f)\n",
+        "    print(f\"\u2713 Found {len(ideas)} ideas\")\n",
+        "    print(f\"\u2713 First idea: {ideas[0]['Name']}\")\n",
+        "except Exception as e:\n",
+        "    print(f\"\u2717 Error reading ideas: {e}\")\n",
+        "\n",
+        "print(\"\\n\" + \"=\" * 70)\n",
+        "print(\"RUNNING....\")\n",
+        "print(\"=\" * 70 + \"\\n\")\n",
+        "\n",
+        "# Run with FULL error capture\n",
+        "result = subprocess.run([\n",
+        "    \"python\", \"launch_scientist.py\",\n",
+        "    \"--model\", model_choice,\n",
+        "    \"--experiment\", \"2d_diffusion\",\n",
+        "    \"--num-ideas\", \"1\",\n",
+        "    \"--skip-idea-generation\",\n",
+        "    \"--skip-novelty-check\",\n",
+        "    \"--parallel\", \"0\"\n",
+        "], capture_output=True, text=True)\n",
+        "\n",
+        "print(\"STDOUT:\")\n",
+        "print(result.stdout)\n",
+        "\n",
+        "print(\"\\nSTDERR:\")\n",
+        "print(result.stderr)\n",
+        "\n",
+        "print(\"\\nReturn code:\", result.returncode)\n",
+        "\n",
+        "if result.returncode != 0:\n",
+        "    print(\"\\n\" + \"=\" * 70)\n",
+        "    print(\"FAILED - See errors above\")\n",
+        "    print(\"=\" * 70)\n",
+        "else:\n",
+        "    print(\"\\n\" + \"=\" * 70)\n",
+        "    print(\"\ud83c\udf89 SUCCESS!\")\n",
+        "    print(\"=\" * 70)\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "paper-header"
+      },
+      "source": [
+        "### Part 5 \u2014 View and Download Generated Paper\n",
+        "\n",
+        "This section locates the most recently generated research paper from the `2d_diffusion` experiment, displays it directly within the notebook, and provides download options for both the paper and the full experiment results.\n",
+        "\n",
+        "#### What this cell does\n",
+        "1. Searches the results directory for generated `.pdf` files.  \n",
+        "2. Displays the most recent paper inline for quick review.  \n",
+        "3. Allows you to download either the paper alone or the entire results folder as a `.zip` archive.\n",
+        "\n",
+        "#### If no PDF is found\n",
+        "- Confirm that the previous pipeline cell finished successfully.  \n",
+        "- Review the console output above for any runtime or file-generation errors.  \n",
+        "\n",
+        "#### To generate a new paper manually\n",
+        "```bash\n",
+        "python launch_scientist.py --model {model_choice} --experiment 2d_diffusion --num-ideas 1\n",
+        "```\n",
+        "---"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "show-example-paper"
+      },
+      "outputs": [],
+      "source": [
+        "# View and Download Generated Paper\n",
+        "\n",
+        "import os\n",
+        "import glob\n",
+        "from IPython.display import IFrame, display\n",
+        "from google.colab import files\n",
+        "\n",
+        "print(\"Looking for generated paper...\")\n",
+        "\n",
+        "results_dir = \"/content/AI-Scientist/results/2d_diffusion\"\n",
+        "pdf_files = glob.glob(f\"{results_dir}/*/*.pdf\")\n",
+        "\n",
+        "if not pdf_files:\n",
+        "    print(\"No PDF found.\")\n",
+        "    print(f\"Checked directory: {results_dir}\")\n",
+        "    print(\"\\nTroubleshooting:\")\n",
+        "    print(\"- Check if the pipeline completed successfully above.\")\n",
+        "    print(\"- Look for error messages in the output.\")\n",
+        "else:\n",
+        "    # Get the most recent PDF\n",
+        "    latest_pdf = max(pdf_files, key=os.path.getctime)\n",
+        "    paper_name = os.path.basename(latest_pdf)\n",
+        "\n",
+        "    print(f\"\u2713 Found paper: {paper_name}\")\n",
+        "    print(f\"Location: {latest_pdf}\")\n",
+        "    print(f\"Size: {os.path.getsize(latest_pdf) / 1024:.1f} KB\")\n",
+        "\n",
+        "    print(\"\\n\" + \"=\" * 70)\n",
+        "    print(\"Displaying generated paper\")\n",
+        "    print(\"=\" * 70 + \"\\n\")\n",
+        "\n",
+        "    display(IFrame(latest_pdf, width=900, height=1200))\n",
+        "\n",
+        "    print(\"\\n\" + \"=\" * 70)\n",
+        "    print(\"Download options\")\n",
+        "    print(\"=\" * 70)\n",
+        "\n",
+        "    # Download the PDF\n",
+        "    print(\"\\nDownloading PDF to your computer...\")\n",
+        "    files.download(latest_pdf)\n",
+        "\n",
+        "    print(\"\\nTo download the full results (code, plots, and paper), run this cell:\")\n",
+        "    print(f\"\"\"\n",
+        "import shutil\n",
+        "shutil.make_archive('/content/full_results', 'zip', '{os.path.dirname(latest_pdf)}')\n",
+        "from google.colab import files\n",
+        "files.download('/content/full_results.zip')\n",
+        "\"\"\")\n",
+        "\n",
+        "    print(\"\\nProcess complete. The AI-generated paper is ready.\")\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "next-steps-header"
+      },
+      "source": [
+        "---\n",
+        "# Run Complete \u2014 AI Scientist \n",
+        "Congratulations! You've successfully run The AI Scientist to generate a complete research paper.\n",
+        "You\u2019ve successfully:\n",
+        "- Installed and configured the AI Scientist framework.\n",
+        "- Validated your environment with the debug pipeline.\n",
+        "- Generated, visualized, and downloaded a full AI-authored research paper.\n",
+        "\n",
+        "### Next recommended steps\n",
+        "- Explore other experiment templates in the `experiments/` folder.\n",
+        "- Modify existing templates or create your own custom experiments.\n",
+        "\n",
+        "### Try other templates:\n",
+        "- `grokking` - Studies sudden generalization in neural networks\n",
+        "- `nanoGPT` - Experiments with small language models\n",
+        "- `probes` - LLM steering and interpretability\n",
+        "\n",
+        "### Learn more:\n",
+        "- \ud83d\udcc4 [Paper](https://arxiv.org/abs/2408.06292)\n",
+        "- \ud83d\udcdd [Blog Post](https://sakana.ai/ai-scientist/)\n",
+        "- \ud83d\udcbb [GitHub](https://github.com/SakanaAI/AI-Scientist)\n",
+        "- \ud83d\udcc2 [Example Papers](https://drive.google.com/drive/folders/1G7A0wTqfXVa-cpexjk0oaXakaSJwffEt)\n",
+        "\n",
+        "### Share your results!\n",
+        "- Tweet about your generated papers with #AIScientist\n",
+        "- Open issues on GitHub if you find bugs\n",
+        "- Contribute new templates!\n",
+        "\n",
+        "---\n",
+        "\n",
+        "**Questions or issues?** Check the [GitHub Issues](https://github.com/SakanaAI/AI-Scientist/issues) or [Discussion](https://github.com/SakanaAI/AI-Scientist/discussions)\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3 (ipykernel)",
+      "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.12.12"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file
diff --git a/requirements/requirements-cpu.txt b/requirements/requirements-cpu.txt
new file mode 100644
index 00000000..10c5eb73
--- /dev/null
+++ b/requirements/requirements-cpu.txt
@@ -0,0 +1,15 @@
+# PyTorch CPU-only (no CUDA)
+# For systems without NVIDIA GPUs or for testing
+
+torch==2.5.1
+torchvision==0.20.1
+torchaudio==2.5.1
+
+# Note: CPU-only will be significantly slower for training
+# Only recommended for:
+# - Testing/development
+# - Idea generation (no GPU needed)
+# - Paper writing (no GPU needed)
+#
+# NOT recommended for:
+# - Running actual experiments (will be very slow)
diff --git a/requirements/requirements-cuda.txt b/requirements/requirements-cuda.txt
new file mode 100644
index 00000000..7f79a565
--- /dev/null
+++ b/requirements/requirements-cuda.txt
@@ -0,0 +1,18 @@
+# PyTorch with CUDA Support
+# For systems with NVIDIA GPUs
+
+# PyTorch 2.5.1 with CUDA 12.1 (recommended)
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.5.1+cu121
+torchvision==0.20.1+cu121
+torchaudio==2.5.1+cu121
+
+# Alternative: CUDA 11.8 (for older GPUs)
+# Uncomment the lines below and comment out the above if using CUDA 11.8
+# --extra-index-url https://download.pytorch.org/whl/cu118
+# torch==2.5.1+cu118
+# torchvision==0.20.1+cu118
+# torchaudio==2.5.1+cu118
+
+# Note: Verify your CUDA version with: nvidia-smi
+# Must match system CUDA version closely (11.8 or 12.1 most common)
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
new file mode 100644
index 00000000..3a9ef242
--- /dev/null
+++ b/requirements/requirements.txt
@@ -0,0 +1,68 @@
+# AI Scientist - Pinned Dependencies
+# Tested on: Ubuntu 22.04, Python 3.10-3.12
+# Last updated: 2025-01-26
+
+# ============================================================================
+# LLM APIs
+# ============================================================================
+anthropic==0.40.0
+aider-chat==0.66.1
+backoff==2.2.1
+openai==1.55.3
+google-generativeai==0.8.3
+
+# ============================================================================
+# Visualization & Document Processing
+# ============================================================================
+matplotlib==3.9.2
+pypdf==5.1.0
+pymupdf4llm==0.0.17
+
+# ============================================================================
+# Machine Learning Core
+# ============================================================================
+# Note: torch is specified separately in requirements-cuda.txt or requirements-cpu.txt
+numpy==1.26.4
+transformers==4.46.3
+datasets==3.1.0
+tiktoken==0.8.0
+tqdm==4.67.1
+
+# ============================================================================
+# Experiment Tracking
+# ============================================================================
+wandb==0.18.7
+
+# ============================================================================
+# Template-Specific Dependencies
+# ============================================================================
+# These are needed for specific templates but don't hurt to have installed
+
+# For 2d_diffusion template
+scikit-learn==1.5.2
+
+# For grokking template
+einops==0.8.0
+
+# For probes template
+pandas==2.2.3
+seaborn==0.13.2
+ipykernel==6.29.5
+transformer-lens==2.7.0
+
+# ============================================================================
+# Notes:
+# ============================================================================
+# - PyTorch is NOT included here. Install separately via:
+#   - requirements-cuda.txt for CUDA support
+#   - requirements-cpu.txt for CPU-only
+#
+# - Template-specific dependencies that require compilation:
+#   - NPEET (2d_diffusion): Install via git clone + pip install .
+#   - transformer_lens (probes): Included via pip
+#
+# - LaTeX (required for paper generation):
+#   - Ubuntu/Debian: sudo apt-get install texlive-full
+#   - Minimal install: See docs/INSTALLATION.md
+#
+# ============================================================================