Python & Data Science
Statistics Under review

A Tour of OpenAI Gym / Gymnasium: Environments You Can Actually Train On

You’ve done it. You built a Deep Q-Network from scratch in Part 5. You watched it learn to balance a pole on a cart. It felt like magic — and it was. But here’s the thing: that CartPole environment was just a toy. It came from a library you imported without thinking much about it. What happens when you want to train your DQN on something harder? Something with pixels? Something with continuous control?

That’s what this article is about. You’re about to meet the tool that powers almost every reinforcement learning experiment in the world: OpenAI Gym (and its modern, maintained fork Gymnasium).

By the end of this tour, you’ll know:

  • The four functions that define every Gymnasium environment
  • The full “environment zoo” — from easy (CartPole) to hard (MuJoCo)
  • How to plug your DQN from Part 5 into any environment and watch it learn
  • Why Gymnasium exists and how to migrate old code

Let’s get started.

1. The Gymnasium API: Four Functions You Need to Know

Think of a Gymnasium environment like a video game console. You plug in a game (that’s make()), press start (reset()), press buttons (step()), and see the screen (render()). That’s it. Four functions. Everything else is just details.

Here’s the core API:

FunctionWhat it doesWhat it returns
make()Creates the environmentAn environment object
reset()Starts a new episodeInitial observation
step(action)Takes one actionobservation, reward, terminated, truncated, info
close()Shuts down the environmentNothing

Now here’s the catch — and this is the hardest part for beginners: the step() function returns five things, not four. The old Gym API returned done as a single boolean. Gymnasium splits it into terminated (the episode ended naturally, like winning or dying) and truncated (the episode was cut short, like hitting a time limit).

Let’s see this in action with the classic CartPole environment.

import gymnasium as gym

# Create the environment with rendering enabled
env = gym.make('CartPole-v1', render_mode='human')

# Start a new episode
observation, info = env.reset()
print(f"Initial observation: {observation}")
print(f"  This means: cart position = {observation[0]:.2f}, cart velocity = {observation[1]:.2f},")
print(f"  pole angle = {observation[2]:.2f}, pole angular velocity = {observation[3]:.2f}")
print()

# Run 200 random steps
for step in range(200):
    # Choose a random action (0 = push left, 1 = push right)
    action = env.action_space.sample()
    
    # Take the action
    observation, reward, terminated, truncated, info = env.step(action)
    
    # Print what happened
    if step < 5 or step % 50 == 0:
        print(f"Step {step}: action={action}, reward={reward:.1f}, terminated={terminated}, truncated={truncated}")
        print(f"  Observation: cart={observation[0]:.2f}, velocity={observation[1]:.2f}, angle={observation[2]:.2f}, angular_vel={observation[3]:.2f}")
    
    # End the episode if the pole fell or we hit the time limit
    if terminated or truncated:
        print(f"Episode ended at step {step}: {'pole fell' if terminated else 'time limit reached'}")
        break

# Clean up
env.close()

When you run this, you’ll see a window pop up with a cart and a pole. The cart moves left and right randomly. The pole wobbles and eventually falls. Each step gives a reward of +1 for keeping the pole balanced. The episode ends when the pole tilts too far (terminated) or you reach 500 steps (truncated).

The observation vector has four numbers. In plain English: where the cart is, how fast it’s moving, what angle the pole is at, and how fast it’s rotating. The action space is two choices: push left or push right. That’s all the information your DQN gets.

2. The Environment Zoo: What’s Actually Available

Environments in Gymnasium are like difficulty levels in a video game. You don’t start on ‘Nightmare’ mode. Let’s tour the tiers from easiest to hardest.

Tier 1: Classic Control

These are the “hello world” environments. Low-dimensional state, fast to train, perfect for debugging.

EnvironmentObservation SpaceAction SpaceWhat the agent learns
CartPole-v14 continuous2 discreteBalance a pole on a cart
MountainCar-v02 continuous3 discreteDrive a car up a hill
Pendulum-v13 continuous1 continuousSwing a pendulum upright
Acrobot-v16 continuous3 discreteSwing a two-link robot up

Tier 2: Toy Text

Discrete state and action spaces. These are great for tabular methods and debugging.

EnvironmentObservation SpaceAction SpaceWhat the agent learns
Blackjack-v13 discrete2 discretePlay blackjack optimally
Taxi-v3500 discrete6 discreteNavigate a taxi to pick up and drop off passengers
FrozenLake-v116 discrete4 discreteWalk across a frozen lake without falling

Tier 3: Box2D

Physics simulation with moderate complexity. Shaped rewards make learning easier.

EnvironmentObservation SpaceAction SpaceWhat the agent learns
LunarLander-v28 continuous4 discreteLand a spacecraft on a landing pad
BipedalWalker-v324 continuous4 continuousMake a two-legged robot walk

Tier 4: Atari

Pixel observations. High-dimensional. Your agent needs a convolutional neural network.

EnvironmentObservation SpaceAction SpaceWhat the agent learns
BreakoutNoFrameskip-v4210x160x3 pixels4 discreteBreak bricks with a ball
PongNoFrameskip-v4210x160x3 pixels6 discretePlay Pong against the computer
SpaceInvadersNoFrameskip-v4210x160x3 pixels6 discreteShoot aliens

Tier 5: MuJoCo

Continuous control with high degrees of freedom. The hardest tier.

EnvironmentObservation SpaceAction SpaceWhat the agent learns
HalfCheetah-v417 continuous6 continuousMake a cheetah run
Ant-v427 continuous8 continuousMake a four-legged ant walk
Humanoid-v4376 continuous17 continuousMake a humanoid stand and walk

Let’s write code that creates each tier and prints what the agent actually sees.

import gymnasium as gym

# Tier 1: Classic Control
print("=== Tier 1: Classic Control ===")
cartpole = gym.make('CartPole-v1')
obs, _ = cartpole.reset()
print(f"CartPole: obs shape={obs.shape}, action space={cartpole.action_space}")
print(f"  Agent sees {obs.shape[0]} continuous numbers, can choose 1 of {cartpole.action_space.n} actions")
cartpole.close()

# Tier 2: Toy Text
print("\n=== Tier 2: Toy Text ===")
taxi = gym.make('Taxi-v3')
obs, _ = taxi.reset()
print(f"Taxi: obs shape=scalar ({obs}), action space={taxi.action_space}")
print(f"  Agent sees a single integer (0-499), can choose 1 of {taxi.action_space.n} actions")
taxi.close()

# Tier 3: Box2D
print("\n=== Tier 3: Box2D ===")
lunar = gym.make('LunarLander-v2')
obs, _ = lunar.reset()
print(f"LunarLander: obs shape={obs.shape}, action space={lunar.action_space}")
print(f"  Agent sees {obs.shape[0]} continuous numbers (position, velocity, angle, leg contact)")
print(f"  Can choose 1 of {lunar.action_space.n} actions (do nothing, fire left, fire main, fire right)")
lunar.close()

# Tier 4: Atari
print("\n=== Tier 4: Atari ===")
breakout = gym.make('BreakoutNoFrameskip-v4', render_mode='rgb_array')
obs, _ = breakout.reset()
print(f"Breakout: obs shape={obs.shape}, action space={breakout.action_space}")
print(f"  Agent sees a {obs.shape[0]}x{obs.shape[1]} pixel image with {obs.shape[2]} color channels")
print(f"  This means the agent needs a convolutional neural network to process the image")
breakout.close()

# Tier 5: MuJoCo
print("\n=== Tier 5: MuJoCo ===")
cheetah = gym.make('HalfCheetah-v4')
obs, _ = cheetah.reset()
print(f"HalfCheetah: obs shape={obs.shape}, action space={cheetah.action_space}")
print(f"  Agent sees {obs.shape[0]} continuous numbers (joint positions, velocities, etc.)")
print(f"  Can output {cheetah.action_space.shape[0]} continuous actions (torques at each joint)")
cheetah.close()

Here’s my recommendation: start with CartPole. When you can solve that consistently, move to LunarLander. Then, if you’re feeling brave, tackle Atari. MuJoCo is for when you’re ready for the big leagues.

3. Actually Training: Plugging a DQN Into a Gymnasium Environment

Now for the fun part. You already have a DQN from Part 5. Let’s wire it to a real Gymnasium environment and watch it learn. The code is surprisingly short.

We’ll use CartPole-v1 because it’s fast and familiar. The training loop follows the same pattern every time:

  1. Reset the environment to get the initial observation
  2. Loop: select an action (epsilon-greedy), take a step, store the experience in the replay buffer, sample a batch, update the DQN
  3. Repeat until the episode ends
  4. Track the reward per episode

Let’s build it.

import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
import random
import matplotlib.pyplot as plt

# --- DQN Network (from Part 5) ---
class DQN(nn.Module):
    """Simple feedforward network for DQN."""
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )
    
    def forward(self, x):
        return self.net(x)

# --- Replay Buffer (from Part 5) ---
class ReplayBuffer:
    """Stores and samples experiences."""
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, terminated):
        self.buffer.append((state, action, reward, next_state, terminated))
    
    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, terminateds = zip(*batch)
        return (
            torch.FloatTensor(np.array(states)),
            torch.LongTensor(np.array(actions)).unsqueeze(1),
            torch.FloatTensor(np.array(rewards)).unsqueeze(1),
            torch.FloatTensor(np.array(next_states)),
            torch.FloatTensor(np.array(terminateds)).unsqueeze(1)
        )
    
    def __len__(self):
        return len(self.buffer)

# --- Training Setup ---
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

policy_net = DQN(state_dim, action_dim)
target_net = DQN(state_dim, action_dim)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()

optimizer = optim.Adam(policy_net.parameters(), lr=0.001)
memory = ReplayBuffer(capacity=10000)

# Hyperparameters
batch_size = 64
gamma = 0.99
epsilon_start = 1.0
epsilon_end = 0.01
epsilon_decay = 0.995
episodes = 500
target_update = 10

# --- Training Loop ---
epsilon = epsilon_start
episode_rewards = []

for episode in range(episodes):
    state, _ = env.reset()
    total_reward = 0
    terminated = False
    truncated = False
    
    while not (terminated or truncated):
        # Epsilon-greedy action selection
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            with torch.no_grad():
                q_values = policy_net(torch.FloatTensor(state).unsqueeze(0))
                action = q_values.argmax().item()
        
        # Take action
        next_state, reward, terminated, truncated, _ = env.step(action)
        total_reward += reward
        
        # Store experience
        memory.push(state, action, reward, next_state, terminated)
        state = next_state
        
        # Train if we have enough experiences
        if len(memory) >= batch_size:
            states, actions, rewards, next_states, terminateds = memory.sample(batch_size)
            
            # Compute target Q-values
            with torch.no_grad():
                next_q_values = target_net(next_states)
                max_next_q = next_q_values.max(1, keepdim=True)[0]
                target_q = rewards + gamma * max_next_q * (1 - terminateds)
            
            # Compute current Q-values
            current_q = policy_net(states).gather(1, actions)
            
            # Loss and update
            loss = nn.MSELoss()(current_q, target_q)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    # Decay epsilon
    epsilon = max(epsilon_end, epsilon * epsilon_decay)
    
    # Update target network
    if episode % target_update == 0:
        target_net.load_state_dict(policy_net.state_dict())
    
    episode_rewards.append(total_reward)
    
    if (episode + 1) % 50 == 0:
        avg_reward = np.mean(episode_rewards[-50:])
        print(f"Episode {episode+1}, Avg Reward (last 50): {avg_reward:.1f}, Epsilon: {epsilon:.3f}")

env.close()

# --- Plot the Training Curve ---
plt.figure(figsize=(10, 6))
plt.plot(episode_rewards, alpha=0.6, label='Episode Reward')
# Moving average
window = 50
if len(episode_rewards) >= window:
    moving_avg = np.convolve(episode_rewards, np.ones(window)/window, mode='valid')
    plt.plot(range(window-1, len(episode_rewards)), moving_avg, 'r-', linewidth=2, label=f'Moving Avg ({window} eps)')
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('DQN Training on CartPole-v1')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

When you run this, you’ll see the reward per episode start around 20 (random behavior) and climb to 500 (the maximum). The moving average line shows the trend clearly. The agent goes from falling after a few seconds to balancing the pole for the full 500 steps.

Here’s the hardest part: hyperparameter tuning. If your agent doesn’t learn, check these first:

  • Learning rate: 0.001 is a good start. Too high and the network diverges. Too low and it learns slowly.
  • Replay buffer size: 10,000 is fine for CartPole. For harder environments, you might need 100,000 or more.
  • Epsilon decay: If epsilon decays too fast, the agent stops exploring before it learns. If it decays too slowly, the agent wastes time on random actions.

4. Beyond the Built-Ins: Custom Environments and Research Use

The Gymnasium API isn’t just for Atari games. It’s a standard interface that researchers use for everything from robotics to theorem proving. Any environment that implements reset(), step(), and render() is compatible.

Here’s what that means in practice:

  • Controlgym: A library that turns industrial control systems (chemical plants, power grids) into Gymnasium environments. You can train an RL agent to control a real chemical reactor using the same code you used for CartPole.
  • gym-saturation: An environment for automated theorem proving. Your agent learns to prove mathematical theorems by choosing which inference rules to apply.

Let’s see how easy it is to create a custom environment. Here’s a minimal skeleton:

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

class SimpleGridWorld(gym.Env):
    """A minimal grid world: navigate from start to goal."""
    
    def __init__(self, grid_size=5):
        super().__init__()
        self.grid_size = grid_size
        
        # Observation: agent's (x, y) position
        self.observation_space = spaces.Box(
            low=0, high=grid_size-1, shape=(2,), dtype=np.int32
        )
        
        # Actions: 0=up, 1=down, 2=left, 3=right
        self.action_space = spaces.Discrete(4)
        
        # Goal position
        self.goal = np.array([grid_size-1, grid_size-1])
    
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # Start at top-left corner
        self.agent_pos = np.array([0, 0])
        return self.agent_pos, {}
    
    def step(self, action):
        # Move agent
        if action == 0:  # up
            self.agent_pos[1] = min(self.agent_pos[1] + 1, self.grid_size - 1)
        elif action == 1:  # down
            self.agent_pos[1] = max(self.agent_pos[1] - 1, 0)
        elif action == 2:  # left
            self.agent_pos[0] = max(self.agent_pos[0] - 1, 0)
        elif action == 3:  # right
            self.agent_pos[0] = min(self.agent_pos[0] + 1, self.grid_size - 1)
        
        # Check if we reached the goal
        terminated = np.array_equal(self.agent_pos, self.goal)
        reward = 1.0 if terminated else -0.01  # Small penalty per step
        truncated = False
        
        return self.agent_pos, reward, terminated, truncated, {}
    
    def render(self):
        # Simple text rendering
        grid = np.full((self.grid_size, self.grid_size), '.')
        grid[self.agent_pos[1], self.agent_pos[0]] = 'A'
        grid[self.goal[1], self.goal[0]] = 'G'
        print('\n'.join([' '.join(row) for row in grid]))
        print()

That’s it. You’ve created a Gymnasium environment. You can train any RL algorithm on it. The API’s power is in its simplicity: one interface, infinite environments.

5. Why Gymnasium (Not the Old Gym)? A Quick Migration Guide

You’ll find old tutorials that use import gym and done. Those won’t work with the current version. Here’s what changed and how to adapt.

The old Gym is archived. No more updates, no bug fixes. Gymnasium is the maintained fork by the Farama Foundation. The migration is mechanical.

Here’s the side-by-side comparison:

# OLD Gym code (will not work with current versions)
import gym

env = gym.make('CartPole-v1')
observation = env.reset()
done = False
while not done:
    action = env.action_space.sample()
    observation, reward, done, info = env.step(action)
    env.render()

# NEW Gymnasium code (current)
import gymnasium as gym

env = gym.make('CartPole-v1', render_mode='human')
observation, info = env.reset()
terminated = False
truncated = False
while not (terminated or truncated):
    action = env.action_space.sample()
    observation, reward, terminated, truncated, info = env.step(action)

The key changes:

  1. import gym → import gymnasium as gym
  2. done → terminated and truncated
  3. env.reset() → observation, info = env.reset()
  4. env.render() → render_mode='human' in make()
  5. env.seed(seed) → env.reset(seed=seed)

If you find an old tutorial, the migration is mechanical. Just add the truncated flag and change the render call. You’ll be up and running in minutes.

Closing: What You Can Do Now

You now have the tool to train any RL algorithm on any standard environment. Let’s recap what you learned:

  • Gymnasium API: Four functions — make(), reset(), step(), close() — define every environment
  • Environment zoo: From easy (CartPole) to hard (MuJoCo), with Classic Control, Toy Text, Box2D, Atari, and MuJoCo tiers
  • Training a DQN: The same code from Part 5 works on any Gymnasium environment with discrete actions
  • Custom environments: The API is just a protocol. Any class that implements reset(), step(), and render() is compatible
  • Migration: Gymnasium is the maintained fork. The changes are mechanical: done → terminated + truncated, and new render API

Next up: Policy Gradients and why they beat DQN on continuous control. You’ll take everything you learned here and apply it to environments like HalfCheetah and LunarLander with continuous action spaces.

Check Your Understanding

Remember: Name the four core Gymnasium functions and what each returns.

Understand: Explain in plain English why the step() function returns both terminated and truncated instead of just done.

Apply: Write a minimal loop that runs 100 random steps on LunarLander-v2 and prints the average reward.

Analyze: Compare the observation spaces of CartPole (4 continuous) and Breakout (210×160×3 pixels). What does this mean for the neural network architecture you’d use?

Evaluate: Given a new environment with a 100-dimensional continuous observation space and 5 discrete actions, which tier from this article does it belong to and why?

Create: Design a custom Gymnasium environment for a simple grid-world with obstacles. Sketch the reset(), step(), and render() methods.

  • Part 5: Deep Q-Networks — The DQN architecture you’ll train on these environments
  • Part 4: Q-Learning From Scratch — The tabular method that inspired DQN
  • Part 1: What Is Reinforcement Learning, Really? — The foundational concepts

Apply What You Learned is for Supporter and Insider subscribers.

Subscribe to unlock the exercises on this post.

See plans
  • Statistics Under review

    The 'Just Run It' Trap

    Learn how to calculate A/B test sample size in Python with statsmodels, avoid the peeking problem, and balance MDE, alpha, and power before you launch.

  • Statistics Under review

    A Self Evaluation Rubric For Code Design Five Dime

    Let's get the definition out of the way in plain English: a rubric is just a list of things to look for, with clear descriptions of what "bad," "okay," and "great" look like. It's not a pass/fail judgment.

  • Statistics Under review

    Building A Simple Research Or Trading Agent As A L

    You've done the hard work. In Part 7, you built a working agent loop: LLM call → tool execution → append result → repeat. It runs. It calls tools. It responds.

  • Statistics Under review

    Think Like a Bayesian: A Guide for Frequentists Who Hate Formulas

    Learn how Bayesian priors, posteriors, and base rates produce clearer, more actionable insights than Frequentist methods — with intuitive Python examples.

Looking for something else?

Search every article by title, summary or topic.