Skip to content

Wrappers

Environment wrappers for MASA-Safe-RL.

This module contains small, composable gymnasium.Wrapper utilities that (1) preserve access to constraint-related objects through wrapper chains, (2) inject monitoring/metrics into info, (3) apply potential-based reward shaping for DFA-based constraints, and (4) provide basic observation/reward normalization and light-weight vector-environment helpers.

Key conventions

  • Constraint-enabled environments expose a _constraint object and (often) label_fn / cost_fn attributes. See masa.common.constraints.base.BaseConstraintEnv.
  • Monitoring wrappers add structured dictionaries under info["constraint"] and/or info["metrics"].
  • Vector wrappers in this file use a simple Python list API: observations, rewards, terminals, truncations, infos are lists of length VecEnvWrapperBase.n_envs.

Notes

For potential-based shaping, the shaped cost inserted into info is of the form

\[ c'_t \;=\; c_t \;+\; \gamma \Phi(q_{t+1}) \;-\; \Phi(q_t), \]

where \(q_t\) is the DFA state, \(c_t\) is the original constraint cost, \(\Phi\) is the potential function, and \(\gamma\) is the shaping discount factor.

API Reference

Base Class

masa.common.wrappers.ConstraintPersistentWrapper 

ConstraintPersistentWrapper(env: Env)

Bases: Wrapper

Base wrapper that persists access to constraint-related attributes.

Many Gymnasium wrappers shadow attributes by changing self.env. This wrapper provides stable properties for:

  • _constraint (if present on the underlying env)
  • cost_fn (if exposed by the constraint)
  • label_fn (if present on the underlying env)

Subclasses can rely on these properties even when stacked with additional wrappers.

Parameters:

Name Type Description Default
env Env

Base environment to wrap.

 required
Source code in masa/common/wrappers.py 
110
111
def __init__(self, env: gym.Env):
    super().__init__(env)

_constraint property 

_constraint

The underlying constraint object, if present.

Returns:

Type Description

The object stored in self.env._constraint if it exists, otherwise

None.

cost_fn property 

cost_fn

Cost function exposed by the underlying constraint, if available.

If self._constraint exists and it exposes cost_fn, this returns that callable-like object (often a masa.common.ltl.DFACostFn). Otherwise returns None.

Returns:

Type Description

A cost function-like object or None.

label_fn property 

label_fn

Labelling function exposed by the underlying environment, if available.

Returns:

Type Description

The object stored in self.env.label_fn if it exists, otherwise

None.

masa.common.wrappers.ConstraintPersistentObsWrapper 

ConstraintPersistentObsWrapper(env: Env)

Bases: ConstraintPersistentWrapper

Base class for wrappers that transform observations while preserving constraint access.

Subclasses must implement _get_obs which maps raw observations to the wrapped observation representation.

Parameters:

Name Type Description Default
env Env

Base environment to wrap.

 required
Source code in masa/common/wrappers.py 
163
164
def __init__(self, env: gym.Env):
    super().__init__(env)

_get_obs 

_get_obs(obs: Any) -> Any

Transform a raw observation into the wrapped observation.

Subclasses must implement this.

Parameters:

Name Type Description Default
obs Any

Raw observation from the underlying environment.

 required

Returns:

Type Description
Any

Transformed observation.

Raises:

Type Description
NotImplementedError

If not implemented by the subclass.
Source code in masa/common/wrappers.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def _get_obs(obs: Any) -> Any:
    """
    Transform a raw observation into the wrapped observation.

    Subclasses must implement this.

    Args:
        obs: Raw observation from the underlying environment.

    Returns:
        Transformed observation.

    Raises:
        NotImplementedError: If not implemented by the subclass.
    """
    raise NotImplementedError

reset 

reset(*, seed: int | None = None, options: Dict[str, Any] | None = None)

Reset the environment and transform the returned observation.

Parameters:

Name Type Description Default
seed int | None

Random seed forwarded to the underlying environment.

 None
options Dict[str, Any] | None

Reset options forwarded to the underlying environment.

 None

Returns:

Type Description

A tuple (obs, info) where obs is transformed via _get_obs.
Source code in masa/common/wrappers.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset the environment and transform the returned observation.

    Args:
        seed: Random seed forwarded to the underlying environment.
        options: Reset options forwarded to the underlying environment.

    Returns:
        A tuple ``(obs, info)`` where ``obs`` is transformed via :meth:`_get_obs`.
    """
    obs, info = self.env.reset(seed=seed, options=options)
    return self._get_obs(obs), info

step 

step(action)

Step the environment and transform the returned observation.

Parameters:

Name Type Description Default
action

Action forwarded to the underlying environment.

 required

Returns:

Type Description

A 5-tuple (obs, reward, terminated, truncated, info) where obs

is transformed via _get_obs.
Source code in masa/common/wrappers.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
def step(self, action):
    """
    Step the environment and transform the returned observation.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(obs, reward, terminated, truncated, info)`` where ``obs``
        is transformed via :meth:`_get_obs`.
    """
    obs, rew, term, trunc, info = self.env.step(action)
    return self._get_obs(obs), rew, term, trunc, info

Helpers

masa.common.wrappers.is_wrapped 

is_wrapped(env: Env, wrapper_class: Wrapper) -> bool

Check whether env is wrapped (anywhere in its wrapper chain) by wrapper_class.

This helper walks through typical wrapper chains:

  • Gymnasium-style wrappers via .env.
  • Vector-env style wrappers via a .venv attribute (commonly used by vectorized environments and some third-party libraries).

Cycle protection is included: if the wrapper chain loops, this function returns False rather than looping forever.

Parameters:

Name Type Description Default
env Env

Environment or wrapper to inspect.

 required
wrapper_class Wrapper

Wrapper type to search for.

 required

Returns:

Type Description
bool

True if an instance of wrapper_class appears in the wrapper chain;
bool

False otherwise.
Source code in masa/common/wrappers.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def is_wrapped(env: gym.Env, wrapper_class: gym.Wrapper) -> bool:
    r"""
    Check whether ``env`` is wrapped (anywhere in its wrapper chain) by
    ``wrapper_class``.

    This helper walks through typical wrapper chains:

    * Gymnasium-style wrappers via ``.env``.
    * Vector-env style wrappers via a ``.venv`` attribute (commonly used by
      vectorized environments and some third-party libraries).

    Cycle protection is included: if the wrapper chain loops, this function
    returns ``False`` rather than looping forever.

    Args:
        env: Environment or wrapper to inspect.
        wrapper_class: Wrapper type to search for.

    Returns:
        ``True`` if an instance of ``wrapper_class`` appears in the wrapper chain;
        ``False`` otherwise.
    """
    current = env
    visited = set()

    while True:
        if id(current) in visited:
            return False
        visited.add(id(current))

        if isinstance(current, wrapper_class):
            return True

        if hasattr(current, "venv"):
            current = current.venv
            continue

        if isinstance(current, gym.Wrapper):
            current = current.env
            continue

        return False

masa.common.wrappers.get_wrapped 

get_wrapped(env: Env, wrapper_class: Wrapper) -> gym.Env

Return the first wrapper instance of type wrapper_class found in env's wrapper chain.

The traversal rules match is_wrapped.

Parameters:

Name Type Description Default
env Env

Environment or wrapper to inspect.

 required
wrapper_class Wrapper

Wrapper type to retrieve.

 required

Returns:

Type Description
Env

The first encountered instance of wrapper_class in the wrapper chain,
Env

or None if it is not present (or if a cycle is detected).
Source code in masa/common/wrappers.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def get_wrapped(env: gym.Env, wrapper_class: gym.Wrapper) -> gym.Env:
    r"""
    Return the first wrapper instance of type ``wrapper_class`` found in ``env``'s
    wrapper chain.

    The traversal rules match :func:`is_wrapped`.

    Args:
        env: Environment or wrapper to inspect.
        wrapper_class: Wrapper type to retrieve.

    Returns:
        The first encountered instance of ``wrapper_class`` in the wrapper chain,
        or ``None`` if it is not present (or if a cycle is detected).
    """

    current = env
    visited = set()

    while True:
        if id(current) in visited:
            return None
        visited.add(id(current))

        if isinstance(current, wrapper_class):
            return current

        if hasattr(current, "venv"):
            current = current.venv
            continue

        if isinstance(current, gym.Wrapper):
            current = current.env
            continue

        return None

Next Steps