Building an RL Environment for an LLM Coding Agent

You already know supervised learning: a fixed dataset of (x, y) pairs, minimize a loss against known targets. Reinforcement learning drops the targets. Nobody tells you the right answer — you only get a scalar reward after acting, and you have to figure out which actions led to it. That single change, supervision by evaluation instead of by demonstration, is the whole story, and it's exactly why RL is the natural tool for "make this LLM agent better at a task" once you can score the outcome.

I wanted to actually understand this rather than nod along to it, so I built one from scratch: a real RL environment for an LLM coding agent, with a safe sandbox, a reward function, a real task distribution, and a real gpt-4o policy making real rollouts. This post walks through what I built, in the order I built it. It stops right before the part where those rewards turn into a policy update — that's a different post.

RL in twenty minutes#

Before touching LLMs, it's worth grounding reset/step on something you can hold in your head completely: a number-guessing game. The agent guesses; the environment says higher/lower and rewards a hit. There are no labels anywhere in this loop — only rewards and observations.

class NumberGuessEnv:
    """Guess the secret number in [1, 100]. Reward 1.0 on a hit, else 0.0."""
 
    def __init__(self, seed=None):
        self._rng = random.Random(seed)
 
    def reset(self):
        self._secret = self._rng.randint(1, 100)
        self._low, self._high = 1, 100
        return {"low": self._low, "high": self._high}   # observation
 
    def step(self, guess):
        if guess == self._secret:
            return {}, 1.0, True, {"secret": self._secret}
        if guess < self._secret:
            self._low = max(self._low, guess + 1)
        else:
            self._high = min(self._high, guess - 1)
        obs = {"low": self._low, "high": self._high}
        return obs, 0.0, False, {}

solved in 9 guesses, secret was 50, return=1.0

Swap the binary-search policy for a random guess and rerun — a worse policy needs more steps for the same reward. That gap is exactly what an RL algorithm climbs. We won't run any gradient updates in this post, but keep the mental image: the environment defines the game; the algorithm climbs it.

View notebook → 00_rl_overview.ipynb

The environment skeleton#

The coding environment uses the exact same shape as the number-guessing toy — that's the point. An RL environment is a contract, not a particular domain:

obs = env.reset()                            # pose a task; obs is what the policy conditions on
obs, reward, done, info = env.step(action)   # score a submission

reset() samples a task from the task distribution and returns an observation — here, a dict with the problem prompt. The hidden unit tests are not in the observation. step() takes the policy's raw output, runs it, and returns the reward plus an info dict of diagnostics.

Here's the actual class — reset, _reward, and step, in full:

class CodingEnv:
    """Single-step coding environment: reset() poses a task, step(code) runs the
    hidden unit tests and returns a reward in [0, 1]. No state transition — a
    contextual bandit, the same shape most production LLM RL on verifiable
    rewards (e.g. GRPO) is built on."""
 
    def __init__(self, tasks=TASKS, timeout=5.0, seed=None):
        self.tasks = tasks
        self.timeout = timeout
        self._rng = random.Random(seed)
        self.current_task = None
 
    def reset(self, task_index=None):
        self.current_task = (
            self._rng.choice(self.tasks) if task_index is None else self.tasks[task_index]
        )
        return {"prompt": self.current_task.prompt,
                "entry_point": self.current_task.entry_point,
                "task_name": self.current_task.name}
 
    def _reward(self, result, code):
        if not code.strip():      # format gate: no code at all -> 0
            return 0.0
        return result.fraction    # dense: tests passed / tests total
 
    def step(self, action):
        code = extract_code(action)
        result = run_tests(code, self.current_task.checks, timeout=self.timeout)
        reward = self._reward(result, code)
        info = {"task_name": self.current_task.name, "passed": result.passed,
                "total": result.total, "timed_out": result.timed_out}
        return {}, reward, True, info   # done is always True — single-step

Here's a full episode against a scripted mock policy, on the easiest task:

obs = env.reset(task_index=0)        # task 0 is `add`
action = mock_policy(obs)            # the 'policy': observation -> code
obs2, reward, done, info = env.step(action)

And the same mock policy over the whole task set — it only knows how to solve add, so the rest score zero:

View notebook → 01_env_skeleton.ipynb

The sandbox#

step() has to execute code an LLM wrote and decide if it's correct. That's the heart of a coding RL environment, and the part most likely to hurt you if done naively. An LLM will eventually emit an infinite loop, a crash, or something actively hostile. The environment has to turn any of that into a clean reward without taking down your process.

def run_tests(code_str, checks, timeout=2.0):
    """Write candidate code + checks to a temp file, run it in a subprocess with
    a hard timeout, and report which checks passed without ever hanging the
    calling process."""
    # ... writes solution.py + a runner script to a temp dir ...
    # ... subprocess.run(..., timeout=timeout) ...
    # ... parses per-check pass/fail out of the subprocess's stdout ...
    return RunResult(passed=..., total=..., fraction=..., timed_out=..., failures=...)

checks = ["assert add(2, 3) == 5", "assert add(-1, 1) == 0"]
cases = {
    "correct":        "def add(a, b):\n    return a + b\n",
    "wrong":          "def add(a, b):\n    return a - b\n",
    "crashes":        "def add(a, b):\n    raise ValueError('boom')\n",
    "syntax error":   "def add(a, b)\n    return a + b\n",
    "infinite loop":  "def add(a, b):\n    while True: pass\n",
}
for label, code_str in cases.items():
    r = run_tests(code_str, checks, timeout=2.0)

Every pathological case becomes a finite result with fraction == 0.0 — the kernel never hangs or dies. The timeout is the single most important line in the file: without it, one bad generation freezes the whole run.

View notebook → 02_sandbox.ipynb

The reward function#

The reward is where you encode what you actually want. Get it subtly wrong and the policy will exploit exactly that gap — reward hacking isn't a hypothetical, it's the default outcome of a misspecified reward. The design here:

def _reward(self, result, code):
    if not code.strip():      # format gate: no code at all -> 0, no partial credit
        return 0.0
    return result.fraction    # dense: (tests passed) / (tests total)

Run it against four attempts at is_palindrome, from nothing to fully correct:

The reward climbs smoothly instead of sitting at zero until the last attempt. That matters: a purely sparse reward (1.0 only if all tests pass, else 0.0) gives the optimizer almost nothing to grip early — most attempts score 0 and look identical. Dense shaping says "you're getting warmer." Re-scoring the same four attempts under a sparse reward collapses three of them to indistinguishable zeros:

And the hackability risk is real, not theoretical — a solution that special-cases the one example in the prompt still earns partial credit, because the hidden tests still catch most of it:

View notebook → 03_reward.ipynb

The task distribution#

An environment is only as good as the problems it can pose. reset() samples from a task distribution, and the breadth and difficulty of that distribution determines what a policy can possibly learn:

That difficulty spread is deliberate, not an afterthought: a task every policy passes, or every policy fails, carries no learning signal at all — there's nothing for an algorithm to distinguish between rollouts on. Difficulty spread is itself a design parameter of the environment.

The punchline is that a standard benchmark is the same object in a different costume. Here's HumanEval's first problem, hand-transcribed into this project's Task schema and run through the exact same reset/step loop:

solution = ("```python\n"
            "def has_close_elements(numbers, threshold):\n"
            "    for i in range(len(numbers)):\n"
            "        for j in range(i+1, len(numbers)):\n"
            "            if abs(numbers[i]-numbers[j]) < threshold:\n"
            "                return True\n"
            "    return False\n```")

Production eval harnesses load HumanEval/MBPP/LiveCodeBench exactly this way: map each problem's prompt, entry point, and tests into the same {prompt, entry_point, checks} contract. Nothing about the environment changes between a handmade task and a real benchmark.

View notebook → 04_tasks.ipynb

A real policy: gpt-4o#

Everything so far used a scripted mock so the environment mechanics stayed in focus. Now the policy becomes a real LLM — and nothing about the environment changes. That's the payoff of a clean reset/step contract: the policy is pluggable.

raw = openai_policy(obs, model="gpt-4o", temperature=0.2)
_, reward, _, info = env.step(raw)

And a full rollout over every task in the set:

gpt-4o aces all four, including the "hard" parse_roman task — a good reminder that the task distribution would need genuinely harder problems (or a weaker model) before reward variance shows up. That variance, when it does appear, isn't noise to be eliminated; it's the raw material a policy-gradient method runs on. We're not going there in this post.

View notebook → 05_real_policy.ipynb

Multi-turn agents#

Until now, every episode was one step — a bandit. Real coding agents don't work that way: they run the tests, read the failure, and revise. That feedback loop introduces genuine state transitions, and with them the thing that makes RL RL: sequential credit assignment — which earlier action deserves credit for the eventual success?

MultiTurnCodingEnv keeps everything from the bandit version, except done is only True once the tests pass or max_turns is hit, and each step() feeds back which hidden tests failed so the model can act on real information:

env = MultiTurnCodingEnv(max_turns=3, seed=0)
obs = env.reset(task_index=idx)
done = False
while not done:
    raw = openai_policy(obs, model="gpt-4o", temperature=0.4)
    obs, reward, done, info = env.step(raw)

Comparing best-of-one against the multi-turn final reward on the same task, across a handful of runs:

Both are 1.0 here for the same reason as above — gpt-4o is strong enough on this task set that revision doesn't have much to fix. The mechanism is what matters: in a multi-turn rollout, the reward arrives at the end, but several actions (turn 1's attempt, turn 2's revision) all contributed to it. Assigning credit back across those turns is the temporal credit assignment problem — the reason discounting, value functions, and advantage estimation exist in full RL. In a single-step bandit, it simply doesn't arise. This is the concrete boundary between a bandit and a sequential MDP.

View notebook → 06_multiturn.ipynb

Recap#

What you now have is a real, verifiable-reward RL environment for an LLM coding agent: a clean reset/step contract, a sandbox that turns any model output into a finite scored result, a reward that's dense without being trivially gameable, a task distribution that mirrors real benchmarks, and a pluggable policy that goes from a scripted mock to real gpt-4o rollouts to a multi-turn agentic loop — without the environment itself ever changing.

What's deliberately not here is turning that reward signal into an actual policy update — GRPO, advantages, the training loop. That's a real, separate piece of machinery, and it deserves its own writeup rather than a rushed section at the end of this one.