Tensorlake Documentation

Train a language model to write correct Python functions with reinforcement learning (RL). Your training process never runs the untrusted code the model generates. The setup has 2 phases: a supervised fine-tuning (SFT) warmup, then Group Sequence Policy Optimization (GSPO) fine-tuning. A hidden pytest suite scores every completion inside an isolated Tensorlake sandbox.

How it works

Training alternates between generating code and scoring it in sandboxes:

SFT warmup (Phase 1): A supervised pass on correct solutions so the model starts generating valid Python. Without it, all completions score 0, reward variance is 0, and the RL trainer has no gradient signal.
GSPO fine-tuning (Phase 2): GRPOTrainer with importance_sampling_level="sequence" generates G completions per step and dispatches them to G parallel sandboxes.
Sandbox reward: Each sandbox runs a hidden pytest suite against the model’s code and returns tests_passed / total as the reward, from 0.0 to 1.0.

Model-generated code is untrusted, so it must not run in your training process. Each completion runs in its own isolated sandbox instead.

GSPO vs GRPO

GSPO and Group Relative Policy Optimization (GRPO) both use clipped importance sampling (IS), but at different granularities:

Algorithm	IS clipping
GRPO	`clip(π_θ(t) / π_old(t))` per token
GSPO	`clip(∏_t π_θ(t) / π_old(t))` once per sequence

For long function bodies, token-level clipping lets noisy individual tokens dominate the gradient. Sequence-level clipping treats the entire trajectory as one unit, which fits code generation better.

Prerequisites

You need Python, the packages below, and a Tensorlake API key:

pip install tensorlake transformers trl datasets torch rich python-dotenv

Create a .env file in your project root with your Tensorlake API key. The script calls load_dotenv() to read it:

TENSORLAKE_API_KEY="your_api_key_here"

TypeScript SDK starter

The Node.js version keeps the same reward oracle. It writes each completion into its own sandbox, runs the hidden pytest suite there, and uses the pass ratio as the reward.

import { Sandbox } from "tensorlake";

const encoder = new TextEncoder();

async function scoreCompletion(
  solutionSource: string,
  hiddenTests: string,
): Promise<number> {
  const sandbox = await Sandbox.create({
    cpus: 1.0,
    memoryMb: 1024,
    timeoutSecs: 300,
    allowInternetAccess: false,
  });

  try {
    await sandbox.writeFile("/workspace/solution.py", encoder.encode(solutionSource));
    await sandbox.writeFile("/workspace/test_hidden.py", encoder.encode(hiddenTests));

    await sandbox.run("python", {
      args: ["-m", "pip", "install", "pytest", "--user", "--break-system-packages"],
    });

    const result = await sandbox.run("python", {
      args: ["-m", "pytest", "-q", "/workspace/test_hidden.py"],
      workingDir: "/workspace",
      timeout: 300,
    });

    const passed = Number(result.stdout.match(/(\d+) passed/)?.[1] ?? 0);
    const failed = Number(result.stdout.match(/(\d+) failed/)?.[1] ?? 0);
    return passed / Math.max(1, passed + failed);
  } finally {
    await sandbox.terminate();
  }
}

const hiddenTests = `
from solution import sum_list
def test_basic():
    assert sum_list([1, 2, 3]) == 6
`;

const completions = [
  "def sum_list(nums):\n    return sum(nums)",
  "def sum_list(nums):\n    return 0",
];

const rewards = await Promise.all(
  completions.map((completion) => scoreCompletion(completion, hiddenTests)),
);

console.log(rewards);
client.close();

This reward function works with the GSPO loop in the full example below. Keep the model and trainer in Python, and move only the sandbox evaluation to TypeScript if your orchestration layer already runs there.

Full example

The script runs 4 stages in order: baseline evaluation, SFT warmup, GSPO fine-tuning, and final evaluation. Pass --smoke for a 5-minute CPU run with 3 tasks, 20 SFT steps, and 1 GSPO epoch.

"""
RL GSPO Reasoner: Code Generation with Hidden Test Suites
===========================================================
Algorithm : GSPO, Group Sequence Policy Optimization (Zheng et al., 2507.18071)
            GRPOConfig(importance_sampling_level="sequence")

Why sandboxes are non-negotiable here
--------------------------------------
The model generates arbitrary Python function bodies.  Running untrusted
model-generated code in the training process directly would be unsafe.
Each completion is executed inside an isolated TensorLake sandbox.
The sandbox runs a hidden pytest suite and returns tests_passed/total as reward.

Training strategy
-----------------
  Phase 1 - SFT warmup (N steps):
    Supervised pass on correct solutions so the model outputs valid Python.
    Without this, all G completions score 0 → reward_std=0 → no gradient.
  Phase 2 - GSPO fine-tuning:
    GRPOTrainer with sequence-level IS.  The reward function dispatches G
    parallel sandboxes per step and prints every completion that scores > 0.

Smoke : --smoke  →  3 functions, 20 SFT steps, 1 GSPO epoch  (~5 min CPU)
Full  : 10 functions, 60 SFT steps, 3 GSPO epochs            (~30 min CPU)
"""

from dotenv import load_dotenv
load_dotenv()

import re
import sys
import textwrap
import torch
from concurrent.futures import ThreadPoolExecutor, as_completed
from datasets import Dataset
from torch.optim import AdamW
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import GRPOTrainer, GRPOConfig
from tensorlake.sandbox import Sandbox
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.rule import Rule
from rich import box
from typing import List

console = Console()

MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
OUTPUT_DIR = "./gspo_coder"
SMOKE      = "--smoke" in sys.argv


# ─── Dataset ──────────────────────────────────────────────────────────────────

TASKS = [
    dict(
        name="sum_list",
        prompt=(
            "Write a Python function:\n\n"
            "def sum_list(nums: list) -> int:\n"
            '    """Return the sum of all integers in nums."""'
        ),
        tests=textwrap.dedent("""\
            from solution import sum_list
            def test_empty():   assert sum_list([]) == 0
            def test_single():  assert sum_list([5]) == 5
            def test_mixed():   assert sum_list([1, 2, 3]) == 6
            def test_neg():     assert sum_list([-1, -2, 3]) == 0
        """),
        solution="def sum_list(nums: list) -> int:\n    return sum(nums)",
    ),
    dict(
        name="is_palindrome",
        prompt=(
            "Write a Python function:\n\n"
            "def is_palindrome(s: str) -> bool:\n"
            '    """Return True if s reads the same forwards and backwards."""'
        ),
        tests=textwrap.dedent("""\
            from solution import is_palindrome
            def test_yes():    assert is_palindrome("racecar") is True
            def test_no():     assert is_palindrome("hello") is False
            def test_empty():  assert is_palindrome("") is True
            def test_single(): assert is_palindrome("a") is True
        """),
        solution="def is_palindrome(s: str) -> bool:\n    return s == s[::-1]",
    ),
    dict(
        name="fizzbuzz",
        prompt=(
            "Write a Python function:\n\n"
            "def fizzbuzz(n: int) -> list:\n"
            '    """Return a list 1..n: "Fizz" div by 3, "Buzz" div by 5,\n'
            '    "FizzBuzz" both, else the number as a string."""'
        ),
        tests=textwrap.dedent("""\
            from solution import fizzbuzz
            def test_basic():
                r = fizzbuzz(15)
                assert r[2]  == "Fizz"
                assert r[4]  == "Buzz"
                assert r[14] == "FizzBuzz"
                assert r[0]  == "1"
            def test_length(): assert len(fizzbuzz(5)) == 5
        """),
        solution=(
            'def fizzbuzz(n: int) -> list:\n'
            '    out = []\n'
            '    for i in range(1, n + 1):\n'
            '        if i % 15 == 0: out.append("FizzBuzz")\n'
            '        elif i % 3 == 0: out.append("Fizz")\n'
            '        elif i % 5 == 0: out.append("Buzz")\n'
            '        else: out.append(str(i))\n'
            '    return out'
        ),
    ),
    dict(
        name="count_vowels",
        prompt=(
            "Write a Python function:\n\n"
            "def count_vowels(s: str) -> int:\n"
            '    """Return the number of vowels (a,e,i,o,u, case-insensitive) in s."""'
        ),
        tests=textwrap.dedent("""\
            from solution import count_vowels
            def test_basic():  assert count_vowels("hello") == 2
            def test_upper():  assert count_vowels("AEIOU") == 5
            def test_none():   assert count_vowels("bcdf") == 0
            def test_empty():  assert count_vowels("") == 0
        """),
        solution=(
            "def count_vowels(s: str) -> int:\n"
            "    return sum(1 for c in s.lower() if c in 'aeiou')"
        ),
    ),
    dict(
        name="flatten",
        prompt=(
            "Write a Python function:\n\n"
            "def flatten(lst: list) -> list:\n"
            '    """Flatten one level of nesting: [[1,2],[3]] -> [1,2,3]."""'
        ),
        tests=textwrap.dedent("""\
            from solution import flatten
            def test_basic():  assert flatten([[1,2],[3,4]]) == [1,2,3,4]
            def test_empty():  assert flatten([]) == []
            def test_single(): assert flatten([[1]]) == [1]
            def test_mixed():  assert flatten([[1,2],[]]) == [1,2]
        """),
        solution=(
            "def flatten(lst: list) -> list:\n"
            "    return [x for sub in lst for x in sub]"
        ),
    ),
    dict(
        name="max_consecutive",
        prompt=(
            "Write a Python function:\n\n"
            "def max_consecutive(nums: list) -> int:\n"
            '    """Return the length of the longest run of equal consecutive elements."""'
        ),
        tests=textwrap.dedent("""\
            from solution import max_consecutive
            def test_basic():  assert max_consecutive([1,1,2,2,2,3]) == 3
            def test_single(): assert max_consecutive([5]) == 1
            def test_empty():  assert max_consecutive([]) == 0
            def test_all():    assert max_consecutive([7,7,7]) == 3
        """),
        solution=(
            "def max_consecutive(nums: list) -> int:\n"
            "    if not nums: return 0\n"
            "    best = cur = 1\n"
            "    for a, b in zip(nums, nums[1:]):\n"
            "        cur = cur + 1 if a == b else 1\n"
            "        best = max(best, cur)\n"
            "    return best"
        ),
    ),
    dict(
        name="second_largest",
        prompt=(
            "Write a Python function:\n\n"
            "def second_largest(nums: list) -> int | None:\n"
            '    """Return the second largest unique value, or None if fewer than 2 unique values."""'
        ),
        tests=textwrap.dedent("""\
            from solution import second_largest
            def test_basic():  assert second_largest([3,1,4,1,5]) == 4
            def test_two():    assert second_largest([2,1]) == 1
            def test_dupes():  assert second_largest([1,1,1]) is None
            def test_empty():  assert second_largest([]) is None
        """),
        solution=(
            "def second_largest(nums: list):\n"
            "    u = sorted(set(nums), reverse=True)\n"
            "    return u[1] if len(u) >= 2 else None"
        ),
    ),
    dict(
        name="run_length_encode",
        prompt=(
            "Write a Python function:\n\n"
            "def run_length_encode(s: str) -> str:\n"
            '    """Run-length encode s: "aaabbc" -> "a3b2c1"."""'
        ),
        tests=textwrap.dedent("""\
            from solution import run_length_encode
            def test_basic():  assert run_length_encode("aaabbc") == "a3b2c1"
            def test_single(): assert run_length_encode("a") == "a1"
            def test_empty():  assert run_length_encode("") == ""
            def test_mixed():  assert run_length_encode("abcd") == "a1b1c1d1"
        """),
        solution=(
            "def run_length_encode(s: str) -> str:\n"
            "    if not s: return ''\n"
            "    out, cur, n = [], s[0], 1\n"
            "    for c in s[1:]:\n"
            "        if c == cur: n += 1\n"
            "        else: out.append(f'{cur}{n}'); cur, n = c, 1\n"
            "    out.append(f'{cur}{n}')\n"
            "    return ''.join(out)"
        ),
    ),
    dict(
        name="rotate_list",
        prompt=(
            "Write a Python function:\n\n"
            "def rotate_list(lst: list, k: int) -> list:\n"
            '    """Return lst rotated right by k positions."""'
        ),
        tests=textwrap.dedent("""\
            from solution import rotate_list
            def test_basic():  assert rotate_list([1,2,3,4,5], 2) == [4,5,1,2,3]
            def test_zero():   assert rotate_list([1,2,3], 0) == [1,2,3]
            def test_empty():  assert rotate_list([], 3) == []
            def test_full():   assert rotate_list([1,2,3], 3) == [1,2,3]
        """),
        solution=(
            "def rotate_list(lst: list, k: int) -> list:\n"
            "    if not lst: return []\n"
            "    k = k % len(lst)\n"
            "    return lst[-k:] + lst[:-k] if k else lst[:]"
        ),
    ),
    dict(
        name="word_frequency",
        prompt=(
            "Write a Python function:\n\n"
            "def word_frequency(text: str) -> dict:\n"
            '    """Return word -> count (case-insensitive, split on whitespace)."""'
        ),
        tests=textwrap.dedent("""\
            from solution import word_frequency
            def test_basic():  assert word_frequency("the cat sat") == {"the":1,"cat":1,"sat":1}
            def test_repeat(): assert word_frequency("a a b") == {"a":2,"b":1}
            def test_case():   assert word_frequency("A a") == {"a":2}
            def test_empty():  assert word_frequency("") == {}
        """),
        solution=(
            "def word_frequency(text: str) -> dict:\n"
            "    d = {}\n"
            "    for w in text.lower().split():\n"
            "        d[w] = d.get(w, 0) + 1\n"
            "    return d"
        ),
    ),
]

SYSTEM_PROMPT = (
    "You are a Python coding assistant. "
    "Write ONLY the function: no imports, no test code, no explanation. "
    "Output raw Python starting with `def`."
)


# ─── Dataset helpers ──────────────────────────────────────────────────────────

def build_dataset(tasks: list) -> Dataset:
    return Dataset.from_dict({
        "prompt": [
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": t["prompt"]},
            ]
            for t in tasks
        ],
        "tests": [t["tests"] for t in tasks],
    })


def _extract_code(text) -> str:
    if isinstance(text, list):
        text = text[0]["content"] if text else ""
    text = text or ""
    m = re.search(r"```(?:python)?\s*(.*?)```", text, re.DOTALL)
    return (m.group(1) if m else text).strip()


# ─── Sandbox reward ───────────────────────────────────────────────────────────

_HARNESS = """\
import sys, os, subprocess, re
sys.path.insert(0, "/tmp/pkgs")
if not os.path.isdir("/tmp/pkgs"):
    subprocess.run(
        ["python3", "-m", "pip", "install", "pytest", "-q", "--target", "/tmp/pkgs"],
        capture_output=True, check=False,
    )
    sys.path.insert(0, "/tmp/pkgs")
os.makedirs("/tmp/sol", exist_ok=True)
open("/tmp/sol/solution.py",  "w").write({code!r})
open("/tmp/sol/test_sol.py",  "w").write({tests!r})
r = subprocess.run(
    ["python3", "-m", "pytest", "/tmp/sol/test_sol.py", "--tb=no", "-q",
     "--import-mode=importlib"],
    capture_output=True, text=True,
    env={{**os.environ, "PYTHONPATH": "/tmp/pkgs:/tmp/sol"}},
)
p = int((re.search(r"(\\d+) passed", r.stdout) or [0,0])[1])
f = int((re.search(r"(\\d+) failed", r.stdout) or [0,0])[1])
t = p + f
print(f"{{p}}/{{t}}")
"""

def _run_sandbox(code: str, tests: str) -> float:
    harness = _HARNESS.format(code=code, tests=tests)
    try:
        box = Sandbox.create(memory_mb=2048)
            ex = box.run("python3", ["-c", harness])
            last = (ex.stdout or "").strip().splitlines()
            last = last[-1] if last else "0/0"
        p, t = (int(x) for x in last.split("/"))
        return p / t if t > 0 else 0.0
    except Exception:
        return 0.0


# ─── Reward function: logs best completion of every batch ───────────────────

_reward_log: List[dict] = []   # accumulates {code, score, step} across training
_step       = [0]              # mutable counter (closure-friendly)

def reward_sandbox(completions, tests: List[str], **kwargs) -> List[float]:
    """
    Reward = fraction of hidden pytest tests that pass (0.0–1.0).
    G completions are dispatched to G parallel sandboxes.
    Every batch whose best score > 0 is printed immediately.
    """
    codes = [_extract_code(c) for c in completions]
    _step[0] += 1

    with ThreadPoolExecutor(max_workers=len(codes)) as pool:
        futures = {pool.submit(_run_sandbox, code, test): i
                   for i, (code, test) in enumerate(zip(codes, tests))}
        scores = [0.0] * len(codes)
        for fut in as_completed(futures):
            i = futures[fut]
            scores[i] = fut.result()
            _reward_log.append({"step": _step[0], "code": codes[i], "score": scores[i]})

    best_i = max(range(len(scores)), key=lambda i: scores[i])
    if scores[best_i] > 0:
        console.print(
            f"\n  [bold green]↑ step {_step[0]}  reward={scores[best_i]:.0%}"
            f"  ({int(scores[best_i]*4)}/4 tests)[/bold green]"
        )
        console.print(Panel(
            codes[best_i],
            title=f"[bold green]Best completion - step {_step[0]}[/bold green]",
            border_style="green",
        ))

    return scores


def print_top_completions(n: int = 3):
    nonzero = [e for e in _reward_log if e["score"] > 0]
    if not nonzero:
        console.print("[yellow]No non-zero rewards recorded during training.[/yellow]")
        return
    top = sorted(nonzero, key=lambda e: e["score"], reverse=True)[:n]
    console.print(Rule(f"[bold green]Top {len(top)} completions by reward[/bold green]", style="green"))
    for rank, entry in enumerate(top, 1):
        color = "green" if entry["score"] >= 0.75 else "yellow"
        console.print(Panel(
            entry["code"],
            title=f"[bold]#{rank}  reward={entry['score']:.0%}  step={entry['step']}[/bold]",
            border_style=color,
        ))


# ─── Phase 1: SFT warmup ──────────────────────────────────────────────────────

def sft_warmup(model, tokenizer, tasks: list, steps: int = 30):
    """
    Brief supervised pass on correct solutions.
    Teaches the model to emit valid Python before GSPO takes over.
    Without this, reward_std=0 every step and GSPO has no gradient signal.
    """
    console.print(Rule("[magenta]Phase 1: SFT warmup[/magenta]", style="magenta"))
    console.print(
        f"[dim]{steps} gradient steps on correct solutions "
        f"({len(tasks)} tasks, cycling).  Goal: non-zero reward_std in Phase 2.[/dim]\n"
    )

    optimizer = AdamW(model.parameters(), lr=2e-5)
    model.train()

    texts = []
    for task in tasks:
        messages = [
            {"role": "system",    "content": SYSTEM_PROMPT},
            {"role": "user",      "content": task["prompt"]},
            {"role": "assistant", "content": task["solution"]},
        ]
        texts.append(tokenizer.apply_chat_template(messages, tokenize=False))

    for step in range(1, steps + 1):
        text   = texts[(step - 1) % len(texts)]
        enc    = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        labels = enc["input_ids"].clone()

        outputs = model(**enc, labels=labels)
        outputs.loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if step % max(1, steps // 5) == 0 or step == steps:
            console.print(f"  SFT step {step:3d}/{steps}  loss={outputs.loss.item():.4f}")

    del optimizer
    console.print("[dim]SFT warmup done.\n[/dim]")


# ─── Evaluation ───────────────────────────────────────────────────────────────

def evaluate(model, tokenizer, tasks: list):
    model.eval()
    device = next(model.parameters()).device

    t = Table(box=box.SIMPLE, show_header=True, header_style="bold white")
    t.add_column("Function",     width=20)
    t.add_column("Tests",        width=7,  justify="right")
    t.add_column("Generated code (first 55 chars)", width=57)
    t.add_column("",             width=5)

    total = 0.0
    for task in tasks:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": task["prompt"]},
        ]
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        enc  = tokenizer(text, return_tensors="pt")
        input_ids      = enc["input_ids"].to(device)
        attention_mask = enc["attention_mask"].to(device)
        with torch.no_grad():
            out = model.generate(
                input_ids, attention_mask=attention_mask,
                max_new_tokens=160, do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        response = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
        code  = _extract_code(response)
        score = _run_sandbox(code, task["tests"])
        total += score

        bar   = "█" * int(score * 5) + "░" * (5 - int(score * 5))
        color = "green" if score == 1.0 else "yellow" if score > 0 else "red"
        t.add_row(
            task["name"],
            f"[{color}]{score:.0%}[/{color}]",
            code.replace("\n", "↵ ")[:55],
            f"[{color}]{bar}[/{color}]",
        )

    console.print(t)
    avg = total / len(tasks)
    console.print(f"  Average test pass rate: [bold cyan]{avg:.1%}[/bold cyan]")
    return avg


# ─── Main ─────────────────────────────────────────────────────────────────────

def train_gspo():
    tasks       = TASKS[:3] if SMOKE else TASKS
    sft_steps   = 20        if SMOKE else 60
    gspo_epochs = 1         if SMOKE else 3

    console.print(Panel(
        "[bold green]RL GSPO: Code Generation with Hidden Test Suites[/bold green]\n\n"
        "[dim]Algorithm   : GSPO (sequence-level IS) - GRPOConfig(importance_sampling_level='sequence')\n"
        "Model       : SmolLM2-135M-Instruct  (135 M params, CPU-friendly)\n"
        "Task        : Implement Python functions from docstrings\n"
        "Reward      : fraction of hidden pytest tests passing (sandbox oracle)\n"
        "Sandboxes   : G parallel TensorLake sandboxes per GSPO step\n"
        "Phase 1     : SFT warmup - correct solutions so model starts generating valid Python\n"
        "Phase 2     : GSPO - refines via reward signal from sandbox test results\n"
        "GPU needed  : No\n"
        f"Mode        : {'SMOKE (3 tasks, 20 SFT steps, 1 GSPO epoch)' if SMOKE else f'Full ({len(tasks)} tasks, {sft_steps} SFT steps, {gspo_epochs} GSPO epochs)'}[/dim]",
        border_style="green",
    ))

    console.print("\n[dim]Loading SmolLM2-135M-Instruct...[/dim]")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.float32)

    split      = max(1, int(0.75 * len(tasks)))
    train_tasks = tasks[:split]
    eval_tasks  = tasks[split:]
    console.print(f"[dim]{split} train tasks / {len(eval_tasks)} eval tasks[/dim]\n")

    # ── Baseline ────────────────────────────────────────────────────────────
    console.print(Rule("[cyan]Baseline: before any training[/cyan]", style="cyan"))
    evaluate(model, tokenizer, eval_tasks)

    # ── Phase 1: SFT warmup ─────────────────────────────────────────────────
    sft_warmup(model, tokenizer, train_tasks, steps=sft_steps)

    console.print(Rule("[cyan]After SFT warmup[/cyan]", style="cyan"))
    evaluate(model, tokenizer, eval_tasks)

    # ── Phase 2: GSPO ────────────────────────────────────────────────────────
    console.print(Rule("[yellow]Phase 2: GSPO fine-tuning[/yellow]", style="yellow"))
    console.print(
        "[dim]Best completions printed live as reward > 0 is observed.\n"
        "reward_std > 0 confirms the policy is exploring.[/dim]\n"
    )

    config = GRPOConfig(
        output_dir=OUTPUT_DIR,
        importance_sampling_level="sequence",   # ← GSPO vs GRPO
        num_generations=2 if SMOKE else 4,
        max_completion_length=200,
        temperature=1.4,   # high temp forces diverse G completions → reward_std > 0
        learning_rate=2e-6,
        num_train_epochs=gspo_epochs,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=2 if SMOKE else 4,
        warmup_steps=5,
        beta=0.001,
        epsilon=0.2,
        logging_steps=1,
        save_steps=999,
        seed=42,
        report_to="none",
        bf16=False,
        fp16=False,
    )

    trainer = GRPOTrainer(
        model=model,
        args=config,
        train_dataset=build_dataset(train_tasks),
        reward_funcs=[reward_sandbox],
        processing_class=tokenizer,
    )

    trainer.train()

    # ── Results ──────────────────────────────────────────────────────────────
    print_top_completions(n=3)

    console.print(Rule("[cyan]After GSPO training[/cyan]", style="cyan"))
    final_acc = evaluate(model, tokenizer, eval_tasks)

    console.print(Panel(
        f"[bold]Result: {final_acc:.0%} average test pass rate on held-out functions[/bold]\n\n"
        "Context:\n"
        "  • Eval functions were [bold]never seen[/bold] during SFT or GSPO training\n"
        "  • Baseline before any training: [red]0%[/red]\n"
        f"  • After GSPO: [bold green]{final_acc:.0%}[/bold green]"
        "  ← model generalised from 7 training functions to unseen ones\n\n"
        "Why 25 % is a reasonable outcome for this setup:\n"
        "  • 135 M params is the [italic]smallest[/italic] publicly available instruct model\n"
        "  • Only 60 SFT steps on 7 reference solutions (~5 min CPU)\n"
        "  • 25 % means 1 / 4 tests pass per function: the model correctly\n"
        "    handles the empty-input edge case on all three unseen functions,\n"
        "    showing the pattern [italic]transferred[/italic] across task types\n"
        "  • Typical zero-shot pass@1 for 135 M models on HumanEval is < 5 %\n\n"
        "Cheap ways to push higher (no extra hardware):\n"
        "  1. [cyan]temperature=1.4[/cyan] (already set) - forces reward_std > 0 so GSPO\n"
        "     has a gradient signal instead of collapsing to all-zero advantages\n"
        "  2. More SFT examples (50+ functions, ~10 min) before GSPO\n"
        "  3. Switch to [cyan]Qwen2.5-0.5B-Instruct[/cyan] (4× more params, same CPU time)",
        title="[bold cyan]Score interpretation[/bold cyan]",
        border_style="cyan",
    ))

    console.print(Panel(
        "[bold]Why GSPO + sandbox here?[/bold]\n\n"
        "1. [cyan]Sandboxes required[/cyan]: model code is untrusted and cannot run in-process.\n\n"
        "2. [cyan]Hidden test suites[/cyan]: the model never sees the tests.\n"
        "   Sandbox is the only oracle → no reward hacking.\n\n"
        "3. [cyan]GSPO over GRPO[/cyan]: long function bodies mean many tokens.\n"
        "   Token-level IS clipping (GRPO) lets noisy tokens dominate the gradient.\n"
        "   Sequence-level clipping (GSPO) clips the whole trajectory once:\n\n"
        "   GRPO: clip( π_θ(t)/π_old(t) ) per token\n"
        "   GSPO: clip( Π_t π_θ(t)/π_old(t) ) once per sequence",
        title="[bold cyan]Design rationale[/bold cyan]",
        border_style="cyan",
    ))


if __name__ == "__main__":
    train_gspo()

What happens step-by-step

Each run has 7 steps, from loading the model to the final held-out evaluation:

Step	Phase	What happens
1	Setup	The script loads the model and tokenizer, then splits the tasks 75/25 into train and eval sets.
2	Baseline	The untrained model attempts the eval tasks, and sandboxes score the output. Expect 0%.
3	SFT warmup	N supervised gradient steps on correct reference solutions teach the model to produce parseable Python before RL begins.
4	After SFT	The eval runs again. Reward variance should now be non-zero, which GSPO needs for a gradient signal.
5	GSPO loop	Each training step generates G completions and dispatches them to G parallel sandboxes. Each sandbox runs the hidden pytest suite and returns a score.
6	Reward signal	`reward_sandbox` collects the scores, logs the best completion, and returns the score list to `GRPOTrainer`.
7	Final eval	The model attempts the held-out functions it never saw during training. Expect a 25% pass rate from a 135M-parameter model.

Key design decisions

Three choices keep the reward signal non-zero and hard to game: a high sampling temperature, an SFT warmup, and hidden tests.

Why `temperature=1.4`

GSPO needs diversity across the G completions in each group to produce a non-zero reward standard deviation. At low temperature all completions are identical, so reward_std is 0. Advantage normalization then produces zero gradients and training stalls. A high temperature forces the model to explore different implementations.

Why SFT warmup is required

Without warmup, a randomly initialized or instruction-tuned model produces malformed Python that scores 0 on every test case. All-zero rewards mean all-zero advantages after normalization, and GSPO has nothing to optimize. Even 20 supervised steps on correct solutions bootstrap non-zero reward variance.

Why sandboxes prevent reward hacking

The model never has access to the test file. Its only feedback is the pass rate the sandbox returns. The model can’t overfit to specific assertion patterns, so it must implement the correct logic.

Overview

Guides

Use Cases

Fine-Tune a Code Model with GSPO

How it works

GSPO vs GRPO

Prerequisites

TypeScript SDK starter

Full example

What happens step-by-step

Key design decisions

Why `temperature=1.4`

Why SFT warmup is required

Why sandboxes prevent reward hacking

What to build next

AI Code Execution

Agentic Swarm Intelligence

​How it works

​GSPO vs GRPO

​Prerequisites

​TypeScript SDK starter

​Full example

​What happens step-by-step

​Key design decisions

​Why temperature=1.4

​Why SFT warmup is required

​Why sandboxes prevent reward hacking

​What to build next

AI Code Execution

Agentic Swarm Intelligence

How it works

GSPO vs GRPO

Prerequisites

TypeScript SDK starter

Full example

What happens step-by-step

Key design decisions

Why `temperature=1.4`

Why SFT warmup is required

Why sandboxes prevent reward hacking

What to build next