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

class masa.common.wrappers.ConstraintPersistentWrapper(env: Env)[source]

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:

env – Base environment to wrap.

Wraps an environment to allow a modular transformation of the step() and reset() methods.

Parameters:

env – The environment to wrap

property _constraint

The underlying constraint object, if present.

Returns:

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

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:

A cost function-like object or None.

property label_fn

Labelling function exposed by the underlying environment, if available.

Returns:

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

class masa.common.wrappers.ConstraintPersistentObsWrapper(env: Env)[source]

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:

env – Base environment to wrap.

Wraps an environment to allow a modular transformation of the step() and reset() methods.

Parameters:

env – The environment to wrap

_get_obs() → Any[source]

Transform a raw observation into the wrapped observation.

Subclasses must implement this.

Parameters:

obs – Raw observation from the underlying environment.

Returns:

Transformed observation.

Raises:

NotImplementedError – If not implemented by the subclass.

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

Reset the environment and transform the returned observation.

Parameters:

  • 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 _get_obs().

step(action)[source]

Step the environment and transform the returned observation.

Parameters:

action – Action forwarded to the underlying environment.

Returns:

A 5-tuple (obs, reward, terminated, truncated, info) where obs is transformed via _get_obs().

Helpers

wrappers.is_wrapped(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:

  • 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.

wrappers.get_wrapped(wrapper_class: Wrapper) → Env

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

The traversal rules match is_wrapped().

Parameters:

  • 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).

Next Steps