Source code for pettingzoo.utils.wrappers.clip_out_of_bounds

from __future__ import annotations

from typing import Any

import numpy as np
from gymnasium.spaces import Box
from typing_extensions import override

from pettingzoo.utils.env import AECEnv
from pettingzoo.utils.env_logger import EnvLogger
from pettingzoo.utils.wrappers.base import BaseWrapper




class ClipOutOfBoundsWrapper(BaseWrapper[Any, Any, Any]):
    """Clips the input action to fit in the continuous action space (emitting a warning if it does so).

    Applied to continuous environments in pettingzoo.
    """

    def __init__(self, env: AECEnv[Any, Any, Any]):
        super().__init__(env)
        assert isinstance(env, AECEnv), (
            "ClipOutOfBoundsWrapper is only compatible with AEC environments."
        )
        assert all(
            isinstance(self.action_space(agent), Box)
            for agent in getattr(self, "possible_agents", [])
        ), "should only use ClipOutOfBoundsWrapper for Box spaces"

    @override
    def step(self, action: np.ndarray | None) -> None:
        space = self.action_space(self.agent_selection)
        assert isinstance(space, Box), (
            "should only use ClipOutOfBoundsWrapper for Box spaces"
        )
        if action is not None and not space.contains(action):
            if np.isnan(action).any():
                EnvLogger.error_nan_action()
            assert space.shape == action.shape, (
                f"action should have shape {space.shape}, has shape {action.shape}"
            )

            EnvLogger.warn_action_out_of_bound(
                action=action, action_space=space, backup_policy="clipping to space"
            )
            action = np.clip(
                action,
                space.low,
                space.high,
            )

        super().step(action)

    @override
    def __str__(self) -> str:
        return str(self.env)