How to build an RL environment for an LLM agent.

Define the capability, not the benchmark

Start from the behavior the deployed agent must produce, then design tasks that exercise it. The reverse order, adopting a public benchmark and training against it, produces models that score well and transfer badly, because public benchmarks encode someone else's task distribution and are often contaminated.

Write the capability as a set of task families with explicit success criteria. For a support agent: resolve a billing dispute given account tools, escalate correctly when policy requires it, recover after a tool returns an error. Each family becomes a generator of task instances, not a fixed list.

Specify the task distribution

A fixed task list is memorized quickly. What generalizes is a distribution: procedural generators with controlled parameters (entities, quantities, orderings, distractors), difficulty knobs for curricula, and adversarial variants seeded from observed policy failures.

• Parameterize every task family: what varies, over what ranges, with what constraints
• Reserve entire parameter regions and scenario types for held-out evaluation, not just random instances
• Balance the mixture so no family dominates the gradient
• Keep a hard-negative pool that grows as the policy exposes weaknesses

For group-based methods such as GRPO, difficulty placement is signal engineering: tasks the policy always fails or always solves produce zero group variance and therefore zero learning signal.

Define state and what the agent can observe

Separate world state from observation. World state is everything the environment tracks: databases, files, ticket status, hidden faults. The observation is the slice the agent sees, rendered into its context. The gap between them is partial observability, and it is usually the point: the deployed agent will also not see everything.

Make the observation renderer an explicit, versioned function of state. Silent changes to how state is rendered into text change the training distribution as surely as changing the tasks.

Define actions and tool contracts

The action space of a language-model agent is text, constrained into structure: tool calls with typed arguments, plus free-form response turns. Specify each tool as a contract: schema, semantics, error behavior, latency, and side effects on world state.

• Validate tool arguments and return realistic errors for invalid calls, rather than crashing the episode
• Reproduce the failure modes of the real tools: timeouts, permission denials, empty results, stale data
• Decide explicitly whether malformed model output ends the episode, costs reward, or returns a parse error the agent can recover from

That last decision matters more than it looks. Recoverable errors teach recovery, which is usually a target capability in itself.

Implement transitions and episode boundaries

The transition function applies an action to world state and returns the next observation. Keep it pure with respect to the environment's random seed: given the same seed and action sequence, the same states must follow. Stochastic elements (flaky tools, background events) draw from the seeded stream so they replay identically.

Define termination as data, not convention: success conditions, hard failure conditions, and a step or token budget as truncation. Distinguish terminated (the task ended) from truncated (the budget ended it), because credit assignment treats them differently and conflating them corrupts reward.

Build reset and replay early

Training will run the environment millions of times. Reset must restore a task's initial state quickly and completely, with no leakage from prior episodes: no leftover files, rows, or cache entries. Container snapshots, database templates, and copy-on-write filesystems are the usual mechanisms.

Replay, reconstructing an episode exactly from its seed and action log, is the debugging tool you will use most. If a trajectory shows a bizarre reward, replay either reproduces it (environment bug) or does not (nondeterminism bug). Both are things to find before training, which is why the reference interface below treats the seed as part of the API.

Build verifiers and rewards

The verifier checks the outcome against ground truth the environment knows: final state matches, tests pass, required steps occurred. The reward maps verifier output and episode costs into the scalar the optimizer maximizes. Keep the two separate: verifiers are reusable facts, reward shaping is training policy.

• Prefer checks on world state over checks on the agent's text: what happened, not what was claimed
• Decompose reward into named components (task success, invalid-call penalty, step cost) and log each component separately
• Red-team the verifier before training: try to construct a passing failure by hand, because the policy will search for exactly that gap

Choosing between programmatic verifiers, learned reward models, and rubric judges is its own decision, covered on the rewards and verifiers page.

Log trajectories as first-class data

Every step of every episode should produce a structured record. Trajectories are simultaneously the training data pipeline, the debugging record, and the source for later supervised distillation, so design the schema before the first rollout rather than reverse-engineering it from logs.

{
  "episode_id": "ep_00142",
  "task_id": "billing-dispute/hard/seed-771",
  "env_version": "1.4.2",
  "policy_checkpoint": "step_01800",
  "step": 7,
  "observation": { "rendered": "...", "tokens": 412 },
  "action": { "type": "tool_call", "tool": "refund_lookup", "args": {} },
  "tool_result": { "status": "ok", "latency_ms": 84 },
  "reward_components": { "task": 0.0, "invalid_call": 0.0, "step_cost": -0.01 },
  "terminated": false,
  "truncated": false
}

Version fields are not optional. When results shift, the first question is always which environment version and which checkpoint produced the trajectory.

A minimal reference interface

The classic reset-and-step shape from Gymnasium transfers well to agent environments, with two additions that agent training keeps needing: a snapshot for replay and state inspection, and explicit truncation.

interface AgentEnvironment<State, Observation, Action> {
  /** Restore initial state for a task. Same seed, same episode. */
  reset(taskId: string, seed: number): Promise<Observation>;

  /** Apply one action and advance world state. */
  step(action: Action): Promise<{
    observation: Observation;
    reward: number;
    rewardComponents: Record<string, number>;
    terminated: boolean;
    truncated: boolean;
    info: Record<string, unknown>;
  }>;

  /** Full world state for replay, debugging, and verifiers. */
  snapshot(): Promise<State>;
}

Nothing requires Gymnasium itself, and agent environments often outgrow its assumptions. What matters is that the contract above is honored exactly, because the rollout system is written against it.

Parallelize rollouts

Online training throughput is set by rollout throughput. Run many environment instances behind a queue, batch the policy's inference across them, and make workers stateless between episodes so instances can be killed and respawned freely. Track weight version per trajectory: samples generated by an outdated policy are a real and quiet source of degraded training.

Budget environment compute honestly. In tool-heavy environments the sandbox, not the model, is often the bottleneck, and a slow reset dominates everything at scale.

Build held-out tasks before training

Hold out at the level that predicts deployment: unseen parameter regions, unseen scenario types, and where possible an evaluation-only task family the training distribution never touches. Evaluate on a fixed cadence during training and select checkpoints on held-out results, never on training reward, which measures the policy's ability to exploit the training distribution.

Environment QA checklist

Run this before the first training job. Every item that fails here costs far more to discover mid-run.

1. Replay determinism: same seed and action log reproduce identical states, observations, and rewards.
2. Reset isolation: back-to-back episodes share no state, verified by fingerprinting world state after reset.
3. Scripted baselines: a known-good action sequence scores high, a known-bad one scores low, a do-nothing agent scores near zero.
4. Verifier red-team: hand-constructed passing failures were attempted and blocked.
5. Error paths: invalid tool calls, timeouts, and malformed model output all produce defined outcomes, not crashes.
6. Truncation accounting: terminated and truncated are never conflated in reward or logs.
7. Distribution audit: task-family mix, difficulty spread, and current-policy success rates are measured, with success neither near 0% nor near 100% on training tasks.
8. Throughput and reset latency measured under realistic parallelism.
9. Trajectory records validate against the schema and carry environment and checkpoint versions.
10. Held-out families exist and are excluded from every training and curriculum path.