Writing Your Own RL Environment

An RL environment is a class with two methods. This page builds one for a decision an LLM engineer actually faces — which model to call, and whether to retry — runs a random agent against it, then a learner that beats the random agent, and ends with the checklist for when neither works.

The interface, and why it is shaped this way

The convention that has settled across the ecosystem is two methods and a five-value return, and every part of it is there for a reason.

The problem we are modelling

A query arrives. You can answer it with a cheap model or a strong one. A checker tells you whether the answer was acceptable. If it was not, you may try again with either model, up to three attempts, after which the episode ends unsuccessfully. You may also give up immediately and escalate.

• Hidden state: whether this query is hard. The agent never sees it, which makes this a partially observed problem, and a realistic one.
• Observation: the number of attempts so far and what failed last time. Two small integers.
• Actions: call cheap, call strong, escalate.
• Reward: +1 for an accepted answer, minus the cost of each call, and -0.5 for escalating. All in the same units, because otherwise there is no trade-off to learn.

The interesting property is that a failed cheap call is informative: it is weak evidence that the query is hard, and a good policy should use it. Nobody wrote that rule down. It is what the agent has to discover.

The environment, complete

import random

CHEAP_COST, STRONG_COST = 0.02, 0.20   # in the same units as a correct answer
CHEAP_P, STRONG_P = 0.55, 0.90         # success on an easy query
HARD_PENALTY = 0.5                     # multiplier on a hard query
P_HARD = 0.40
MAX_ATTEMPTS = 3

CHEAP, STRONG, ESCALATE = 0, 1, 2


class RoutingEnv:
    """One query per episode. Choose a model, see whether the checker accepted
    the answer, retry or escalate. Difficulty is hidden from the agent."""

    n_actions = 3

    def __init__(self, seed=None):
        self.rng = random.Random(seed)
        self.hard = False
        self.attempts = 0
        self.last_failure = 0          # 0 none, 1 cheap failed, 2 strong failed

    def reset(self, seed=None):
        if seed is not None:
            self.rng.seed(seed)
        self.hard = self.rng.random() < P_HARD
        self.attempts = 0
        self.last_failure = 0
        return self._obs(), {"hard": self.hard}

    def _obs(self):
        return (self.attempts, self.last_failure)

    def step(self, action):
        if action == ESCALATE:
            return self._obs(), -0.5, True, False, {"outcome": "escalated"}

        cost = CHEAP_COST if action == CHEAP else STRONG_COST
        p = CHEAP_P if action == CHEAP else STRONG_P
        if self.hard:
            p *= HARD_PENALTY

        solved = self.rng.random() < p
        self.attempts += 1
        if not solved:
            self.last_failure = 1 if action == CHEAP else 2

        reward = (1.0 if solved else 0.0) - cost
        terminated = solved
        truncated = (not solved) and self.attempts >= MAX_ATTEMPTS
        outcome = "solved" if solved else ("out of attempts" if truncated else "retry")
        return self._obs(), reward, terminated, truncated, {"outcome": outcome}

Three decisions in that class are worth defending. Difficulty is drawn once per episode and never revealed, which is what makes the failure signal meaningful. Cost is subtracted at the moment the call is made rather than at the end, so a policy that burns attempts pays for them even if it eventually succeeds. And running out of attempts sets truncated rather than terminated, because the task did not finish — you simply stopped.

A random agent, and why you run it first

Before any learning, run a policy that ignores the observation entirely. It is the baseline every later number is compared against, and it is the fastest way to find a broken environment.

def run(env, policy, episodes=5000, seed=0):
    rng = random.Random(seed)
    total, solved, calls = 0.0, 0, 0
    for ep in range(episodes):
        obs, _ = env.reset(seed=seed + ep)
        done, success = False, False
        while not done:
            action = policy(obs, rng)
            obs, reward, terminated, truncated, info = env.step(action)
            total += reward
            calls += 1
            success = terminated
            done = terminated or truncated
        solved += int(success)
    return {
        "mean_return": total / episodes,
        "solve_rate": solved / episodes,
        "calls_per_episode": calls / episodes,
    }


def random_policy(obs, rng):
    return rng.randrange(RoutingEnv.n_actions)


def always_cheap(obs, rng):
    return CHEAP


def always_strong(obs, rng):
    return STRONG


for name, pol in [("random", random_policy),
                  ("always cheap", always_cheap),
                  ("always strong", always_strong)]:
    print(name, run(RoutingEnv(), pol))

Run those three before writing a learner. They tell you the range the learner has to beat, and they catch the most common environment bugs immediately: if always-strong does not beat random, something in the reward or the transition is wrong, and no amount of training will reveal which.

1. Random gives the floor. A learner that does not beat it is not learning.
2. Always-cheap and always-strong give the fixed-policy baselines. In a well-posed routing problem the learner should beat both, because it can condition on the failure signal and they cannot.
3. If a fixed policy beats your learner after training, the problem is the learner or the reward scale, not the environment. That distinction is worth an hour of debugging on its own.

A learner that beats it

The observation space is six states and there are three actions, so a table is not merely adequate, it is the correct choice. Reaching for a neural network here would add a hundred lines and several failure modes for no benefit.

def q_learning(env, episodes=200_000, alpha=0.1, gamma=1.0,
               eps_start=0.5, eps_end=0.01, seed=0):
    rng = random.Random(seed)
    Q = {}

    def row(obs):
        return Q.setdefault(obs, [0.0] * RoutingEnv.n_actions)

    for ep in range(episodes):
        eps = eps_end + (eps_start - eps_end) * (1 - ep / episodes)
        obs, _ = env.reset(seed=seed + ep)
        done = False
        while not done:
            r_obs = row(obs)
            if rng.random() < eps:
                a = rng.randrange(RoutingEnv.n_actions)
            else:
                a = max(range(RoutingEnv.n_actions), key=lambda i: r_obs[i])
            nxt, reward, terminated, truncated, _ = env.step(a)
            done = terminated or truncated
            target = reward if terminated else reward + gamma * max(row(nxt))
            r_obs[a] += alpha * (target - r_obs[a])
            obs = nxt
    return Q


Q = q_learning(RoutingEnv())
for obs in sorted(Q):
    values = ", ".join(f"{v:+.3f}" for v in Q[obs])
    best = ["cheap", "strong", "escalate"][max(range(3), key=lambda i: Q[obs][i])]
    print(f"attempts={obs[0]} last_failure={obs[1]}  [{values}]  -> {best}")


def learned_policy(obs, rng):
    r = Q.get(obs)
    return CHEAP if r is None else max(range(3), key=lambda i: r[i])

print("learned", run(RoutingEnv(), learned_policy))

Note the target line: it bootstraps on truncated but not on terminated. In this particular environment the two give the same number, because a truncated episode lands in a state with no successor and therefore a value of zero. Write it correctly anyway. In any environment where the step limit cuts off a continuing task — an agent stopped at 40 tool calls, a robot stopped at 1,000 timesteps — treating truncation as termination tells the learner that the state it was in was worthless, and the error propagates backwards through every state leading to it.

With these numbers, the interesting behaviour is at the second attempt. After a cheap failure the agent has weak evidence the query is hard, and the arithmetic of whether to try cheap again or pay for the strong model is close enough that changing CHEAP_COST from 0.02 to 0.10 flips it. That sensitivity is the point of building the environment: the answer depends on numbers you can measure, and the environment is where you put them.

Making it a Gymnasium environment

Everything above works with no library at all. Wrapping it in the standard base class is worth doing when you want to use an existing algorithm implementation, because those expect the space objects rather than raw tuples.

import numpy as np
import gymnasium as gym
from gymnasium import spaces


class RoutingGymEnv(gym.Env):
    metadata = {"render_modes": []}

    def __init__(self):
        super().__init__()
        self.action_space = spaces.Discrete(RoutingEnv.n_actions)
        self.observation_space = spaces.MultiDiscrete([MAX_ATTEMPTS + 1, 3])
        self.inner = RoutingEnv()

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        obs, info = self.inner.reset(seed=seed)
        return np.array(obs, dtype=np.int64), info

    def step(self, action):
        obs, reward, terminated, truncated, info = self.inner.step(int(action))
        return np.array(obs, dtype=np.int64), float(reward), terminated, truncated, info

Two things to check against the version you install rather than against this page. The library ships a checker utility that validates an environment against the current API and reports mismatches; run it before spending a day debugging an algorithm. And registration — giving the environment a string id so it can be created by name — has a specific call signature that has changed between releases, so take it from the current documentation.

Five mistakes that make an environment untrainable

1. The observation omits something the optimal policy needs. If two situations requiring different actions look identical to the agent, no algorithm can separate them, and the symptom is a policy that plateaus well below what you know is achievable. Test by writing the best policy you can by hand using only the observation. If you cannot, the agent cannot either.
2. Rewards are on incomparable scales. A success worth 1 and a cost worth 400 means the agent will never attempt anything. Put every term in the same units deliberately, then print the distribution of episode returns and check the range is sane.
3. Truncation is treated as termination. Covered above, and worth repeating because it is silent. The learner bootstraps from zero at a state that was not actually worthless, so the values near the step limit are wrong and the error propagates backwards.
4. The environment is not reproducible. If reset(seed=n) does not produce the same episode twice, you cannot tell a code change from noise, and every debugging session becomes a fight with the random number generator. Seed explicitly and never call the global random functions inside an environment.
5. Reward is only available at the very end of a long episode. Legal, and often untrainable. Either shorten the episode, add intermediate reward that is genuinely correlated with the goal, or accept the sample requirement described on credit assignment.

The general rule behind all five: debug the environment with fixed policies before introducing a learner. A random agent, a greedy agent and a hand-written best-guess policy will find most environment bugs in minutes, whereas a learning algorithm will hide them for a week and then fail for reasons you attribute to hyperparameters.