02_lipogram_grpo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# \ud83d\udd24 Teach a Model a *Lipogram* \u2014 GRPO + reward design\n",
    "\n",
    "Train **Qwen3.5-2B** to answer questions **without ever using the letter \u201ce\u201d** \u2014 a *lipogram*.\n",
    "Inspired by the [Thinking Machines \u201cInkling\u201d demo](https://thinkingmachines.ai/news/introducing-inkling/).\n",
    "\n",
    "Why this makes a great RL tutorial: the constraint is **verifiable** (just check for \u201ce\u201d), so the\n",
    "reward is easy to write \u2014 and that's exactly why it's the perfect place to learn **reward design** and\n",
    "the classic **reward-hacking** failure mode. A lazy reward gets gamed (the model spits out short,\n",
    "letter-free gibberish). We'll *watch* that happen, then fix it.\n",
    "\n",
    "Unlike notebook 01 (which used an OpenEnv server for the reward), here the reward is a plain **Python\n",
    "function you own** \u2014 the other way rewards enter GRPO. Only a **Hugging Face token** is needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%pip install -q -U \"trl>=0.19\" datasets accelerate trackio 2>/dev/null\n",
    "print('\u2705 installed')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## \ud83d\udd11 Log in to Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login, whoami\n",
    "try:\n",
    "    who = whoami()  # already logged in (HF_TOKEN in the job)\n",
    "except Exception:\n",
    "    notebook_login()\n",
    "    who = whoami()\n",
    "print('\ud83d\udc64 logged in as:', who['name'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Settings"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1'\n",
    "USERNAME          = whoami()['name']\n",
    "\n",
    "MODEL             = 'Qwen/Qwen3.5-2B'      # small text model; any chat LLM works\n",
    "FORBIDDEN_LETTER  = 'e'                     # the letter the model must avoid\n",
    "DATASET           = 'yahma/alpaca-cleaned'  # open-ended instructions -> avoiding 'e' is genuinely hard\n",
    "\n",
    "NUM_TRAIN_SAMPLES = 2000\n",
    "MAX_STEPS         = 40                       # short demo; bump to 300+ to really move it\n",
    "NUM_GENERATIONS   = 8\n",
    "MAX_COMPLETION    = 256\n",
    "LEARNING_RATE     = 1e-5\n",
    "REWARD_K          = 4.0                      # steepness of the letter penalty (higher = harsher)\n",
    "MIN_WORDS         = 6                        # anti-hack: answers below this score 0\n",
    "\n",
    "USE_TRACKIO       = True\n",
    "TRACKIO_PROJECT   = 'lipogram-grpo'\n",
    "TRACKIO_SPACE_ID  = f'{USERNAME}/trackio-lipogram'\n",
    "print(f'\ud83e\udde0 {MODEL}  \u00b7  forbid \"{FORBIDDEN_LETTER}\"  \u00b7  {MAX_STEPS} steps \u00d7 {NUM_GENERATIONS} gens')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The reward \u2014 and how it gets *hacked* \ud83c\udfaf\n",
    "\n",
    "**This is the whole point of the notebook.** The obvious reward is \u201cpenalize the letter e.\u201d But if that's\n",
    "*all* you reward, the model discovers it can win by producing **short, repetitive, letter-free gibberish**\n",
    "\u2014 maximal reward, zero usefulness. That's **reward hacking**.\n",
    "\n",
    "So our reward is a **product of three terms**, so all must hold:\n",
    "\n",
    "```\n",
    "reward = lipogram_score  \u00d7  length_guard  \u00d7  diversity_guard\n",
    "```\n",
    "- **lipogram_score** \u2014 steep `exp(-k \u00b7 e_per_word)`: 1.0 with zero \u201ce\u201d, dropping fast per \u201ce\u201d.\n",
    "- **length_guard** \u2014 0 below `MIN_WORDS` (kills 1-word hacks), decays for rambling.\n",
    "- **diversity_guard** \u2014 penalizes low unique-word ratio (kills \u201cno no no no\u201d).\n",
    "\n",
    "Run the cell \u2014 it scores a few hand-written answers so you can *see* the naive-vs-guarded difference."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "L = FORBIDDEN_LETTER.lower()\n",
    "\n",
    "def naive_reward(text):\n",
    "    \"\"\"The lazy reward: just 1 - fraction of words containing the letter. Gets hacked.\"\"\"\n",
    "    w = text.split();  wc = max(len(w), 1)\n",
    "    return 1.0 - sum(1 for x in w if L in x.lower()) / wc\n",
    "\n",
    "def guarded_reward(text, k=REWARD_K, min_words=MIN_WORDS, max_words=120, div_floor=0.4):\n",
    "    \"\"\"Composite reward: avoid the letter AND stay a real, varied answer.\"\"\"\n",
    "    w = text.split();  wc = len(w)\n",
    "    if wc < min_words:              # empty / one-word letter-free hack\n",
    "        return 0.0\n",
    "    n = text.lower().count(L)\n",
    "    lipo   = math.exp(-k * (n / wc))               # 1.0 at zero 'e', steep penalty per 'e'\n",
    "    length = 1.0 if wc <= max_words else max(0.0, 1 - (wc - max_words) / max_words)\n",
    "    div    = min(1.0, (len(set(x.lower() for x in w)) / wc) / div_floor)\n",
    "    return lipo * length * div\n",
    "\n",
    "examples = {\n",
    "    'normal answer (lots of e)': 'The weather here is really pleasant and the trees are green everywhere.',\n",
    "    'HACK: one word, no e'     : 'Ok.',\n",
    "    'HACK: repetition, no e'   : 'blah blah blah blah blah blah blah blah',\n",
    "    'GOOD: real answer, no e'  : 'My dog ran fast across a wide grassy hill on a warm sunny day, glad and calm.',\n",
    "}\n",
    "print(f\"{'completion':30s} {'naive':>7} {'guarded':>8}\")\n",
    "for name, t in examples.items():\n",
    "    print(f'{name:30s} {naive_reward(t):7.2f} {guarded_reward(t):8.2f}')\n",
    "print('\\n\ud83d\udc46 naive rewards the gibberish hacks ~1.0; guarded gives them ~0 and only the real e-free answer wins.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The GRPO reward function\n",
    "GRPO calls a reward function with a batch of `completions` and returns one score each. We wrap the\n",
    "guarded reward above."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "SUFFIX = (f\"\\n\\nRules: (1) never use the letter '{FORBIDDEN_LETTER}' anywhere in your reply; \"\n",
    "          f\"(2) keep it brief and to the point.\")\n",
    "\n",
    "def lipogram_reward(completions, **kwargs):\n",
    "    return [guarded_reward((c[0]['content'] or '').strip()) for c in completions]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. The dataset\n",
    "Open-ended alpaca instructions, each with the \u201cdon't use e\u201d rule appended. GRPO only needs the\n",
    "**prompt** (the reward scores the model's own answer \u2014 there's no gold label to match)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from datasets import load_dataset, Dataset\n",
    "raw = load_dataset(DATASET, split=f'train[:{NUM_TRAIN_SAMPLES}]')\n",
    "prompts = []\n",
    "for r in raw:\n",
    "    q = (r.get('instruction') or '').strip()\n",
    "    inp = (r.get('input') or '').strip()\n",
    "    if inp: q = f'{q}\\n\\n{inp}'\n",
    "    if q and len(q) <= 600:\n",
    "        prompts.append([{'role':'user','content': q + SUFFIX}])\n",
    "train_dataset = Dataset.from_dict({'prompt': prompts})\n",
    "EVAL_QUESTIONS = [\n",
    "    'What are the benefits of regular exercise?',\n",
    "    'Explain how photosynthesis works.',\n",
    "    'Give three tips for staying productive.',\n",
    "    'Describe your favorite season and why you like it.',\n",
    "    'How do airplanes stay in the air?',\n",
    "]\n",
    "print(f'\ud83d\udcda {len(train_dataset)} training prompts')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Baseline \u2014 how often is it e-free *now*?\n",
    "Spoiler: almost never. Normal English is saturated with \u201ce\u201d."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "tok = AutoTokenizer.from_pretrained(MODEL)\n",
    "def load_model():\n",
    "    return AutoModelForCausalLM.from_pretrained(MODEL, dtype='bfloat16', device_map='cuda')\n",
    "\n",
    "def evaluate(model, questions=EVAL_QUESTIONS, show=True):\n",
    "    model.eval(); ok = 0; samples = []\n",
    "    for q in questions:\n",
    "        msgs = [{'role':'user','content': q + SUFFIX}]\n",
    "        text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)\n",
    "        ins = tok(text, return_tensors='pt').to(model.device)\n",
    "        with torch.no_grad():\n",
    "            out = model.generate(**ins, max_new_tokens=MAX_COMPLETION, do_sample=False)\n",
    "        ans = tok.decode(out[0][ins['input_ids'].shape[1]:], skip_special_tokens=True).strip()\n",
    "        e = ans.lower().count(L); ok += (e == 0 and len(ans.split()) >= MIN_WORDS)\n",
    "        samples.append((q, ans, e))\n",
    "    rate = ok / len(questions)\n",
    "    if show:\n",
    "        print(f'lipogram success: {rate*100:.0f}%  (answers with ZERO \"{FORBIDDEN_LETTER}\")')\n",
    "        for q, a, e in samples[:2]: print(f'  Q: {q}\\n  A: {a[:160]}  [e={e}]\\n')\n",
    "    return rate\n",
    "\n",
    "base_model = load_model()\n",
    "print('=== BEFORE training ==='); baseline = evaluate(base_model)\n",
    "del base_model; torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Train \ud83c\udfcb\ufe0f\n",
    "`GRPOTrainer` samples `NUM_GENERATIONS` answers per prompt, scores them with `lipogram_reward`, and\n",
    "nudges the model toward the higher-reward (e-free, real) ones. Short demo run."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from trl import GRPOConfig, GRPOTrainer\n",
    "cfg = GRPOConfig(\n",
    "    output_dir='lipogram-grpo',\n",
    "    model_init_kwargs={'dtype':'bfloat16'},\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    num_generations=NUM_GENERATIONS,\n",
    "    per_device_train_batch_size=NUM_GENERATIONS,\n",
    "    steps_per_generation=1, gradient_accumulation_steps=1,\n",
    "    max_steps=MAX_STEPS, max_completion_length=MAX_COMPLETION,\n",
    "    temperature=1.0,\n",
    "    chat_template_kwargs={'enable_thinking': False},\n",
    "    bf16=True, gradient_checkpointing=True, gradient_checkpointing_kwargs={'use_reentrant': False},\n",
    "    logging_steps=1, save_strategy='no',\n",
    "    report_to=('trackio' if USE_TRACKIO else 'none'),\n",
    "    project=TRACKIO_PROJECT, trackio_space_id=(TRACKIO_SPACE_ID if USE_TRACKIO else None),\n",
    "    run_name='lipogram-grpo',\n",
    ")\n",
    "trainer = GRPOTrainer(model=MODEL, args=cfg, train_dataset=train_dataset, reward_funcs=lipogram_reward)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Did it learn to avoid \u201ce\u201d?"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "print('=== AFTER training ==='); after = evaluate(trainer.model)\n",
    "print(f'\\nlipogram success:  {baseline*100:.0f}%  \u2192  {after*100:.0f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recap\n",
    "- An RL **reward is just a function** of the model's output \u2014 here, a Python function you fully control.\n",
    "- The **naive** reward was hackable (gibberish scored ~1.0); the **guarded** reward (letter \u00d7 length \u00d7\n",
    "  diversity) made the model produce *real, e-free* answers.\n",
    "- Same GRPO machinery as notebook 01 \u2014 the only difference is **where the reward comes from**\n",
    "  (an OpenEnv server there, an in-process function here).\n",
    "- 40 steps is a demo; a real lipogram model needs a few hundred steps + a stronger LR.\n",
    "\n",
    "**Takeaway for your own envs: the reward is the whole game.** Design it so the *only* way to score high\n",
    "is to actually do the task."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}