Gymnasium Custom Env Step() Returns Invalid Shape: 5 Fixes

Q: Q: Why does my observation pass observation_space.contains() but training still fails?

The contains() method checks bounds and shape but may not enforce dtype strictly in all Gymnasium versions. Always explicitly match dtype by using np.asarray(obs, dtype=np.float32) or specifying dtype in array creation. Some RL libraries also have their own internal checks that are stricter.

Q: Q: Should I use Box with infinite bounds or finite bounds for continuous observations?

Use infinite bounds (-np.inf, np.inf) unless you have a genuine physical constraint. Finite bounds require clipping logic that can mask bugs and create discontinuities at edges. If your state is truly bounded (robot joint angles, normalized features), then bounded spaces help—just remember to clip i

Disclosure: As an Amazon Associate, I earn from qualifying purchases. Some links in this post are affiliate links — they cost you nothing extra.

⚡ Key Takeaways

Explicit dtype=np.float32 in both observation_space definition and returned arrays prevents silent dtype mismatch bugs.
Always return obs.copy() from step() and reset() to avoid view aliasing that corrupts replay buffers.
Single-element observations collapse from shape (1,) to () unless you use np.array([value]) with explicit shape.
Float accumulation can push bounded observations past limits—use np.clip() before returning.
Stable Baselines3 check_env() only runs a few steps and misses long-horizon bugs, so write custom validation tests.

The Error That Breaks Every Custom Gymnasium Environment

Your step() function returns a perfectly valid numpy array. Stable Baselines3 throws AssertionError: The observation returned by thestep()method does not match the given observation space. You stare at your code for 20 minutes, print the shape, it looks correct, and yet the environment refuses to work.

This specific error has cost me more debugging hours than any algorithm hyperparameter ever did. The shape mismatch between observation_space and actual returned observations is the single most common bug when building custom Gymnasium environments, and it fails silently in ways that will make you question your sanity.

Here’s the working pattern I now use for every custom environment:

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

class RobustCustomEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.state_dim = 4

        # Bug fix #1: Always use explicit dtype
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(self.state_dim,),
            dtype=np.float32  # CRITICAL: explicit dtype
        )
        self.action_space = spaces.Discrete(2)
        self._state = None

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # Bug fix #2: Cast to correct dtype immediately
        self._state = np.zeros(self.state_dim, dtype=np.float32)
        return self._state.copy(), {}  # Bug fix #3: Return copy

    def step(self, action):
        # Bug fix #4: Validate before any processing
        assert self._state is not None, "Call reset() before step()"

        # Your environment logic here
        self._state = self._state + 0.1 * np.random.randn(self.state_dim)

        reward = 1.0
        terminated = False
        truncated = False

        # Bug fix #5: Force correct shape and dtype at return
        obs = np.asarray(self._state, dtype=np.float32).reshape(self.state_dim,)
        return obs, reward, terminated, truncated, {}

That’s the template. Now let me show you the five bugs it prevents—each one I’ve hit in production environments.

Two girls in school uniforms walking inside a gymnasium with playful decor. — Photo by Ron Lach on Pexels

Bug #1: dtype Mismatch Causes Silent Assertion Failures

This is the sneakiest one. Your shape looks correct, but the dtype is wrong.

# BROKEN: dtype defaults to float64
class BrokenDtypeEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.observation_space = spaces.Box(
            low=-1.0, high=1.0, shape=(4,)
        )  # dtype defaults to float64!
        self.action_space = spaces.Discrete(2)

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # This is float64 by default
        return np.zeros(4), {}

    def step(self, action):
        # Also float64
        obs = np.random.randn(4)
        return obs, 0.0, False, False, {}

Here’s what happens when you run Stable Baselines3’s check_env:

from stable_baselines3.common.env_checker import check_env

env = BrokenDtypeEnv()
check_env(env)  # Passes! But...

It passes. But when you actually train:

from stable_baselines3 import PPO

model = PPO("MlpPolicy", env)
model.learn(total_timesteps=1000)
# UserWarning: Your observation space dtype is float64, 
# but your observation dtype is float64. SB3 expects float32.

The warning appears, but training continues—with subtle numerical instabilities. The policy network internally casts to float32, introducing rounding that accumulates over millions of steps.

The fix is dead simple but easy to forget:

self.observation_space = spaces.Box(
    low=-1.0, high=1.0, shape=(4,), dtype=np.float32
)

And then ensure every return matches:

return np.zeros(4, dtype=np.float32), {}

Enjoying this article? Get more like it delivered to your inbox. Subscribe to the newsletter

Bug #2: Nested Observations Lose Their Shape

Dict and Tuple observation spaces are powerful but error-prone. The most common failure mode: one component has the wrong shape.

# BROKEN: image shape is (H, W, C) but observation_space expects (C, H, W)
class BrokenDictEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.observation_space = spaces.Dict({
            "image": spaces.Box(0, 255, shape=(3, 64, 64), dtype=np.uint8),
            "vector": spaces.Box(-np.inf, np.inf, shape=(10,), dtype=np.float32)
        })
        self.action_space = spaces.Discrete(4)

    def step(self, action):
        obs = {
            "image": np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8),  # WRONG: (H,W,C)
            "vector": np.random.randn(10).astype(np.float32)
        }
        return obs, 0.0, False, False, {}

The error message you get:

AssertionError: The observation returned by the `step()` method does not match 
the given observation space. Expected: Box(0, 255, (3, 64, 64), uint8), 
actual: (64, 64, 3)

The fix requires careful attention to channel ordering. If you’re using images, Gymnasium typically expects channels-last (H, W, C), but some policies (especially CNNs from PyTorch) expect channels-first (C, H, W). Pick one and stick with it:

# FIXED: Consistent channel-last format
self.observation_space = spaces.Dict({
    "image": spaces.Box(0, 255, shape=(64, 64, 3), dtype=np.uint8),
    "vector": spaces.Box(-np.inf, np.inf, shape=(10,), dtype=np.float32)
})

def step(self, action):
    obs = {
        "image": np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8),
        "vector": np.random.randn(10).astype(np.float32)
    }
    return obs, 0.0, False, False, {}

I’d recommend a validation helper that you call at the end of both reset() and step():

def _validate_obs(self, obs):
    """Debug helper - remove in production for speed."""
    try:
        assert self.observation_space.contains(obs), \
            f"Obs not in space. Got shapes: {self._get_shapes(obs)}"
    except Exception as e:
        print(f"Observation validation failed: {e}")
        print(f"Expected: {self.observation_space}")
        if isinstance(obs, dict):
            for k, v in obs.items():
                print(f"  {k}: shape={v.shape}, dtype={v.dtype}")
        raise
    return obs

def _get_shapes(self, obs):
    if isinstance(obs, dict):
        return {k: v.shape for k, v in obs.items()}
    return obs.shape

Bug #3: Returning Views Instead of Copies

This one is subtle and won’t show up until you run vectorized environments or replay buffers.

# BROKEN: Returns view of internal state
class ViewBugEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.observation_space = spaces.Box(-np.inf, np.inf, (4,), dtype=np.float32)
        self.action_space = spaces.Discrete(2)
        self._state = np.zeros(4, dtype=np.float32)

    def step(self, action):
        self._state += action * 0.1
        # DANGER: returning the actual internal array, not a copy
        return self._state, 1.0, False, False, {}

Why is this dangerous? Because the caller might store that reference:

env = ViewBugEnv()
obs1, _ = env.reset()
obs2, _, _, _, _ = env.step(0)
obs3, _, _, _, _ = env.step(1)

print(obs1)  # These are all the SAME array!
print(obs2)  # All pointing to env._state
print(obs3)  # Whatever the final value is

Replay buffers store observations. If you return views, every stored observation becomes a pointer to your current state. Your buffer contains $N$ identical observations—the last one.

The fix is simple:

return self._state.copy(), 1.0, False, False, {}

Or use np.array(self._state) which also creates a copy. The Gymnasium documentation mentions this, but it’s buried and easy to miss.

A coach instructing a diverse group of children during a gym class in a school sports hall. — Photo by cottonbro studio on Pexels

Bug #4: Dimension Collapse on Single-Element Observations

This one trips up everyone who builds environments with configurable state dimensions.

# BROKEN when state_dim=1
class CollapseDimEnv(gym.Env):
    def __init__(self, state_dim=1):
        super().__init__()
        self.state_dim = state_dim
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, 
            shape=(state_dim,),  # Shape is (1,)
            dtype=np.float32
        )
        self.action_space = spaces.Discrete(2)

    def step(self, action):
        # This creates a scalar when state_dim=1!
        value = np.random.randn()
        if self.state_dim == 1:
            obs = np.float32(value)  # Shape is () not (1,)!
        else:
            obs = np.random.randn(self.state_dim).astype(np.float32)
        return obs, 0.0, False, False, {}

The observation space expects shape (1,), but np.float32(value) creates a 0-dimensional array with shape (). They’re not the same:

a = np.float32(0.5)
print(a.shape)  # ()

b = np.array([0.5], dtype=np.float32)
print(b.shape)  # (1,)

The validation error you’ll see:

AssertionError: Expected shape (1,), got shape ()

Always construct arrays with explicit shape:

obs = np.array([value], dtype=np.float32)  # Forces (1,) shape
# or
obs = np.zeros(self.state_dim, dtype=np.float32)
obs[0] = value

Bug #5: Bounds Violations From Float Precision

If your observation space has finite bounds, floating-point arithmetic can push values just outside:

# BROKEN: Position can exceed bounds due to float precision
class BoundsViolationEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.observation_space = spaces.Box(
            low=np.array([0.0, 0.0], dtype=np.float32),
            high=np.array([1.0, 1.0], dtype=np.float32),
            dtype=np.float32
        )
        self.action_space = spaces.Discrete(4)
        self._pos = np.array([0.5, 0.5], dtype=np.float32)

    def step(self, action):
        delta = 0.1
        # After many steps, accumulated float errors push _pos to 1.0000001
        if action == 0: self._pos[0] += delta
        if action == 1: self._pos[0] -= delta
        if action == 2: self._pos[1] += delta
        if action == 3: self._pos[1] -= delta

        # Position can be 1.0000001 or -0.0000001
        return self._pos.copy(), 0.0, False, False, {}

After enough steps, you’ll hit:

AssertionError: Observation [1.0000001, 0.5] is not within observation space 
bounds [0.0, 1.0]

The fix is to clip explicitly:

def step(self, action):
    # ... movement logic ...

    # Clip to bounds
    obs = np.clip(
        self._pos, 
        self.observation_space.low, 
        self.observation_space.high
    ).astype(np.float32)

    return obs.copy(), 0.0, False, False, {}

For unbounded spaces (using np.inf), this isn’t an issue. But if you’re modeling physical systems with real constraints—robot joint angles, positions in a grid, normalized features—bounded spaces are semantically correct and help with policy stability. Just remember to enforce them.

How to Debug When You’re Stuck

When check_env passes but training fails, here’s my debugging checklist:

import gymnasium as gym
import numpy as np

def diagnose_env(env_class, num_steps=100):
    """Run diagnostic on custom environment."""
    env = env_class()

    print(f"Observation space: {env.observation_space}")
    print(f"Expected dtype: {env.observation_space.dtype}")
    print(f"Expected shape: {env.observation_space.shape}")
    print()

    # Test reset
    obs, info = env.reset()
    print(f"reset() obs shape: {obs.shape if hasattr(obs, 'shape') else type(obs)}")
    print(f"reset() obs dtype: {obs.dtype if hasattr(obs, 'dtype') else type(obs)}")
    print(f"reset() in space: {env.observation_space.contains(obs)}")
    print()

    # Test multiple steps
    shapes_seen = set()
    dtypes_seen = set()
    bounds_violations = 0

    for i in range(num_steps):
        action = env.action_space.sample()
        obs, reward, term, trunc, info = env.step(action)

        shape = obs.shape if hasattr(obs, 'shape') else 'scalar'
        dtype = obs.dtype if hasattr(obs, 'dtype') else type(obs).__name__

        shapes_seen.add(str(shape))
        dtypes_seen.add(str(dtype))

        if not env.observation_space.contains(obs):
            bounds_violations += 1
            if bounds_violations <= 3:  # Only print first few
                print(f"Step {i}: Bounds violation! obs={obs}")

        if term or trunc:
            obs, info = env.reset()

    print(f"Shapes seen: {shapes_seen}")
    print(f"Dtypes seen: {dtypes_seen}")
    print(f"Bounds violations: {bounds_violations}/{num_steps}")

    env.close()

# Usage
diagnose_env(YourCustomEnv)

This catches most issues. Run it before integrating with any RL library.

Why check_env() Doesn’t Catch Everything

Stable Baselines3’s check_env is useful but limited. It only runs a few steps. Time-dependent bugs—like the bounds violation from accumulated float errors or the view-vs-copy issue that manifests in replay buffers—won’t appear.

The checks it performs include:

Observation/action space validity
Reset returns correct format
Step returns 5-tuple with correct types
Observation matches space (for a few samples)
Render mode validity

What it doesn’t check:

Long-horizon bounds stability
Copy semantics (view aliasing)
Dtype consistency under all code paths
Dict observation key consistency across episodes

I’ve started adding this to my test suites:

import pytest

def test_env_long_horizon():
    env = MyCustomEnv()
    obs, _ = env.reset()

    for _ in range(10000):  # Long episode
        action = env.action_space.sample()
        obs, _, term, trunc, _ = env.step(action)

        # Explicit checks every step
        assert obs.dtype == np.float32, f"dtype is {obs.dtype}"
        assert obs.shape == env.observation_space.shape
        assert env.observation_space.contains(obs), f"OOB: {obs}"

        if term or trunc:
            obs, _ = env.reset()

def test_env_copy_semantics():
    """Ensure observations are copies, not views."""
    env = MyCustomEnv()

    observations = []
    obs, _ = env.reset()
    observations.append(obs.copy())  # Store a copy

    for _ in range(100):
        obs, _, term, trunc, _ = env.step(env.action_space.sample())
        observations.append(obs.copy())
        if term or trunc:
            break

    # If env returns views, early observations would have been mutated
    # Check that observations are different
    unique_obs = len(set(tuple(o.flatten()) for o in observations))
    assert unique_obs > 1, "All observations are identical - returning views?"

The Reset Protocol Trap

Gymnasium 0.26+ changed the reset signature from reset() -> obs to reset(seed=None, options=None) -> (obs, info). This causes shape errors that look like observation space bugs but aren’t:

# Old style (Gym 0.21-0.25)
def reset(self):
    return self.initial_state()

# New style (Gymnasium 0.26+)
def reset(self, seed=None, options=None):
    super().reset(seed=seed)
    return self.initial_state(), {}

If you forget the info dict, you’ll see:

TypeError: cannot unpack non-iterable ndarray object

Or worse, if the caller expects (obs, info) and gets just obs, the tuple unpacking interprets the array elements as the two return values.

Always inherit from gym.Env and call super().reset(seed=seed) to handle seeding properly. The migration guide covers this but, honestly, the errors are cryptic enough that many people don’t connect them to the reset signature.

Performance Tip: Pre-allocate Everything

Once your environment works correctly, allocating arrays inside step() becomes the bottleneck. Pre-allocate in __init__:

class FastEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.observation_space = spaces.Box(-np.inf, np.inf, (100,), dtype=np.float32)
        self.action_space = spaces.Discrete(4)

        # Pre-allocate
        self._obs_buffer = np.zeros(100, dtype=np.float32)
        self._state = np.zeros(100, dtype=np.float32)

    def step(self, action):
        # Modify in-place
        np.copyto(self._obs_buffer, self._state)
        # ... update self._state ...

        # Return copy of buffer
        return self._obs_buffer.copy(), 0.0, False, False, {}

This avoids allocation overhead in the hot path. For environments called millions of times during training, this matters.

And if you’re doing late-night RL debugging sessions, Dark Chocolate Espresso Beans are genuinely helpful—caffeine plus antioxidants, no brewing required.

FAQ

Q: Why does my observation pass observation_space.contains() but training still fails?

The contains() method checks bounds and shape but may not enforce dtype strictly in all Gymnasium versions. Always explicitly match dtype by using np.asarray(obs, dtype=np.float32) or specifying dtype in array creation. Some RL libraries also have their own internal checks that are stricter.

Q: Should I use Box with infinite bounds or finite bounds for continuous observations?

Use infinite bounds (-np.inf, np.inf) unless you have a genuine physical constraint. Finite bounds require clipping logic that can mask bugs and create discontinuities at edges. If your state is truly bounded (robot joint angles, normalized features), then bounded spaces help—just remember to clip in step().

Q: How do I debug Dict/Tuple observation spaces when the error message is unhelpful?

Iterate through each key/index manually and check containment individually:

for key, subspace in env.observation_space.spaces.items():
    print(f"{key}: expected {subspace.shape}, got {obs[key].shape}")
    print(f"  dtype expected {subspace.dtype}, got {obs[key].dtype}")
    print(f"  contains: {subspace.contains(obs[key])}")

This tells you exactly which component fails.

Use explicit dtypes, return copies not views, validate obsessively during development, and write long-horizon tests. These five bugs account for probably 90% of the “observation space mismatch” errors I’ve seen in custom Gymnasium environments. The remaining 10% are usually typos.

What I still haven’t figured out is a clean pattern for dynamic observation spaces—environments where the observation shape changes based on the number of entities in the scene. Gymnasium technically supports this with unbounded Sequence spaces, but RL algorithms that assume fixed-size tensors choke on it. If you’ve found a good solution, I’d genuinely love to hear it.

Did you find this helpful?

Your support keeps this blog running and ad-free content coming.

☕ Buy me a coffee