Gymnasium Custom Environment: 7 Patterns That Save Hours

Why Most Custom Environments Break in Subtle Ways

Your custom Gymnasium environment passes check_env(), trains for 10k steps without errors, then crashes with a cryptic shape mismatch at step 47,293. Or worse — it trains fine but never converges, and you spend days tweaking PPO hyperparameters when the real bug is in your reset() method.

I’ve debugged dozens of custom environments over the past two years. The pattern is always the same: the obvious mistakes (wrong action space type, missing render() method) get caught immediately. The subtle ones — observation normalization inconsistencies, edge cases in step() transitions, incorrect terminal state handling — waste days of debugging and GPU time.

This post covers seven patterns that catch 90% of custom environment bugs before they reach training. Not the basics you’d find in the official Gymnasium tutorial, but the production patterns I wish I’d known earlier.

Pattern 1: Normalize Observations in the Environment, Not the Wrapper

Most tutorials tell you to use NormalizeObservation wrapper. That works for built-in environments, but breaks down when you add domain-specific constraints.

Here’s what happens: you build a robotic arm environment where joint angles are in $[-\pi, \pi]$ radians, velocities in $[-5, 5]$ rad/s, and end-effector position in $[0, 1]$ meters (normalized workspace). You wrap it with NormalizeObservation, which computes running mean and std across all observation dimensions equally.

But that’s wrong. Joint angles are periodic — wrapping from $-\pi$ to $\pi$ is valid. Normalizing them with running statistics breaks the periodicity. End-effector position is already normalized. Only velocities need standardization.

Better approach: normalize inside the environment, dimension-specific.

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

class RoboticArmEnv(gym.Env):
    def __init__(self):
        super().__init__()
        # Joint: [angle, velocity] x 3 joints + end_effector [x, y, z]
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(9,), dtype=np.float32
        )
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(3,), dtype=np.float32
        )

        # Domain knowledge: typical velocity ranges from real robot logs
        self.velocity_scale = 5.0  # max observed velocity (rad/s)

    def _get_obs(self):
        """Convert raw physics state to normalized observation."""
        angles = self.state[:3]  # [-pi, pi]
        velocities = self.state[3:6]  # unbounded
        ee_pos = self.state[6:9]  # [0, 1] (pre-normalized workspace)

        # Normalize angles to [-1, 1] preserving periodicity
        norm_angles = angles / np.pi

        # Clip and normalize velocities based on domain knowledge
        norm_velocities = np.clip(velocities / self.velocity_scale, -1, 1)

        # ee_pos already in [0,1], map to [-1, 1] for consistency
        norm_ee_pos = 2 * ee_pos - 1

        return np.concatenate([norm_angles, norm_velocities, norm_ee_pos]).astype(np.float32)

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # Initialize with small random perturbations
        self.state = np.zeros(9)
        self.state[:3] = self.np_random.uniform(-0.1, 0.1, size=3)  # near zero angles
        self.state[6:9] = np.array([0.5, 0.5, 0.5])  # center of workspace
        return self._get_obs(), {}

    def step(self, action):
        # Physics update (simplified)
        self.state[:3] += 0.01 * action  # angle update
        self.state[3:6] = action * 2.0  # velocity from action

        # Wrap angles to [-pi, pi]
        self.state[:3] = np.arctan2(np.sin(self.state[:3]), np.cos(self.state[:3]))

        obs = self._get_obs()
        reward = -np.linalg.norm(self.state[6:9] - np.array([0.8, 0.5, 0.5]))
        terminated = False
        truncated = False
        return obs, reward, terminated, truncated, {}

The key insight: _get_obs() applies dimension-specific normalization based on domain knowledge, not blind statistics. Angles stay periodic, velocities get clipped to reasonable ranges (preventing outlier explosion during early random exploration), and workspace coordinates remain interpretable.

This approach also makes the observation space bounds honest. Setting low=-np.inf, high=np.inf signals to the algorithm that observations are pre-normalized, not raw sensor values with unknown ranges.

Pattern 2: Separate Termination from Truncation (And Get It Right)

Gymnasium 0.26+ split the old done flag into terminated (episode ended naturally) and truncated (hit time limit or external stop). Most environments I’ve reviewed still get this wrong.

The distinction matters for value bootstrapping. When terminated=True, the value of the next state is zero (episode is over). When truncated=True, the agent should bootstrap from the value estimate of the truncated state (there IS a next state, we just stopped observing it).

Common mistake: treating out-of-bounds as termination.

# WRONG: treats boundary violation as terminal state
def step(self, action):
    self.state += action

    reward = -np.linalg.norm(self.state - self.goal)

    # Out of bounds = episode over?
    if np.any(np.abs(self.state) > 10):
        return self.state, -10, True, False, {}  # terminated=True

    return self.state, reward, False, False, {}

But going out of bounds isn’t a natural episode termination — it’s a constraint violation. The world doesn’t end. The robot just hit a limit switch. Better:

# BETTER: constrain state, penalize, keep episode running
def step(self, action):
    self.state += action

    # Constrain to valid region
    penalty = 0
    if np.any(np.abs(self.state) > 10):
        self.state = np.clip(self.state, -10, 10)
        penalty = -5  # one-time penalty for hitting boundary

    reward = -np.linalg.norm(self.state - self.goal) + penalty

    # Only terminate on goal reached
    terminated = np.linalg.norm(self.state - self.goal) < 0.1
    truncated = False  # let TimeLimit wrapper handle this

    return self.state, reward, terminated, truncated, {}

Now the value function learns that hitting boundaries is bad (penalty) but not terminal. The episode continues, forcing the agent to recover. This matches physical reality better.

For time limits, use Gymnasium’s TimeLimit wrapper instead of manually counting steps:

from gymnasium.wrappers import TimeLimit

env = RoboticArmEnv()
env = TimeLimit(env, max_episode_steps=1000)

The wrapper sets truncated=True at step 1000, signaling to the algorithm that bootstrapping is valid.

Pattern 3: Make reset() Stochastic (But Reproducible)

Deterministic reset() is the silent killer of generalization. Your agent learns to exploit the exact initial state instead of the task.

I trained a maze-solving agent that achieved 95% success rate in training. Deployed it with a different random seed — 12% success. The agent had memorized the starting position, not learned wall-following.

Gymnasium’s reset(seed=...) API makes randomization reproducible:

def reset(self, seed=None, options=None):
    super().reset(seed=seed)  # sets self.np_random

    # Bad: always start at origin
    # self.agent_pos = np.array([0, 0])

    # Good: random start in valid region
    while True:
        pos = self.np_random.uniform(-5, 5, size=2)
        if self._is_valid_position(pos):  # not inside wall
            self.agent_pos = pos
            break

    # Also randomize goal (if task allows)
    self.goal_pos = self.np_random.uniform(-5, 5, size=2)

    return self._get_obs(), {}

Critical: use self.np_random (seeded RNG from parent class), NOT np.random (global RNG). This makes the environment reproducible when called with the same seed sequence, which is essential for debugging and benchmarking.

For curriculum learning, use the options dict:

def reset(self, seed=None, options=None):
    super().reset(seed=seed)

    difficulty = options.get('difficulty', 1.0) if options else 1.0

    # Harder = goal farther from start
    max_dist = 5.0 * difficulty

    self.agent_pos = np.array([0, 0])
    self.goal_pos = self.np_random.uniform(-max_dist, max_dist, size=2)

    return self._get_obs(), {}

# Training loop
obs, info = env.reset(options={'difficulty': current_success_rate})

This lets you gradually increase difficulty as the agent improves, without modifying the environment class.

Pattern 4: Return Dense Info Dicts (Even If You Don’t Use Them Yet)

The info dict returned by step() is optional, but skipping it is a mistake. You’ll need debugging metrics eventually — add them from day one.

Minimum useful info:

def step(self, action):
    prev_state = self.state.copy()
    self.state += action * 0.1

    dist_to_goal = np.linalg.norm(self.state - self.goal)
    reward = -dist_to_goal

    terminated = dist_to_goal < 0.1
    truncated = False

    info = {
        'distance_to_goal': dist_to_goal,
        'state_change': np.linalg.norm(self.state - prev_state),
        'is_success': terminated,  # for success rate tracking
        'action_norm': np.linalg.norm(action),
    }

    return self._get_obs(), reward, terminated, truncated, info

When training diverges at episode 8374, you can now plot distance_to_goal and see it suddenly spiked, suggesting an exploration outlier. Without info, you’d be blind.

Stable Baselines3 automatically logs is_success if present, giving you success rate curves for free. I covered this in a previous post on reward shaping — sparse reward environments often need success rate as the primary metric, not mean reward.

Pattern 5: Test Edge Cases in reset(), Not Just step()

Most developers test step() logic carefully but treat reset() as trivial. Then their agent dies at step 0 with a shape mismatch.

Edge cases to verify:

import pytest
import numpy as np

def test_reset_observation_shape():
    env = MyEnv()
    obs, info = env.reset(seed=42)
    assert obs.shape == env.observation_space.shape
    assert obs.dtype == env.observation_space.dtype

def test_reset_determinism():
    env1, env2 = MyEnv(), MyEnv()
    obs1, _ = env1.reset(seed=123)
    obs2, _ = env2.reset(seed=123)
    np.testing.assert_array_equal(obs1, obs2)

def test_reset_after_termination():
    """Reset should work after terminated=True."""
    env = MyEnv()
    obs1, _ = env.reset(seed=1)

    # Force termination
    while True:
        obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
        if terminated or truncated:
            break

    # Reset should not crash
    obs2, _ = env.reset(seed=2)
    assert obs2.shape == env.observation_space.shape

def test_reset_observation_bounds():
    """Initial observation should respect declared bounds."""
    env = MyEnv()
    for _ in range(100):  # sample multiple resets
        obs, _ = env.reset()
        assert env.observation_space.contains(obs), f"Reset obs {obs} outside space {env.observation_space}"

The last test catches a bug I see constantly: declaring observation_space = Box(low=-1, high=1, ...) but returning observations outside that range. RL algorithms use these bounds for normalization — if they’re wrong, training silently degrades.

Pattern 6: Use Spaces Correctly (Box vs Discrete vs Dict)

Gymnasium has six space types. Most custom environments only use Box and Discrete, but choosing wrong kills sample efficiency.

Box: Continuous values. Use for positions, velocities, images.

spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32)

Discrete: Integer from [0, n). Use for categorical actions (move N/S/E/W), not discretized continuous actions.

spaces.Discrete(4)  # 4 discrete actions: {0, 1, 2, 3}

MultiBinary: Binary vector. Use for multi-label classification or on/off switches.

spaces.MultiBinary(5)  # [0,1,1,0,1] - five independent binary flags

Dict: Heterogeneous observations. Use when observation contains both continuous and discrete parts.

spaces.Dict({
    'image': spaces.Box(low=0, high=255, shape=(64, 64, 3), dtype=np.uint8),
    'goal': spaces.Box(low=-10, high=10, shape=(2,), dtype=np.float32),
    'inventory': spaces.MultiBinary(10),  # which items are held
})

Common mistake: using Box for categorical data.

# WRONG: action_space = Box(low=0, high=3.99, shape=(1,))
# Agent might output 1.7, what does that mean?

# RIGHT: action_space = Discrete(4)

Another: using Discrete for discretized continuous actions.

# BAD: discretize velocity into 10 bins
action_space = Discrete(10)

# Convert action to velocity
velocity = (action - 5) * 0.2  # maps {0..9} to {-1..0.8}

This works, but forces the agent to learn ordinal structure (action 5 is “between” 4 and 6). Policy gradient methods learn this slowly. If the underlying action is continuous, keep it continuous:

# BETTER: let agent output continuous velocity directly
action_space = Box(low=-1, high=1, shape=(1,), dtype=np.float32)

Use Discrete only when actions are truly categorical (turn left vs right vs jump) with no ordering.

Pattern 7: Implement render() Even If You Don’t Need It

Skipping render() is tempting. “I’ll add it later when I need to debug.” Then later arrives at 2am after six hours of debugging, and you have no idea what your agent is doing.

Minimum viable render():

def render(self):
    if self.render_mode == 'human':
        print(f"Agent: {self.agent_pos}, Goal: {self.goal_pos}, Distance: {np.linalg.norm(self.agent_pos - self.goal_pos):.2f}")
    elif self.render_mode == 'rgb_array':
        # Return a simple matplotlib figure as image
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(5, 5))
        ax.set_xlim(-10, 10)
        ax.set_ylim(-10, 10)
        ax.plot(self.agent_pos[0], self.agent_pos[1], 'bo', markersize=10, label='Agent')
        ax.plot(self.goal_pos[0], self.goal_pos[1], 'r*', markersize=15, label='Goal')
        ax.legend()
        fig.canvas.draw()
        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        plt.close(fig)
        return image

This takes 10 minutes to write and saves hours of debugging. When your reward curve plateaus, just record a video:

from gymnasium.wrappers import RecordVideo

env = MyEnv(render_mode='rgb_array')
env = RecordVideo(env, video_folder='./videos/', episode_trigger=lambda x: x % 10 == 0)

Now you get MP4s of every 10th episode, showing exactly what failure mode the agent is stuck in.

For text-based environments (no spatial layout), render_mode='human' with print statements is fine. For spatial environments, rgb_array is worth the matplotlib boilerplate.

When These Patterns Actually Matter

Honestly? If you’re just following a tutorial and training on CartPole, none of this matters. The real value shows up when:

• You’re adapting a research paper environment (say, ProcGen or a MuJoCo variant) to your specific task
• Your environment has domain-specific physics (robotic constraints, financial market rules, etc.)
• You’re stuck debugging why SAC converges but PPO doesn’t — often the culprit is observation normalization or terminal state handling

I’m still not entirely sure why terminated vs truncated matters so much for SAC specifically (theory says off-policy algorithms should be robust to this), but empirically it does. My best guess: SAC’s entropy term is sensitive to value estimate errors near episode boundaries, and incorrect bootstrapping amplifies this.

And if you’re building an environment for others to use (releasing it as a benchmark), all seven patterns become mandatory. Users will find edge cases you never tested.

FAQ

Q: Should I use gym.make() or instantiate my environment class directly?

Direct instantiation (env = MyEnv()) is fine for development. Use gym.make() only if you plan to register your environment for others:

from gymnasium.envs.registration import register

register(
    id='MyEnv-v0',
    entry_point='my_module:MyEnv',
    max_episode_steps=1000,
)

env = gym.make('MyEnv-v0')

Registration lets you apply wrappers by default and version your environment (v0, v1, etc.).

Q: How do I debug why my environment trains on PPO but not SAC?

Usually observation scaling. SAC is more sensitive to input ranges than PPO (batch normalization in the critic helps, but not always). Print obs.min(), obs.max(), obs.std() after reset() and confirm they’re in a reasonable range (roughly [-3, 3] if Gaussian-normalized). If you see values like 847.3 or 0.00012, your normalization is broken.

Q: My environment works in single-process mode but crashes with SubprocVecEnv. Why?

Your environment likely stores non-picklable state (like a lambda function, a TensorFlow session, or a threading.Lock). SubprocVecEnv uses multiprocessing.Process, which pickles the environment. Move non-picklable objects to render() or lazy-initialize them in reset().

The Part Nobody Mentions

Most Gymnasium tutorials show you how to build an environment that works. What they skip: building an environment where RL algorithms actually converge.

The difference is subtle. A broken environment crashes with a stack trace. A poorly designed environment trains for 5 million steps, occasionally shows promising reward spikes, then plateaus at 60% success rate forever. You tweak PPO’s clip range, adjust the learning rate schedule, add reward shaping — none of it helps, because the root cause is observation normalization or terminal state handling.

These seven patterns won’t magically make your agent solve the task. But they remove the silent failure modes that waste weeks of debugging and hyperparameter tuning. If your environment is built on these foundations, and your agent still doesn’t converge, at least you know the bug is in your algorithm or reward function — not in some edge case in reset() that fires once every 10,000 episodes.

I’m currently experimenting with auto-tuning observation normalization bounds using percentile statistics from the first 10k steps (like Optuna, but for environment design). Not sure if it’ll work at scale, but the manual approach definitely doesn’t scale when you have 20+ observation dimensions.

Oh, and if you’re reading environment code at 3am trying to figure out why your robot keeps spinning in circles — grab some Dark Chocolate Espresso Beans. The caffeine-to-sanity ratio is unbeatable.