lit_sb3.vec_env.base_vec_env

[docs] module lit_sb3.vec_env.base_vec_env

from abc import ABC, abstractmethod
from collections.abc import Iterable, Sequence
from copy import deepcopy
import inspect
from typing import Any, Optional
import warnings

import cloudpickle
import gymnasium as gym
from gymnasium import spaces
import numpy as np

# Define type aliases here to avoid circular import
# Used when we want to access one or more VecEnv
VecEnvIndices = None | int | Iterable[int]
# VecEnvObs is what is returned by the reset() method
# it contains the observation for each env
VecEnvObs = np.ndarray | dict[str, np.ndarray] | tuple[np.ndarray, ...]
# VecEnvStepReturn is what is returned by the step() method
# it contains the observation, reward, done, info for each env
VecEnvStepReturn = tuple[VecEnvObs, np.ndarray, np.ndarray, list[dict]]

[docs]
def tile_images(images_nhwc: Sequence[np.ndarray]) -> np.ndarray:  # pragma: no cover
    """Tile N images into one big PxQ image.

    (P,Q) are chosen to be as close as possible, and if N
    is square, then P=Q.

    Args:
        images_nhwc: list or array of images, ndim=4 once turned into array.
            n = batch index, h = height, w = width, c = channel

    Returns:
        img_HWc, ndim=3
    """
    img_nhwc = np.asarray(images_nhwc)
    n_images, height, width, n_channels = img_nhwc.shape
    # new_height was named H before
    new_height = int(np.ceil(np.sqrt(n_images)))
    # new_width was named W before
    new_width = int(np.ceil(float(n_images) / new_height))
    img_nhwc = np.array(list(img_nhwc) + [img_nhwc[0] * 0 for _ in range(n_images, new_height * new_width)])
    # img_HWhwc
    out_image = img_nhwc.reshape((new_height, new_width, height, width, n_channels))
    # img_HhWwc
    out_image = out_image.transpose(0, 2, 1, 3, 4)
    # img_Hh_Ww_c
    out_image = out_image.reshape((new_height * height, new_width * width, n_channels))
    return out_image


class VecEnv(ABC):
[docs]
    """An abstract asynchronous, vectorized environment.

    Args:
        num_envs: Number of environments
        observation_space: Observation space
        action_space: Action space
    """

    def __init__(
        self,
        num_envs: int,
        observation_space: spaces.Space,
        action_space: spaces.Space,
    ) -> None:
        self.num_envs = num_envs
        self.observation_space = observation_space
        self.action_space = action_space
        # store info returned by the reset method
        self.reset_infos: list[dict[str, Any]] = [{} for _ in range(num_envs)]
        # seeds to be used in the next call to env.reset()
        self._seeds: list[int | None] = [None for _ in range(num_envs)]
        # options to be used in the next call to env.reset()
        self._options: list[dict[str, Any]] = [{} for _ in range(num_envs)]

        try:
            render_modes = self.get_attr("render_mode")
        except AttributeError:
            warnings.warn(
                "The `render_mode` attribute is not defined in your environment. It will be set to None.", stacklevel=2
            )
            render_modes = [None for _ in range(num_envs)]

        assert all(render_mode == render_modes[0] for render_mode in render_modes), (
            "render_mode mode should be the same for all environments"
        )
        self.render_mode = render_modes[0]

        render_modes = (
            ["human", "rgb_array"]
            if self.render_mode == "rgb_array"
            else [self.render_mode]
            if self.render_mode is not None
            else []
        )

        self.metadata = {"render_modes": render_modes}

    def _reset_seeds(self) -> None:
        """Reset the seeds that are going to be used at the next reset."""
        self._seeds = [None for _ in range(self.num_envs)]

    def _reset_options(self) -> None:
        """Reset the options that are going to be used at the next reset."""
        self._options = [{} for _ in range(self.num_envs)]

    @abstractmethod
    def reset(self) -> VecEnvObs:
[docs]
        """Reset all the environments and return an array of.

        observations, or a tuple of observation arrays.

        If step_async is still doing work, that work will
        be cancelled and step_wait() should not be called
        until step_async() is invoked again.

        Returns:
            observation
        """
        raise NotImplementedError()

    @abstractmethod
    def step_async(self, actions: np.ndarray) -> None:
[docs]
        """Tell all the environments to start taking a step.

        with the given actions.
        Call step_wait() to get the results of the step.

        You should not call this if a step_async run is
        already pending.
        """
        raise NotImplementedError()

    @abstractmethod
    def step_wait(self) -> VecEnvStepReturn:
[docs]
        """Wait for the step taken with step_async().

        Returns:
            observation, reward, done, information
        """
        raise NotImplementedError()

    @abstractmethod
    def close(self) -> None:
[docs]
        """Clean up the environment's resources."""
        raise NotImplementedError()

    def has_attr(self, attr_name: str) -> bool:
[docs]
        """Check if an attribute exists for a vectorized environment.

        Args:
            attr_name: The name of the attribute to check

        Returns:
            True if 'attr_name' exists in all environments
        """
        # Default implementation, will not work with things that cannot be pickled:
        # https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/issues/49
        try:
            self.get_attr(attr_name)
        except AttributeError:
            return False
        else:
            return True

    @abstractmethod
[docs]
    def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> list[Any]:
        """Return attribute from vectorized environment.

        Args:
            attr_name: The name of the attribute whose value to return
            indices: Indices of envs to get attribute from

        Returns:
            List of values of 'attr_name' in all environments
        """
        raise NotImplementedError()

    @abstractmethod
[docs]
    def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
        """Set attribute inside vectorized environments.

        Args:
            attr_name: The name of attribute to assign new value
            value: Value to assign to `attr_name`
            indices: Indices of envs to assign value
        """
        raise NotImplementedError()

    @abstractmethod
[docs]
    def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> list[Any]:
        """Call instance methods of vectorized environments.

        Args:
            method_name: The name of the environment method to invoke.
            indices: Indices of envs whose method to call
            method_args: Any positional arguments to provide in the call
            method_kwargs: Any keyword arguments to provide in the call

        Returns:
            List of items returned by the environment's method call
        """
        raise NotImplementedError()

    @abstractmethod
[docs]
    def env_is_wrapped(self, wrapper_class: type[gym.Wrapper], indices: VecEnvIndices = None) -> list[bool]:
        """Check if environments are wrapped with a given wrapper.

        Args:
            method_name: The name of the environment method to invoke.
            indices: Indices of envs whose method to call
            method_args: Any positional arguments to provide in the call
            method_kwargs: Any keyword arguments to provide in the call

            wrapper_class: The wrapper class parameter.

        Returns:
            True if the env is wrapped, False otherwise, for each env queried.
        """
        raise NotImplementedError()

    def step(self, actions: np.ndarray) -> VecEnvStepReturn:
[docs]
        """Step the environments with the given action.

        Args:
            actions: the action

        Returns:
            observation, reward, done, information
        """
        self.step_async(actions)
        return self.step_wait()

    def get_images(self) -> Sequence[np.ndarray | None]:
[docs]
        """Return RGB images from each environment when available."""
        raise NotImplementedError

    def render(self, mode: str | None = None) -> np.ndarray | None:
[docs]
        """Gym environment rendering.

        Args:
            mode: the rendering type
        """
        if mode == "human" and self.render_mode != mode:
            # Special case, if the render_mode="rgb_array"
            # we can still display that image using opencv
            if self.render_mode != "rgb_array":
                warnings.warn(
                    f"You tried to render a VecEnv with mode='{mode}' "
                    "but the render mode defined when initializing the environment must be "
                    f"'human' or 'rgb_array', not '{self.render_mode}'.",
                    stacklevel=2,
                )
                return None

        elif mode and self.render_mode != mode:
            warnings.warn(
                f"""Starting from gymnasium v0.26, render modes are determined during the initialization of the environment.
                We allow to pass a mode argument to maintain a backwards compatible VecEnv API, but the mode ({mode})
                has to be the same as the environment render mode ({self.render_mode}) which is not the case.""",
                stacklevel=2,
            )
            return None

        mode = mode or self.render_mode

        if mode is None:
            warnings.warn(
                "You tried to call render() but no `render_mode` was passed to the env constructor.", stacklevel=2
            )
            return None

        # mode == self.render_mode == "human"
        # In that case, we try to call `self.env.render()` but it might
        # crash for subprocesses
        if self.render_mode == "human":
            self.env_method("render")
            return None

        if mode == "rgb_array" or mode == "human":
            # call the render method of the environments
            images = self.get_images()
            # Create a big image by tiling images from subprocesses
            bigimg = tile_images(images)  # type: ignore[arg-type]

            if mode == "human":
                # Display it using OpenCV
                import cv2  # ty:ignore[unresolved-import]

                cv2.imshow("vecenv", bigimg[:, :, ::-1])
                cv2.waitKey(1)
            else:
                return bigimg

        else:
            # Other render modes:
            # In that case, we try to call `self.env.render()` but it might
            # crash for subprocesses
            # and we don't return the values
            self.env_method("render")
        return None

    def seed(self, seed: int | None = None) -> Sequence[None | int]:
[docs]
        """Sets the random seeds for all environments, based on a given seed.

        Each individual environment will still get its own seed, by incrementing the given seed.
        WARNING: since gym 0.26, those seeds will only be passed to the environment
        at the next reset.

        Args:
            seed: The random seed. May be None for completely random seeding.

        Returns:
            Returns a list containing the seeds for each individual env.
            Note that all list elements may be None, if the env does not return anything when being seeded.
        """
        if seed is None:
            # To ensure that subprocesses have different seeds,
            # we still populate the seed variable when no argument is passed
            seed = int(np.random.randint(0, np.iinfo(np.uint32).max, dtype=np.uint32))

        self._seeds = [seed + idx for idx in range(self.num_envs)]
        return self._seeds
[docs]
    def set_options(self, options: list[dict] | dict | None = None) -> None:
        """Set environment options for all environments.

        If a dict is passed instead of a list, the same options will be used for all environments.
        WARNING: Those options will only be passed to the environment at the next reset.

        Args:
            options: A dictionary of environment options to pass to each environment at the next reset.
        """
        if options is None:
            options = {}
        # Use deepcopy to avoid side effects
        if isinstance(options, dict):
            self._options = deepcopy([options] * self.num_envs)
        else:
            self._options = deepcopy(options)

    @property
    def unwrapped(self) -> "VecEnv":
[docs]
        if isinstance(self, VecEnvWrapper):
            return self.venv.unwrapped
        else:
            return self
[docs]
    def getattr_depth_check(self, name: str, already_found: bool) -> str | None:
        """Check if an attribute reference is being hidden in a recursive call to __getattr__.

        Args:
            name: name of attribute to check for
            already_found: whether this attribute has already been found in a wrapper

        Returns:
            name of module whose attribute is being shadowed, if any.
        """
        if hasattr(self, name) and already_found:
            return f"{type(self).__module__}.{type(self).__name__}"
        else:
            return None

    def _get_indices(self, indices: VecEnvIndices) -> Iterable[int]:
        """Convert a flexibly-typed reference to environment indices to an implied list of indices.

        Args:
            indices: refers to indices of envs.

        Returns:
            the implied list of indices.
        """
        if indices is None:
            indices = range(self.num_envs)
        elif isinstance(indices, int):
            indices = [indices]
        return indices


class VecEnvWrapper(VecEnv):
[docs]
    """Vectorized environment base class.

    Args:
        venv: the vectorized environment to wrap
        observation_space: the observation space (can be None to load from venv)
        action_space: the action space (can be None to load from venv)
    """

    def __init__(
        self,
        venv: VecEnv,
        observation_space: spaces.Space | None = None,
        action_space: spaces.Space | None = None,
    ) -> None:
        self.venv = venv

        super().__init__(
            num_envs=venv.num_envs,
            observation_space=observation_space or venv.observation_space,
            action_space=action_space or venv.action_space,
        )
        self.class_attributes = dict(inspect.getmembers(self.__class__))

    def step_async(self, actions: np.ndarray) -> None:
[docs]
        self.venv.step_async(actions)

    @abstractmethod
    def reset(self) -> VecEnvObs:
[docs]
        pass

    @abstractmethod
    def step_wait(self) -> VecEnvStepReturn:
[docs]
        pass

    def seed(self, seed: int | None = None) -> Sequence[None | int]:
[docs]
        return self.venv.seed(seed)
[docs]
    def set_options(self, options: list[dict] | dict | None = None) -> None:
        return self.venv.set_options(options)

    def close(self) -> None:
[docs]
        return self.venv.close()

    def render(self, mode: str | None = None) -> np.ndarray | None:
[docs]
        return self.venv.render(mode=mode)

    def get_images(self) -> Sequence[np.ndarray | None]:
[docs]
        return self.venv.get_images()

    def has_attr(self, attr_name: str) -> bool:
[docs]
        return self.venv.has_attr(attr_name)
[docs]
    def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> list[Any]:
        return self.venv.get_attr(attr_name, indices)
[docs]
    def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
        return self.venv.set_attr(attr_name, value, indices)
[docs]
    def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> list[Any]:
        return self.venv.env_method(method_name, *method_args, indices=indices, **method_kwargs)
[docs]
    def env_is_wrapped(self, wrapper_class: type[gym.Wrapper], indices: VecEnvIndices = None) -> list[bool]:
        return self.venv.env_is_wrapped(wrapper_class, indices=indices)

    def __getattr__(self, name: str) -> Any:
        """Find attribute from wrapped venv(s) if this wrapper does not have it.

        Useful for accessing attributes from venvs which are wrapped with multiple wrappers
        which have unique attributes of interest.
        """
        blocked_class = self.getattr_depth_check(name, already_found=False)
        if blocked_class is not None:
            own_class = f"{type(self).__module__}.{type(self).__name__}"
            error_str = (
                f"Error: Recursive attribute lookup for {name} from {own_class} is "
                f"ambiguous and hides attribute from {blocked_class}"
            )
            raise AttributeError(error_str)

        return self.getattr_recursive(name)

    def _get_all_attributes(self) -> dict[str, Any]:
        """Get all (inherited) instance and class attributes.

        Returns:
            all_attributes
        """
        all_attributes = self.__dict__.copy()
        all_attributes.update(self.class_attributes)
        return all_attributes

    def getattr_recursive(self, name: str) -> Any:
[docs]
        """Recursively check wrappers to find attribute.

        Args:
            name: name of attribute to look for

        Returns:
            attribute
        """
        all_attributes = self._get_all_attributes()
        if name in all_attributes:  # attribute is present in this wrapper
            attr = getattr(self, name)
        elif hasattr(self.venv, "getattr_recursive"):
            # Attribute not present, child is wrapper. Call getattr_recursive rather than getattr
            # to avoid a duplicate call to getattr_depth_check.
            attr = self.venv.getattr_recursive(name)  # ty:ignore[call-non-callable]
        else:  # attribute not present, child is an unwrapped VecEnv
            attr = getattr(self.venv, name)

        return attr
[docs]
    def getattr_depth_check(self, name: str, already_found: bool) -> str | None:
        """See base class.

        Returns:
            name of module whose attribute is being shadowed, if any.
        """
        all_attributes = self._get_all_attributes()
        if name in all_attributes and already_found:
            # this venv's attribute is being hidden because of a higher venv.
            shadowed_wrapper_class: str | None = f"{type(self).__module__}.{type(self).__name__}"
        elif name in all_attributes and not already_found:
            # we have found the first reference to the attribute. Now check for duplicates.
            shadowed_wrapper_class = self.venv.getattr_depth_check(name, True)
        else:
            # this wrapper does not have the attribute. Keep searching.
            shadowed_wrapper_class = self.venv.getattr_depth_check(name, already_found)

        return shadowed_wrapper_class


class CloudpickleWrapper:
[docs]
    """Uses cloudpickle to serialize contents (otherwise multiprocessing tries to use pickle).

    Args:
        var: the variable you wish to wrap for pickling with cloudpickle
    """

    def __init__(self, var: Any) -> None:
        self.var = var

    def __getstate__(self) -> Any:
        return cloudpickle.dumps(self.var)

    def __setstate__(self, var: Any) -> None:
        self.var = cloudpickle.loads(var)