Basic Usage¶

This page shows the minimal way to use MASA without masa.common.utils.make_env, by manually constructing a Gymnasium environment and wrapping it in the recommended order:

TimeLimit \(\rightarrow\) LabelledEnv \(\rightarrow\) BaseConstraintEnv \(\rightarrow\) ConstraintMonitor \(\rightarrow\) RewardMonitor

This is the same order enforced by make_env (notably, TimeLimit must come first).

Overview¶

MASA components reason over labels (atomic predicates) derived from observations. The wrapper masa.common.labelled_env.LabelledEnv computes these labels on every gymnasium.Env.reset and gymnasium.Env.step and stores them in info["labels"].

Constraints then consume these labels and expose consistent metrics, while the monitor wrappers attach step/episode summaries to the info dictionary for logging and debugging.

Minimal environment construction¶

import gymnasium as gym

# Core MASA wrappers
from masa.common.wrappers import TimeLimit, ConstraintMonitor, RewardMonitor
from masa.common.labelled_env import LabelledEnv

# Simple Media Streaming environment
from masa.env.tabular.media_streaming import MediaStreaming

# Example constraint wrapper (a BaseConstraintEnv implementation)
from masa.common.constraints.cmdp import CumulativeCostEnv

# --- 1) Define label and cost functions ---

def label_fn(obs):
    """
    Example labelling function for MediaStreaming-like observations.

    Returns:
        set[str]: Atomic predicates holding in the current observation.
    """
    labels = set()
    # These keys are illustrative; adapt to your observation structure.
    try:
        if int(obs) == 0:
            labels.add("unsafe")
    except:
        return set()
    return labels

def cost_fn(labels):
    """
    Example 0/1 cost: unsafe if the 'unsafe' predicate holds.
    """
    return 1.0 if "unsafe" in labels else 0.0

# --- 1.5) Or use default label_fn and cost_fn supplied by the environment (recommended)
from masa.env.tabular.media_streaming import label_fn, cost_fn

# --- 2) Build the environment and wrap in the correct order ---

env = MediaStreaming()

# Recommended: apply TimeLimit first (episode length is enforced before anything else).
env = TimeLimit(env, max_episode_steps=1_000)

# Attach labels to info["labels"] at every reset/step.
env = LabelledEnv(env, label_fn)

# Apply a constraint wrapper (example: cumulative cost/budget style constraint).
# Typical kwargs are shown; consult the constraint's docstring / Constraints API reference.
env = CumulativeConstraintEnv(env, cost_fn=cost_fn, budget=25.0)

# Finally, attach monitoring wrappers for constraints and reward logging.
env = ConstraintMonitor(env)
env = RewardMonitor(env)

Random-agent interaction loop (Gymnasium-style)¶

import numpy as np

num_episodes = 3

for ep in range(num_episodes):
    obs, info = env.reset(seed=ep)

    ep_return = 0.0
    ep_len = 0

    # Your monitors/constraints may attach additional fields; labels are always in info["labels"]
    labels = info.get("labels", set())
    print(f"[episode {ep}] reset labels={labels}")

    terminated = truncated = False
    while not (terminated or truncated):
        action = env.action_space.sample()

        obs, reward, terminated, truncated, info = env.step(action)
        ep_return += float(reward)
        ep_len += 1

        labels = info.get("labels", set())

        # Common pattern: monitors expose step metrics in info (names may vary by constraint).
        constraint = info.get("constraint", {})
        if isinstance(constraint, dict) and "step" in constraint:
            step_cost = constraint["step"].get("cost", 0.0)
            violated = constraint["step"].get("violated", False)

        if violated:
            print(f"  step={ep_len:04d} VIOLATION labels={labels} cost={step_cost}")

    # Episode-end metrics are often attached on the final transition by the monitors.
    # Again, keys vary; print what you care about.
    constraint = info.get("constraint", {})
    if isinstance(constraint, dict) and "episode" in constraint:
        ep_cost = constraint["episode"].get("cum_cost", None)
        ep_satisfied = constraint["episode"].get("satisfied", None)

    print(
        f"[episode {ep}] return={ep_return:.2f} len={ep_len} "
        f"episode_cost={ep_cost} episode_satisfied={ep_satisfied}"
   )

Training with PPO¶

Below is a minimal example showing how to initialize and train PPO (provided by MASA) using the wrapped environment. The specific PPO constructor and train API may include additional options (e.g., logging, eval env, saving); the snippet mirrors the general style used in MASA runs.

from masa.algorithms.on_policy import PPO

# (Optional) create a separate evaluation environment with the same wrapper stack.
def make_eval_env():
    eval_env = MediaStreaming()
    eval_env = TimeLimit(eval_env, max_episode_steps=1_000)
    eval_env = LabelledEnv(eval_env, label_fn)
    eval_env = CumulativeConstraintEnv(eval_env, cost_fn=cost_fn, budget=25.0)
    eval_env = ConstraintMonitor(eval_env)
    eval_env = RewardMonitor(eval_env)
    return eval_env

eval_env = make_eval_env()

# Initialize PPO.
# Common kwargs (device, seed, logging) follow the same pattern as other MASA algorithms.
algo = PPO(
    env,
    seed=0,
    device="auto",
    verbose=1,
    eval_env=eval_env,            # optional
    tensorboard_logdir=None,      # optional
)

# Train PPO. MASA algorithms automatically support eval/log frequencies and windowed stats.
algo.train(
    total_timesteps=200_000,
    num_eval_episodes=10,         # optional
    eval_freq=10_000,             # optional
    log_freq=2_000,               # optional
    stats_window_size=100,        # optional
)

Recording Videos¶

The central environment factories can optionally wrap the completed environment stack with video recording. Recording is disabled by default and requires an image-producing render mode such as "rgb_array".

from masa.common.utils import make_env
from masa.envs.discrete.conveyor_belt import cost_fn, label_fn

def record_every_episode(episode_id):
    return True

env = make_env(
    "conveyor_belt",
    "cmdp",
    100,
    label_fn=label_fn,
    cost_fn=cost_fn,
    budget=10.0,
    env_kwargs={"render_mode": "rgb_array"},
    record_video=True,
    record_video_episode_trigger=record_every_episode,
    video_folder="videos",
)

For renderable PettingZoo parallel environments, make_marl_env applies MASA's PettingZoo video wrapper in the same outermost position:

from masa.common.constraints.multi_agent.cmg import Budget
from masa.common.utils import make_marl_env

def record_every_episode(episode_id):
    return True

env = make_marl_env(
    "renderable_marl_env",
    "cmg",
    budgets=[Budget(amount=10.0, agents=("player_0", "player_1"), name="shared")],
    env_kwargs={"render_mode": "rgb_array"},
    record_video=True,
    record_video_episode_trigger=record_every_episode,
    video_folder="videos",
)

The PettingZoo wrapper supports renderable parallel environments and plugins whose render() returns RGB numpy arrays. The built-in matrix games expose the central wrapper path, but do not yet include concrete image renderers.

API Reference for environment factories¶

masa.common.utils.make_env ¶

make_env(env_id: str, constraint: str, max_episode_steps: int, *, label_fn: Optional[LabelFn] = None, constraint_kwargs: Optional[dict[str, Any]] = None, env_kwargs: Optional[dict[str, Any]] = None, record_video: bool = False, record_video_episode_trigger: Optional[Callable[[int], bool]] = None, video_folder: str = 'videos', video_kwargs: Optional[dict[str, Any]] = None, **kw) -> gym.Env

Construct a fully wrapped MASA environment using the canonical wrapper order.

This helper creates a Gymnasium environment and applies MASA wrappers in the recommended and enforced order:

TimeLimit \(\rightarrow\) LabelledEnv \(\rightarrow\) BaseConstraintEnv \(\rightarrow\) ConstraintMonitor \(\rightarrow\) RewardMonitor

The resulting environment exposes labels, constraint metrics, and reward summaries exclusively via the Gymnasium info dictionary. Observations and rewards themselves are left unchanged.

Parameters:

Name	Type	Description	Default
`env_id`	`str`	Environment identifier registered in `ENV_REGISTRY`.	required
`constraint`	`str`	Constraint identifier registered in `CONSTRAINT_REGISTRY`.	required
`max_episode_steps`	`int`	Maximum number of steps per episode. Applied via `TimeLimit` as the outermost wrapper.	required
`label_fn`	`Optional[LabelFn]`	Optional function mapping observations to atomic predicate labels. If provided, labels are computed on every `reset` and `step` and stored under `info["labels"]`. If omitted, the base environment's `label_fn` attribute is used.	`None`
`constraint_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to the constraint wrapper constructor. If `cost_fn` is omitted and the base environment exposes one, it is forwarded automatically.	`None`
`env_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to the base environment constructor.	`None`
`record_video`	`bool`	Whether to wrap the resulting environment with Gymnasium's `RecordVideo`. Defaults to `False`.	`False`
`record_video_episode_trigger`	`Optional[Callable[[int], bool]]`	Optional predicate called with the episode id to decide whether to record that episode. This is forwarded as `episode_trigger` to `RecordVideo`.	`None`
`video_folder`	`str`	Output directory for recorded videos when `record_video=True`.	`'videos'`
`video_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to `RecordVideo`.	`None`

Returns:

Type	Description
`Env`	A fully wrapped Gymnasium environment compatible with MASA algorithms,
`Env`	monitors, and logging utilities.

Notes

Wrapper order is fixed and enforced.
Constraints are reset automatically on environment reset.
All semantic metadata (labels, costs, violations, metrics) is communicated via the info dictionary.

See Also

masa.common.labelled_env.LabelledEnv
masa.common.constraints.base.BaseConstraintEnv
masa.common.wrappers.ConstraintMonitor
masa.common.wrappers.RewardMonitor

Source code in masa/common/utils.py

def make_env(
    env_id: str, 
    constraint: str, 
    max_episode_steps: int, 
    *,
    label_fn: Optional[LabelFn] = None, 
    constraint_kwargs: Optional[dict[str, Any]] = None,
    env_kwargs: Optional[dict[str, Any]] = None,
    record_video: bool = False,
    record_video_episode_trigger: Optional[Callable[[int], bool]] = None,
    video_folder: str = "videos",
    video_kwargs: Optional[dict[str, Any]] = None,
    **kw
) -> gym.Env:
    r"""
    Construct a fully wrapped MASA environment using the canonical wrapper order.

    This helper creates a Gymnasium environment and applies MASA wrappers in the
    **recommended and enforced order**:

    :class:`~gymnasium.wrappers.TimeLimit` :math:`\rightarrow` 
    :class:`~masa.common.labelled_env.LabelledEnv` :math:`\rightarrow` 
    :class:`~masa.common.constraints.base.BaseConstraintEnv`  :math:`\rightarrow` 
    :class:`~masa.common.wrappers.ConstraintMonitor` :math:`\rightarrow` 
    :class:`~masa.common.wrappers.RewardMonitor`

    The resulting environment exposes labels, constraint metrics, and reward
    summaries exclusively via the Gymnasium ``info`` dictionary. Observations
    and rewards themselves are left unchanged.

    Args:
        env_id:
            Environment identifier registered in ``ENV_REGISTRY``.
        constraint:
            Constraint identifier registered in ``CONSTRAINT_REGISTRY``.
        max_episode_steps:
            Maximum number of steps per episode. Applied via ``TimeLimit`` as
            the outermost wrapper.
        label_fn:
            Optional function mapping observations to atomic predicate labels.
            If provided, labels are computed on every ``reset`` and ``step`` and
            stored under ``info["labels"]``. If omitted, the base environment's 
            ``label_fn`` attribute is used.
        constraint_kwargs:
           Optional keyword arguments forwarded to the constraint wrapper
           constructor. If ``cost_fn`` is omitted and the base environment
           exposes one, it is forwarded automatically.
        env_kwargs:
            Optional keyword arguments forwarded to the base environment
            constructor.
        record_video:
            Whether to wrap the resulting environment with Gymnasium's
            :class:`~gymnasium.wrappers.RecordVideo`. Defaults to ``False``.
        record_video_episode_trigger:
            Optional predicate called with the episode id to decide whether to
            record that episode. This is forwarded as ``episode_trigger`` to
            :class:`~gymnasium.wrappers.RecordVideo`.
        video_folder:
            Output directory for recorded videos when ``record_video=True``.
        video_kwargs:
            Optional keyword arguments forwarded to
            :class:`~gymnasium.wrappers.RecordVideo`.

    Returns:
        A fully wrapped Gymnasium environment compatible with MASA algorithms,
        monitors, and logging utilities.

    Notes:
        - Wrapper order is fixed and enforced.
        - Constraints are reset automatically on environment reset.
        - All semantic metadata (labels, costs, violations, metrics) is communicated
          via the ``info`` dictionary.

    See Also: 
        - :class:`masa.common.labelled_env.LabelledEnv` 
        - :class:`masa.common.constraints.base.BaseConstraintEnv` 
        - :class:`masa.common.wrappers.ConstraintMonitor` 
        - :class:`masa.common.wrappers.RewardMonitor`
    """

    env_id = resolve_registered_id(
        env_id,
        set(registry.ENV_REGISTRY.keys()),
        format_env_id,
        "env",
    )
    constraint = resolve_registered_id(
        constraint,
        set(registry.CONSTRAINT_REGISTRY.keys()),
        format_constraint_id,
        "constraint",
    )
    env_ctor = registry.get_env(env_id)
    constraint_ctor = registry.get_constraint(constraint)
    env = env_ctor(**dict(env_kwargs or {}))
    # must wrap time limit first
    env = TimeLimit(env, max_episode_steps)
    label_fn = label_fn if label_fn is not None else getattr(env, "label_fn", None)
    if label_fn is not None:
        env = LabelledEnv(env, label_fn)
    constraint_kwargs = dict(constraint_kwargs or {})
    if "cost_fn" not in constraint_kwargs:
        cost_fn = getattr(env, "cost_fn", None)
        if cost_fn is None:
            cost_fn = getattr(env.unwrapped, "cost_fn", None)
        if cost_fn is not None:
            constraint_kwargs["cost_fn"] = cost_fn
    env = constraint_ctor(env, **constraint_kwargs)
    env = ConstraintMonitor(env)
    env = RewardMonitor(env)
    if record_video:
        env = ConstraintPersistentGymnasiumWrapper(
            env,
            GymnasiumRecordVideo,
            video_folder=video_folder,
            **_resolve_video_kwargs(video_kwargs, record_video_episode_trigger),
        )
    return env

masa.common.utils.make_marl_env ¶

make_marl_env(env_id: str, constraint: str, *, label_fn: Optional[dict[str, LabelFn] | LabelFn] = None, constraint_kwargs: Optional[dict[str, Any]] = None, env_kwargs: Optional[dict[str, Any]] = None, record_video: bool = False, record_video_episode_trigger: Optional[Callable[[int], bool]] = None, video_folder: str = 'videos', video_kwargs: Optional[dict[str, Any]] = None, **kw) -> ParallelEnv

Construct a fully wrapped MASA multi-agent environment.

This helper creates a PettingZoo parallel environment and applies the standard MARL wrapper order:

LabelledParallelEnv \(\rightarrow\) ConstrainedMarkovGameEnv

Parameters:

Name	Type	Description	Default
`env_id`	`str`	Multi-agent environment identifier registered in `MARL_ENV_REGISTRY`.	required
`constraint`	`str`	Multi-agent constraint identifier registered in `MARL_CONSTRAINT_REGISTRY`.	required
`label_fn`	`Optional[dict[str, LabelFn] \| LabelFn]`	Optional labelling function, or per-agent mapping of labelling functions. If omitted, the base environment's `label_fn` attribute is used.	`None`
`constraint_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to the constraint wrapper constructor. If `cost_fn` is omitted and the base environment exposes one, it is forwarded automatically.	`None`
`env_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to the base environment constructor.	`None`
`record_video`	`bool`	Whether to wrap the resulting PettingZoo parallel environment with `RecordVideoParallel`. Defaults to `False`.	`False`
`record_video_episode_trigger`	`Optional[Callable[[int], bool]]`	Optional predicate called with the episode id to decide whether to record that episode. This is forwarded as `episode_trigger` to `RecordVideoParallel`.	`None`
`video_folder`	`str`	Output directory for recorded videos when `record_video=True`.	`'videos'`
`video_kwargs`	`Optional[dict[str, Any]]`	Optional keyword arguments forwarded to `RecordVideoParallel`.	`None`

Returns:

Type	Description
`ParallelEnv`	A wrapped PettingZoo parallel environment compatible with MASA MARL
`ParallelEnv`	constraints.

Source code in masa/common/utils.py

def make_marl_env(
    env_id: str,
    constraint: str,
    *,
    label_fn: Optional[dict[str, LabelFn] | LabelFn] = None,
    constraint_kwargs: Optional[dict[str, Any]] = None,
    env_kwargs: Optional[dict[str, Any]] = None,
    record_video: bool = False,
    record_video_episode_trigger: Optional[Callable[[int], bool]] = None,
    video_folder: str = "videos",
    video_kwargs: Optional[dict[str, Any]] = None,
    **kw
) -> ParallelEnv:
    r"""
    Construct a fully wrapped MASA multi-agent environment.

    This helper creates a PettingZoo parallel environment and applies the
    standard MARL wrapper order:

    :class:`~masa.common.multi_agent.labelled_pz_env.LabelledParallelEnv` :math:`\rightarrow`
    :class:`~masa.common.constraints.multi_agent.cmg.ConstrainedMarkovGameEnv`

    Args:
        env_id:
            Multi-agent environment identifier registered in ``MARL_ENV_REGISTRY``.
        constraint:
            Multi-agent constraint identifier registered in
            ``MARL_CONSTRAINT_REGISTRY``.
        label_fn:
            Optional labelling function, or per-agent mapping of labelling
            functions. If omitted, the base environment's ``label_fn`` attribute
            is used.
        constraint_kwargs:
            Optional keyword arguments forwarded to the constraint wrapper
            constructor. If ``cost_fn`` is omitted and the base environment
            exposes one, it is forwarded automatically.
        env_kwargs:
            Optional keyword arguments forwarded to the base environment
            constructor.
        record_video:
            Whether to wrap the resulting PettingZoo parallel environment with
            :class:`~masa.common.pettingzoo_record_video.RecordVideoParallel`.
            Defaults to ``False``.
        record_video_episode_trigger:
            Optional predicate called with the episode id to decide whether to
            record that episode. This is forwarded as ``episode_trigger`` to
            :class:`~masa.common.pettingzoo_record_video.RecordVideoParallel`.
        video_folder:
            Output directory for recorded videos when ``record_video=True``.
        video_kwargs:
            Optional keyword arguments forwarded to
            :class:`~masa.common.pettingzoo_record_video.RecordVideoParallel`.

    Returns:
        A wrapped PettingZoo parallel environment compatible with MASA MARL
        constraints.
    """

    env_id = resolve_registered_id(
        env_id,
        set(registry.MARL_ENV_REGISTRY.keys()),
        format_env_id,
        "env",
    )
    constraint = resolve_registered_id(
        constraint,
        set(registry.MARL_CONSTRAINT_REGISTRY.keys()),
        format_constraint_id,
        "constraint",
    )
    env_ctor = registry.get_marl_env(env_id)
    constraint_ctor = registry.get_marl_constraint(constraint)
    raw_env = env_ctor(**dict(env_kwargs or {}))
    resolved_label_fn = label_fn if label_fn is not None else getattr(raw_env, "label_fn", None)
    if resolved_label_fn is None:
        raise ValueError(
            f"MARL env '{env_id}' does not expose a default label_fn. "
            "Pass label_fn=... to make_marl_env."
        )
    constraint_kwargs = dict(constraint_kwargs or {})
    if "cost_fn" not in constraint_kwargs:
        cost_fn = getattr(raw_env, "cost_fn", None)
        if cost_fn is not None:
            constraint_kwargs["cost_fn"] = cost_fn
    env = LabelledParallelEnv(raw_env, resolved_label_fn)
    env = constraint_ctor(env, **constraint_kwargs)
    if record_video:
        env = RecordVideoParallel(
            env,
            video_folder=video_folder,
            **_resolve_video_kwargs(video_kwargs, record_video_episode_trigger),
        )
    return env

Next Steps¶

Constraints API Reference - View the common constraints supported by MASA.