Spaces

Overview

PufferLib’s space module provides compatibility layers for both Gym and Gymnasium, along with utilities for converting single-agent spaces to multi-agent joint spaces.

Space types

All space types are tuples containing both Gym and Gymnasium versions:

Discrete

import pufferlib.spaces

Discrete = (gym.spaces.Discrete, gymnasium.spaces.Discrete)

Discrete space for categorical actions or observations. Example:

import gymnasium
import pufferlib.spaces

action_space = gymnasium.spaces.Discrete(4)  # 4 discrete actions (0, 1, 2, 3)
isinstance(action_space, pufferlib.spaces.Discrete)  # True

Box

import pufferlib.spaces

Box = (gym.spaces.Box, gymnasium.spaces.Box)

Continuous or bounded integer space. Example:

import gymnasium
import numpy as np
import pufferlib.spaces

# Continuous action space
action_space = gymnasium.spaces.Box(
    low=-1.0, high=1.0, shape=(2,), dtype=np.float32
)

# Image observation space
obs_space = gymnasium.spaces.Box(
    low=0, high=255, shape=(84, 84, 3), dtype=np.uint8
)

isinstance(obs_space, pufferlib.spaces.Box)  # True

MultiDiscrete

import pufferlib.spaces

MultiDiscrete = (gym.spaces.MultiDiscrete, gymnasium.spaces.MultiDiscrete)

Multiple discrete actions. Example:

import gymnasium
import numpy as np
import pufferlib.spaces

# 3 discrete actions: first has 4 options, second has 3, third has 2
action_space = gymnasium.spaces.MultiDiscrete([4, 3, 2])
isinstance(action_space, pufferlib.spaces.MultiDiscrete)  # True

MultiBinary

import pufferlib.spaces

MultiBinary = (gym.spaces.MultiBinary, gymnasium.spaces.MultiBinary)

Multiple binary (0 or 1) values. Example:

import gymnasium
import pufferlib.spaces

action_space = gymnasium.spaces.MultiBinary(5)  # 5 binary actions
isinstance(action_space, pufferlib.spaces.MultiBinary)  # True

Tuple

import pufferlib.spaces

Tuple = (gym.spaces.Tuple, gymnasium.spaces.Tuple)

Tuple of multiple spaces. Example:

import gymnasium
import numpy as np
import pufferlib.spaces

space = gymnasium.spaces.Tuple((
    gymnasium.spaces.Discrete(3),
    gymnasium.spaces.Box(low=0, high=1, shape=(2,), dtype=np.float32),
))
isinstance(space, pufferlib.spaces.Tuple)  # True

Dict

import pufferlib.spaces

Dict = (gym.spaces.Dict, gymnasium.spaces.Dict)

Dictionary of multiple named spaces. Example:

import gymnasium
import numpy as np
import pufferlib.spaces

space = gymnasium.spaces.Dict({
    'image': gymnasium.spaces.Box(low=0, high=255, shape=(84, 84, 3), dtype=np.uint8),
    'position': gymnasium.spaces.Box(low=-np.inf, high=np.inf, shape=(2,), dtype=np.float32),
    'action_mask': gymnasium.spaces.MultiBinary(10),
})
isinstance(space, pufferlib.spaces.Dict)  # True

joint_space()

Convert a single-agent space to a multi-agent joint space.

import pufferlib.spaces

joint = pufferlib.spaces.joint_space(space, n)

space

gymnasium.spaces.Space

required

The single-agent space to convert. Must be Discrete, MultiDiscrete, or Box.

int

required

Number of agents.

joint_space

gymnasium.spaces.Space

The joint space for n agents.

Behavior by space type

Discrete
MultiDiscrete
Box

Converts Discrete(k) to MultiDiscrete([k] * n)

import gymnasium
import pufferlib.spaces

single = gymnasium.spaces.Discrete(4)
joint = pufferlib.spaces.joint_space(single, 8)
# Result: MultiDiscrete([4, 4, 4, 4, 4, 4, 4, 4])
print(joint.shape)  # (8,)

Converts MultiDiscrete([k1, k2, ...]) to Box with shape (n, len(space))

import gymnasium
import numpy as np
import pufferlib.spaces

single = gymnasium.spaces.MultiDiscrete([3, 4, 5])
joint = pufferlib.spaces.joint_space(single, 8)
# Result: Box(low=0, high=[[2,3,4], ...], shape=(8, 3))
print(joint.shape)  # (8, 3)

Converts Box(shape=(s1, s2, ...)) to Box(shape=(n, s1, s2, ...))

import gymnasium
import numpy as np
import pufferlib.spaces

single = gymnasium.spaces.Box(
    low=0, high=255, shape=(84, 84, 3), dtype=np.uint8
)
joint = pufferlib.spaces.joint_space(single, 8)
# Result: Box(low=0, high=255, shape=(8, 84, 84, 3))
print(joint.shape)  # (8, 84, 84, 3)

Usage examples

Creating environments with joint spaces

import numpy as np
import gymnasium
import pufferlib
import pufferlib.spaces

class MultiAgentEnv(pufferlib.PufferEnv):
    def __init__(self, buf=None):
        # Define single-agent spaces
        self.single_observation_space = gymnasium.spaces.Box(
            low=0, high=1, shape=(10,), dtype=np.float32
        )
        self.single_action_space = gymnasium.spaces.Discrete(4)
        self.num_agents = 8
        
        super().__init__(buf)
        
        # After super().__init__(), joint spaces are automatically created:
        # self.observation_space = joint_space(single_observation_space, num_agents)
        # self.action_space = joint_space(single_action_space, num_agents)
        
        print(self.observation_space.shape)  # (8, 10)
        print(self.action_space.shape)        # (8,) - MultiDiscrete([4]*8)

Type checking

import gymnasium
import pufferlib.spaces

space = gymnasium.spaces.Discrete(4)

# Check if space is a Discrete space (works with both Gym and Gymnasium)
if isinstance(space, pufferlib.spaces.Discrete):
    print(f"Discrete space with {space.n} actions")

# Check if space is a Box space
box_space = gymnasium.spaces.Box(low=0, high=1, shape=(5,))
if isinstance(box_space, pufferlib.spaces.Box):
    print(f"Box space with shape {box_space.shape}")

Converting spaces manually

import numpy as np
import gymnasium
import pufferlib.spaces

# Single agent with discrete actions
single_action = gymnasium.spaces.Discrete(6)

# Create joint space for 16 agents
joint_action = pufferlib.spaces.joint_space(single_action, 16)
print(joint_action)  # MultiDiscrete([6 6 6 6 6 6 6 6 6 6 6 6 6 6 6 6])

# Sample actions for all agents
actions = joint_action.sample()
print(actions.shape)  # (16,)
print(actions)  # e.g., [3 1 5 2 0 4 1 3 2 5 0 1 4 3 2 1]

# Single agent with box observations
single_obs = gymnasium.spaces.Box(
    low=-np.inf, high=np.inf, shape=(3, 3), dtype=np.float32
)

# Create joint space for 16 agents
joint_obs = pufferlib.spaces.joint_space(single_obs, 16)
print(joint_obs.shape)  # (16, 3, 3)

# Sample observations for all agents
observations = joint_obs.sample()
print(observations.shape)  # (16, 3, 3)

Working with complex action spaces

import numpy as np
import gymnasium
import pufferlib.spaces

# MultiDiscrete: e.g., [move_direction, attack_target, item_to_use]
single_action = gymnasium.spaces.MultiDiscrete([4, 10, 5])

# Joint space for 32 agents
joint_action = pufferlib.spaces.joint_space(single_action, 32)
print(joint_action.shape)  # (32, 3)
print(type(joint_action))  # Box (MultiDiscrete becomes Box in joint space)

# Sample actions
actions = joint_action.sample()
print(actions.shape)  # (32, 3)
print(actions[0])     # e.g., [2, 7, 3] - first agent's action

Implementation details

The joint_space() function is used internally by PufferEnv to automatically create observation_space and action_space from single_observation_space and single_action_space.

Currently, only Discrete, MultiDiscrete, and Box spaces are supported by joint_space(). Other space types (Tuple, Dict, MultiBinary) will raise a ValueError.

Common patterns

PufferEnv with discrete actions

import gymnasium
import pufferlib
import pufferlib.spaces

class MyEnv(pufferlib.PufferEnv):
    def __init__(self, buf=None):
        self.single_observation_space = gymnasium.spaces.Box(
            low=0, high=255, shape=(84, 84, 3), dtype=np.uint8
        )
        self.single_action_space = gymnasium.spaces.Discrete(18)  # 18 actions
        self.num_agents = 16
        super().__init__(buf)
        # self.action_space is now MultiDiscrete([18]*16)
        # self.observation_space is now Box(shape=(16, 84, 84, 3))

PufferEnv with continuous actions

import numpy as np
import gymnasium
import pufferlib
import pufferlib.spaces

class ContinuousEnv(pufferlib.PufferEnv):
    def __init__(self, buf=None):
        self.single_observation_space = gymnasium.spaces.Box(
            low=-np.inf, high=np.inf, shape=(20,), dtype=np.float32
        )
        self.single_action_space = gymnasium.spaces.Box(
            low=-1, high=1, shape=(4,), dtype=np.float32
        )
        self.num_agents = 8
        super().__init__(buf)
        # self.action_space is now Box(shape=(8, 4))
        # self.observation_space is now Box(shape=(8, 20))

Core API

Training

Emulation

Utilities

Overview

Space types

Discrete

Box

MultiDiscrete

MultiBinary

Tuple

Dict

joint_space()

Behavior by space type

Usage examples

Creating environments with joint spaces

Type checking

Converting spaces manually

Working with complex action spaces

Implementation details

Common patterns

PufferEnv with discrete actions

PufferEnv with continuous actions

Build docs developers (and LLMs) love

Core API

Training

Emulation

Utilities

​Overview

​Space types

​Discrete

​Box

​MultiDiscrete

​MultiBinary

​Tuple

​Dict

​joint_space()

​Behavior by space type

​Usage examples

​Creating environments with joint spaces

​Type checking

​Converting spaces manually

​Working with complex action spaces

​Implementation details

​Common patterns

​PufferEnv with discrete actions

​PufferEnv with continuous actions

Build docs developers (and LLMs) love

Overview

Space types

Discrete

Box

MultiDiscrete

MultiBinary

Tuple

Dict

joint_space()

Behavior by space type

Usage examples

Creating environments with joint spaces

Type checking

Converting spaces manually

Working with complex action spaces

Implementation details

Common patterns

PufferEnv with discrete actions

PufferEnv with continuous actions