marlbench

Setting up MARL environments can take a surprising amount of time. Each one has its own API, and using more than one environment in the same codebase usually means writing additional code to provide a common API.

marlbench does this work for you. It provides a common interface for several MARL environments, together with environment wrappers and vectorized environments.

This repository contains environment tools only. It does not include MARL algorithms.

How to use

Install it using uv pip install marlbench, or clone the repo.

Full examples are provided in CleanMARL

What is included?

Supported environments

The SMACv2 scenario configuration files are already included in marlbench/configs/smacv2/.

The supported MAgent2 environments are: adversarial_pursuit_v4,battle_v4, battlefield_v5, combined_arms_v6, gather_v5,tiger_deer_v4.

Common API

All interfaces provide the following methods:

obs, info = env.reset(seed=42)
obs, reward, terminated, truncated, info = env.step(actions)

state = env.get_state()
avail_actions = env.get_avail_actions()
agent_mask = env.get_agent_mask() # which agents are active

obs_size = env.get_obs_size()
state_size = env.get_state_size()
action_size = env.get_action_size()

actions = env.sample()
env.close()

Example

The following example runs one LBF episode:

from marlbench.interfaces.lbf import LBFInterface


env = LBFInterface(
    env_name="Foraging-2s-10x10-3p-3f-coop-v3",
    max_episode_steps=150,
    reward_aggr="sum",
    disable_env_checker=True,
)

obs, info = env.reset()
done = False

while not done:
    actions = env.sample()
    obs, reward, terminated, truncated, info = env.step(actions)

    state = env.get_state()
    avail_actions = env.get_avail_actions()
    agent_mask = env.get_agent_mask()

    done = terminated or truncated

env.close()

Other environments can be created in the same way:

from marlbench.interfaces.magent import MAgent2Interface
from marlbench.interfaces.mamujoco import MAmujocoInterface
from marlbench.interfaces.pz import PettingZooInterface
from marlbench.interfaces.rware import RWAREInterface
from marlbench.interfaces.smac import SMACInterface
from marlbench.interfaces.smaclite import SMACliteInterface
from marlbench.interfaces.smacv2 import SMACv2Interface


rware_env = RWAREInterface(
    env_name="rware-tiny-2ag-v2",
    reward_aggr="sum",
)

smaclite_env = SMACliteInterface(env_name="MMM2")

pettingzoo_env = PettingZooInterface(env_name="pursuit_v4",family="sisl")

mamujoco_env = MAmujocoInterface(env_name="Ant-p2x4")

magent_env = MAgent2Interface(env_name="adversarial_pursuit_v4")

smac_env = SMACInterface(env_name="3m")

smacv2_env = SMACv2Interface(env_name="terran_5_vs_5")

Wrappers

The following wrappers are available:

Wrappers can be combined:

import numpy as np

from marlbench.interfaces.lbf import LBFInterface
from marlbench.wrappers.obs_wrappers import AddAgentID, NormalizeObservation
from marlbench.wrappers.reward_wrappers import NormalizeReward, TransformReward
from marlbench.wrappers.common import TimeLimit


env = LBFInterface(
    env_name="Foraging-2s-10x10-3p-3f-coop-v3",
    reward_aggr="none",
)

env = TimeLimit(env, max_episode_steps=150)
env = NormalizeObservation(env, normalize_state=True)
env = AddAgentID(env)
env = TransformReward(env, lambda reward: np.clip(reward, -1.0, 1.0))
env = NormalizeReward(env, gamma=0.99)

The order matters. Observations are normalized before the agent IDs are added, so the IDs remain zero or one. Rewards are clipped before they are normalized.

Vector environments

Two vector environment classes are available:

Both classes receive a list of functions that create environments.

SyncVectorEnv

from marlbench.interfaces.lbf import LBFInterface
from marlbench.vec_envs.sync_vec import SyncVectorEnv
N_ENVS = 4
def make_lbf():
    return LBFInterface(
        env_name="Foraging-2s-10x10-3p-3f-coop-v3",
        max_episode_steps=150,
        reward_aggr="sum")
env_fns = [make_lbf for _ in range(N_ENVS)]
env = SyncVectorEnv(env_fns,auto_reset=False)
observations, infos = env.reset(seed=42)
observations, rewards, dones, truncated, infos = env.step(env.sample())
env.close()

SubprocVectorEnv

SubprocVectorEnv uses the same API:

from marlbench.vec_envs.subproc_vec import SubprocVectorEnv
N_ENVS = 4
def make_lbf():
    return LBFInterface(
        env_name="Foraging-2s-10x10-3p-3f-coop-v3",
        max_episode_steps=150,
        reward_aggr="sum")
if __name__ == "__main__":
    env_fns = [make_lbf for _ in range(N_ENVS)]
    env = SubprocVectorEnv(
        env_fns,
        start_method="spawn",
        auto_reset=False)
    observations, infos = env.reset(seed=42)
    observations, rewards, dones, truncated, infos = env.step(env.sample())
    env.close()

Every returned array has n_envs as its first dimension:

How to handle episodes with different lengths

Episodes do not end at the same time. With parallel environments running, one may terminate before the others. When auto_reset=False, an environment that returns done or truncated becomes inactive. It is not stepped again: it keeps returning its final observation with a reward of zero until you reset it. env.get_env_mask() tells you which environments are still running. You can use reset(indices=...) to reset a subset of environment. It accepts an integer, a list of integers, or a boolean mask of shape (n_envs,).

In constrast auto_reset=True automatically resets a finished environment inside step(), and the observation returned by step() then belongs to the new episode. The terminal data is moved into that environment's info dict. It can be accessed using final_obs,final_state, final_avail_actions, final_agent_mask,final_info. dones, truncated and rewards still describe the step that ended the episode, so episode statistics keep working.