Core Wrappers¶

def step(self, action):
    """
    Step the environment and apply time-limit truncation.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(observation, reward, terminated, truncated, info)``. If
        the time limit is reached, ``truncated`` is forced to ``True``.
    """
    observation, reward, terminated, truncated, info = self.env.step(action)
    self._elapsed_steps += 1

    if self._elapsed_steps >= self._max_episode_steps:
        truncated = True

    return observation, reward, terminated, truncated, info

def reset(self, **kwargs):
    """
    Reset the environment and the elapsed step counter.

    Args:
        **kwargs: Forwarded to the underlying environment's ``reset``.

    Returns:
        The underlying environment's ``reset`` return value.
    """
    self._elapsed_steps = 0
    return self.env.reset(**kwargs)

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset and populate initial constraint metadata in ``info``.

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

    Returns:
        A tuple ``(obs, info)``. The returned ``info`` includes
        ``info["constraint"]["type"]`` and ``info["constraint"]["step"]``.
    """
    obs, info = self.env.reset(seed=seed, options=options)
    info = dict(info or {})
    info.setdefault("constraint", {})["type"] = self._constraint_env.constraint_type
    info["constraint"]["step"] = self._step_metrics()
    return obs, info

def step(self, action):
    """
    Step and populate constraint metrics in ``info``.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(observation, reward, terminated, truncated, info)``.
        The returned ``info`` includes ``constraint`` fields described in the
        class docstring.
    """
    observation, reward, terminated, truncated, info = self.env.step(action)
    info = dict(info or {})
    info.setdefault("constraint", {})["type"] = self._constraint_env.constraint_type
    info["constraint"]["step"] = self._step_metrics()
    if terminated or truncated:
        info["constraint"]["episode"] = self._episode_metrics()
    return observation, reward, terminated, truncated, info

def _step_metrics(self) -> Dict[str, float]:
    """
    Read step-level constraint metrics.

    Returns:
        A dictionary of step-level metrics. If the underlying constraint raises
        an exception, returns an empty dictionary.
    """
    try:
        return dict(self._constraint_env.constraint_step_metrics())
    except Exception:
        return {}

def _episode_metrics(self) -> Dict[str, float]:
    """
    Read episode-level constraint metrics.

    Returns:
        A dictionary of episode-level metrics. If the underlying constraint raises
        an exception, returns an empty dictionary.
    """
    try:
        return dict(self._constraint_env.constraint_episode_metrics())
    except Exception:
        return {}

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset reward counters and forward ``reset`` to the underlying env.

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

    Returns:
        A tuple ``(obs, info)`` from the underlying environment.
    """
    obs, info = self.env.reset(seed=seed, options=options)
    self.total_reward = 0.0
    self.total_steps = 0
    return obs, info

def step(self, action):
    """
    Step the environment and update reward metrics.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(observation, reward, terminated, truncated, info)``.
        On episode end, ``info["metrics"]["episode"]`` is populated with
        episode return and length.
    """
    observation, reward, terminated, truncated, info = self.env.step(action)
    self.total_reward += reward
    self.total_steps += 1
    info = dict(info or {})
    info.setdefault("metrics", {})
    info["metrics"]["step"] = {"reward": reward}
    if terminated or truncated:
        info["metrics"]["episode"] = self._episode_metrics()
    return observation, reward, terminated, truncated, info

def _episode_metrics(self):
    """
    Compute episode-level reward metrics.

    Returns:
        A dictionary with keys:

        * ``"ep_reward"``: total episode reward.
        * ``"ep_length"``: episode length in steps.
    """
    return {"ep_reward": self.total_reward, "ep_length": self.total_steps}