Designing Custom Gymnasium Environments for RL Training

Building a custom Gymnasium environment is where reinforcement learning stops being a tutorial exercise and starts being an engineering problem. Most public benchmarks are clean, well-behaved MDPs. Your real problem — whether it's warehouse routing, battery dispatch, or a robotic manipulation task — is not. I've built environments for supply-chain simulators and for continuous control loops that had to run at 2,000 steps per second across 32 parallel workers without leaking state or silently breaking the Markov property. The difference between an environment that trains and one that wastes a week of GPU time usually comes down to a handful of design decisions made early. This article covers those decisions: how to model the API contract correctly, how to choose spaces and rewards that actually learn, and how to make the whole thing fast, testable, and reproducible.

The Gymnasium API Contract: What Step, Reset, and Spaces Actually Enforce

Gymnasium (the maintained fork of OpenAI Gym) defines a deceptively simple interface. You subclass gym.Env, define observation_space and action_space, and implement reset() and step(action). In my experience, most bugs come from misinterpreting what that contract promises to the training loop.

Since Gymnasium 0.26, step returns five values, not four: observation, reward, terminated, truncated, info. The split between terminated and truncated matters. terminated means the MDP reached a terminal state defined by your task (e.g., robot fell, goal reached, battery depleted). truncated means you stopped the episode for an external reason, typically a time limit. If you conflate them, algorithms like PPO that bootstrap value estimates will learn incorrect returns. If you have spent time with the formalism in Reinforcement Learning Fundamentals: MDPs, Policies and Value Functions, this is the practical consequence of the infinite-horizon vs. episodic return distinction.

reset also has a strict contract: it must return (observation, info) and it must be able to take an optional seed and options dict. The seed handling is not optional if you want reproducibility. Use self.np_random created via super().reset(seed=seed) instead of the global numpy.random or Python's random module.

A Minimal, Correct Skeleton

Here is the skeleton I start every new environment from. It handles seeding, spaces, and the five-tuple return correctly:

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

class EnergyDispatchEnv(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"]}

    def __init__(self, render_mode=None, max_steps=1000):
        super().__init__()
        self.max_steps = max_steps
        self.render_mode = render_mode
        
        # Action: continuous charge/discharge rate in [-1, 1]
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32)
        
        # Observation: [state_of_charge, price, demand, time_of_day]
        self.observation_space = spaces.Box(
            low=np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32),
            high=np.array([1.0, 10.0, 5.0, 1.0], dtype=np.float32),
            dtype=np.float32
        )
        self._state = None
        self._step_count = 0

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)  # sets self.np_random
        self._step_count = 0
        soc = self.np_random.uniform(0.3, 0.7)
        price = self.np_random.uniform(0.5, 2.0)
        # ... initialize remaining state
        self._state = np.array([soc, price, 1.2, 0.25], dtype=np.float32)
        return self._state, {}

    def step(self, action):
        action = np.clip(action, self.action_space.low, self.action_space.high)
        self._step_count += 1
        
        # ... transition logic here
        reward = self._compute_reward(action)
        terminated = self._is_failure()  # e.g., SOC out of bounds
        truncated = self._step_count >= self.max_steps
        
        obs = self._get_obs()
        info = {"soc": float(obs[0])}
        
        return obs, reward, terminated, truncated, info

I've found that forcing every environment to go through this template catches 80% of integration errors before they reach the trainer. In particular, clipping actions inside step prevents NaNs when an exploratory policy samples near the bounds, and always returning a dict for info keeps wrappers like RecordEpisodeStatistics happy.

Common Violations I Still See in Code Review

Three mistakes show up repeatedly: returning done as a single boolean, mutating and returning the same observation array object (which vectorized trainers will buffer incorrectly), and using np.random directly inside step. The last one breaks determinism even when you pass a seed to reset. If you need reproducible rollouts for debugging or for PyTorch Documentation determinism checks downstream, this detail is non-negotiable.

Crafting Observation and Action Spaces That Don't Fight Your Optimizer

Your choice of spaces is not just an API detail; it determines what your policy network sees and how gradients flow. A poorly normalized observation will make PPO or SAC struggle even if the underlying dynamics are simple.

For observations, prefer spaces.Box with explicit low/high bounds and keep values roughly in the range [-1, 1] or [0, 1]. I've had environments where electricity prices in [0, 200] and state-of-charge in [0, 1] were concatenated raw. The value network put all its capacity into price because its magnitude dominated the loss. Normalizing price by dividing by a known max, or better, using a running normalization wrapper, stabilized training overnight. Use spaces.Dict only when you need heterogeneous modalities (e.g., image + vector). It is powerful but forces your policy to use a Dict feature extractor, which complicates sharing code across algorithms discussed in Deep RL Algorithms: DQN, PPO and SAC Compared.

For actions, be deliberate about discrete vs. continuous. Discrete(n) is the right choice when actions are truly categorical (e.g., select which machine to schedule). Box is correct for continuous control, but you need to decide if the raw action is a delta, a target setpoint, or a torque. I've found that action-as-delta (change in setpoint) trains more smoothly for systems with inertia because the policy starts near zero action, which corresponds to holding position, a sensible prior.

Validation Inside the Environment

Add assertions that your observation is always inside observation_space. Gymnasium provides a utility for this, but a simple check in development is valuable:

def _get_obs(self):
    # Normalize to [0, 1] range expected by policy
    raw_price = self._internal_state["price"]  # e.g., 0..150 $/MWh
    norm_price = raw_price / 150.0
    norm_price = np.clip(norm_price, 0.0, 1.0)
    
    obs = np.array([
        self._internal_state["soc"],  # already 0..1
        norm_price,
        self._internal_state["demand"] / 5.0,
        self._internal_state["time"] / 24.0
    ], dtype=np.float32)
    
    # Fail fast during development
    assert self.observation_space.contains(obs), f"Obs {obs} outside space"
    return obs

def step(self, action):
    # Defensive: ensure action is valid even if policy misbehaves
    if not self.action_space.contains(action):
        action = np.clip(action, self.action_space.low, self.action_space.high)
    # ... rest of step

Leave these asserts enabled for your first million steps. Once training is stable, you can gate them behind a DEBUG flag for throughput, but they have saved me from subtle bugs where a division by zero produced an inf observation that silently became NaN after normalization.

Shaping Rewards and Termination Logic Without Teaching Bad Habits

Reward design is where most custom environments fail. A sparse reward (e.g., +1 at goal, 0 otherwise) is faithful to the task but gives the agent almost no gradient. A dense, hand-shaped reward learns faster but can create perverse incentives if you are not careful. In one early logistics environment, I gave a small positive reward for moving closer to the target. The agent learned to oscillate back and forth near the target to accumulate distance rewards forever, never actually terminating.

My rule of thumb is: start with the true objective as the reward, then add shaping terms that are potential-based. Potential-based shaping, where F(s, s') = gamma * Phi(s') - Phi(s), is guaranteed not to change the optimal policy, but in practice even simple shaping works if you keep the scale balanced. Keep rewards roughly in [-1, 1] or [-10, 10] per step. If one component is 100x larger than another, the optimizer will ignore the smaller one.

Termination and truncation also need explicit design. Use terminated for true failure/success states that should be learned as absorbing states. Use truncated for time limits. For many industrial control problems, I set a generous time limit (e.g., 2048 steps) and treat exceeding it as truncated, not terminated. This lets the value function bootstrap correctly. If you mark time-limit exits as terminated, the agent thinks the world ends after N steps, which biases value estimates near the horizon. Recent work surveyed on arXiv Machine Learning consistently shows that incorrect handling of truncation is a leading source of non-reproducible results in custom environments.

The table below summarizes patterns I use to decide on reward structure before writing code:

Always log reward components separately in info. Instead of returning a single scalar, compute info["reward_components"] = {"task": ..., "energy": ..., "safety": ...}. This makes it possible to see during training whether the agent is improving the true objective or just gaming one term. I also run a simple sanity check before any RL training: roll out a random policy for 10 episodes and plot returns, terminated vs. truncated rates, and observation histograms. If a random policy gets high reward, your reward is too easy to exploit.

Wrappers, Vectorization, and Squeezing Throughput Out of Your Environment

A correct single-instance environment is not enough. Modern training runs 8 to 128 copies in parallel, and the overhead of Python loops can dominate runtime if you are not careful. Gymnasium provides two mechanisms for this: wrappers and vector environments.

Wrappers are small decorators that modify observations, rewards, or termination without editing the base class. I rely on a short, ordered stack: TimeLimit for horizon control, RecordEpisodeStatistics for logging, and optionally NormalizeObservation or ClipAction. Order matters. Normalization should wrap the base env, but RecordEpisodeStatistics should be outermost so it sees the final rewards. Don't write your own normalization when the standard wrapper already handles running mean/variance correctly.

For vectorization, SyncVectorEnv is sufficient up to about 16 envs and is easier to debug because it runs in the same process. Beyond that, AsyncVectorEnv (or the newer gymnasium.vector.AsyncVectorEnv with shared memory) pays off, but it uses multiprocessing and can hide state-leak bugs. In my experience, the threshold where async helps is when your step takes more than ~0.5 ms. For fast simulators, sync is often faster due to IPC overhead.

Putting It Together: Registration and Vector Creation

import gymnasium as gym
from gymnasium.wrappers import TimeLimit, RecordEpisodeStatistics, NormalizeObservation
from gymnasium.envs.registration import register

# Register once at import time
register(
    id="EnergyDispatch-v0",
    entry_point="my_pkg.envs:EnergyDispatchEnv",
    max_episode_steps=2048,
)

def make_env(seed=0):
    def _thunk():
        env = gym.make("EnergyDispatch-v0", max_steps=2048)
        env = RecordEpisodeStatistics(env)  # outermost: logs episodic return
        # NormalizeObservation must be applied after creation, before vectorization
        # env = NormalizeObservation(env)  # uncomment after debugging
        return env
    return _thunk

# Vectorized training with 32 parallel envs
vec_env = gym.vector.SyncVectorEnv([make_env(seed=i) for i in range(32)])
obs, info = vec_env.reset(seed=42)
# vec_env auto-batches obs to shape (32, obs_dim) and rewards to (32,)

One detail that trips up teams: gym.make with max_episode_steps automatically wraps with TimeLimit which sets truncated=True on timeout. If you also manually implement a step count inside your env and set truncated, you will double-truncate. Pick one place to enforce the horizon. I prefer letting registration handle it and removing internal step counting entirely.

For performance, profile your step in isolation. I've seen 30% speedups from replacing a for loop over NumPy arrays with a vectorized operation, and another 2x from caching static arrays instead of allocating new ones each step. If you are bridging to a physics simulator for Sim-to-Real Transfer: Training RL Agents in Simulation for Real-World Robotics, batch your simulator calls. Stepping 32 envs one by one through a Python binding is often 10x slower than a batched C++ call.

Testing for Determinism, Correctness, and Silent Simulation Bugs

Custom environments have a large surface area for silent bugs that don't crash but corrupt learning. I maintain a small test harness that every environment must pass before training.

First, the Gymnasium checker: gym.utils.env_checker.check_env(env). It validates spaces, dtypes, and the five-tuple return. Run it once; it is not meant for production loops.

Second, determinism. Reset with the same seed and apply the same action sequence, assert that observations and rewards match exactly. If they don't, you are using global RNG state or unseeded noise. Third, check the Markov property manually: the same (state, action) should produce the same distribution over next_state, regardless of history. A common violation is storing episode-level counters in the observation without including them in state, or using a reward that depends on the previous action not encoded in the observation.

What I Log on Every Run

Beyond standard returns, I log three diagnostics: episode length histogram, terminated vs. truncated ratio, and observation marginal distributions. If episode lengths collapse to 1 or 2 steps, your termination condition is too aggressive. If truncated is always 100%, the agent never learns to succeed or fail — it just times out. And if one observation dimension drifts outside its declared bounds after 1M steps, your dynamics have a leak.

I also keep one deterministic golden trajectory: a fixed seed and a fixed action sequence with expected observations checked into the repo as a JSON snapshot. It acts as a regression test when someone modifies the dynamics. Without it, a refactor that subtly changes physics can go unnoticed for days of training.

Packaging and Versioning Environments for Reproducible Training Pipelines

An environment is part of your model. If you update the reward scale or the observation normalization without versioning the environment ID, you invalidate all previous checkpoints and make hyperparameter comparisons meaningless. I follow the Gym convention: EnergyDispatch-v0, EnergyDispatch-v1, etc., and bump the version whenever the transition, reward, or space changes. The code for v0 stays importable so old experiments can be reproduced.

Package the environment as an installable Python package with a pyproject.toml and an entry point. Avoid relative imports that only work from a notebook directory. Register environments on import in your package's __init__.py so gym.make("MyEnv-v0") works anywhere after pip install -e .. Pin gymnasium, numpy, and any simulator dependencies in a lock file. I have debugged too many failures where a minor NumPy release changed clipping behavior.

For reproducibility, store the full environment config alongside model checkpoints: the env ID, the seed, the wrapper stack and their parameters, and the max_episode_steps. When using record-and-replay for debugging, set PYTHONHASHSEED=0 and torch.backends.cudnn.deterministic=True if your policy uses PyTorch — otherwise even a deterministic environment can appear stochastic due to non-deterministic GPU kernels. Document the render mode and whether you use rgb_array for video logging; that choice affects whether you can reconstruct visual rollouts later.

Finally, keep training and evaluation environments separate. Create a maker function make_train_env() that includes exploration noise and reward shaping, and make_eval_env() that uses the sparse true objective and no normalization wrappers that depend on training statistics. Evaluating on the training wrapper is one of the easiest ways to overestimate performance.

Frequently Asked Questions