"""Base class for vectorized environments."""
from __future__ import annotations
from enum import Enum
from typing import TYPE_CHECKING, Any, Generic, cast
import numpy as np
from typing_extensions import TypeVar
import gymnasium as gym
from gymnasium.core import (
ActType,
ObsType,
RenderFrame,
WrapperActType,
WrapperObsType,
)
from gymnasium.logger import warn
from gymnasium.utils import seeding
if TYPE_CHECKING:
from typing import Self
from gymnasium.envs.registration import EnvSpec
ArrayType = TypeVar("ArrayType", default=Any)
__all__ = [
"VectorEnv",
"VectorWrapper",
"VectorObservationWrapper",
"VectorActionWrapper",
"VectorRewardWrapper",
"ArrayType",
"AutoresetMode",
]
[docs]
class AutoresetMode(Enum):
"""The autoreset modes of a vector environment, deciding when a sub-environment is reset after its episode ends.
The mode of a vector environment is given by ``metadata["autoreset_mode"]`` and, for :class:`SyncVectorEnv` and
:class:`AsyncVectorEnv`, is set with the ``autoreset_mode`` argument (an ``AutoresetMode`` or its string value).
See https://gymnasium.farama.org/api/vector/autoreset/ for details and example training loops for each mode.
Examples:
The first sub-environment's episode is truncated on the first step, the second sub-environment's continues.
>>> import numpy as np
>>> import gymnasium as gym
>>> from gymnasium.vector import AutoresetMode, SyncVectorEnv
>>> env_fns = [lambda: gym.make("CartPole-v1", max_episode_steps=1), lambda: gym.make("CartPole-v1")]
Next-step (default), the sub-environment is reset on the following step, ignoring its action, with zero reward
>>> envs = SyncVectorEnv(env_fns, autoreset_mode=AutoresetMode.NEXT_STEP)
>>> envs.metadata["autoreset_mode"]
<AutoresetMode.NEXT_STEP: 'NextStep'>
>>> _ = envs.reset(seed=0)
>>> obs, rewards, terminations, truncations, infos = envs.step(np.array([1, 1]))
>>> rewards, truncations
(array([1., 1.]), array([ True, False]))
>>> obs, rewards, terminations, truncations, infos = envs.step(np.array([1, 1]))
>>> obs[0] # the reset observation
array([0.03132702, 0.04127556, 0.01066358, 0.02294966], dtype=float32)
>>> rewards, truncations
(array([0., 1.]), array([False, False]))
>>> envs.close()
Same-step, the sub-environment is reset within the same step, with the final observation and info in ``infos``
>>> envs = SyncVectorEnv(env_fns, autoreset_mode="SameStep")
>>> _ = envs.reset(seed=0)
>>> obs, rewards, terminations, truncations, infos = envs.step(np.array([1, 1]))
>>> obs[0] # the reset observation
array([0.03132702, 0.04127556, 0.01066358, 0.02294966], dtype=float32)
>>> rewards, truncations
(array([1., 1.]), array([ True, False]))
>>> infos["_final_obs"]
array([ True, False])
>>> infos["final_obs"][0]
array([ 0.01323574, 0.17272775, -0.04686959, -0.3551522 ], dtype=float32)
>>> envs.close()
Disabled, the sub-environment must be reset by the user with a ``reset_mask``, other sub-environments are unaffected
>>> envs = SyncVectorEnv(env_fns, autoreset_mode=AutoresetMode.DISABLED)
>>> _ = envs.reset(seed=0)
>>> obs, rewards, terminations, truncations, infos = envs.step(np.array([1, 1]))
>>> obs
array([[ 0.01323574, 0.17272775, -0.04686959, -0.3551522 ],
[ 0.00208309, 0.24066003, -0.03468674, -0.2588293 ]],
dtype=float32)
>>> obs, infos = envs.reset(options={"reset_mask": np.logical_or(terminations, truncations)})
>>> obs
array([[ 0.03132702, 0.04127556, 0.01066358, 0.02294966],
[ 0.00208309, 0.24066003, -0.03468674, -0.2588293 ]],
dtype=float32)
>>> envs.close()
"""
#: Sub-environments are reset on the step after their episode ends (the default).
NEXT_STEP = "NextStep"
#: Sub-environments are reset within the step that their episode ends.
SAME_STEP = "SameStep"
#: Sub-environments are not reset automatically, use ``reset(options={"reset_mask": mask})``.
DISABLED = "Disabled"
[docs]
class VectorEnv(Generic[ObsType, ActType, ArrayType]):
"""Base class for vectorized environments to run multiple independent copies of the same environment in parallel.
Vector environments can provide a linear speed-up in the steps taken per second through sampling multiple
sub-environments at the same time. Gymnasium contains two generalised Vector environments: :class:`AsyncVectorEnv`
and :class:`SyncVectorEnv` along with several custom vector environment implementations.
For :func:`reset` and :func:`step` batches `observations`, `rewards`, `terminations`, `truncations` and
`info` for each sub-environment, see the example below. For the `rewards`, `terminations`, and `truncations`,
the data is packaged into a NumPy array of shape `(num_envs,)`. For `observations` (and `actions`, the batching
process is dependent on the type of observation (and action) space, and generally optimised for neural network
input/outputs. For `info`, the data is kept as a dictionary such that a key will give the data for all sub-environment.
For creating environments, :func:`make_vec` is a vector environment equivalent to :func:`make` for easily creating
vector environments that contains several unique arguments for modifying environment qualities, number of environment,
vectorizer type, vectorizer arguments.
To avoid having to wait for all sub-environments to end before resetting, sub-environments are reset individually
when their episode ends (`terminated or truncated is True`), with the :class:`AutoresetMode` deciding when:
on the next step (``AutoresetMode.NEXT_STEP``, the default for Gymnasium's implementations), within the same step
(``AutoresetMode.SAME_STEP``) or only when the user calls ``reset(options={"reset_mask": mask})``
(``AutoresetMode.DISABLED``). The mode used by a vector environment should be available in ``metadata["autoreset_mode"]``.
As the mode changes what :meth:`step` returns at the end of an episode, training code must handle it correctly.
Warning, some vector implementations, wrappers or training algorithms will only support particular autoreset modes.
See https://gymnasium.farama.org/api/vector/autoreset/ for each mode's behaviour and example training loops,
and https://farama.org/Vector-Autoreset-Mode for background.
Note:
The info parameter of :meth:`reset` and :meth:`step` was originally implemented before v0.25 as a list
of dictionary for each sub-environment. However, this was modified in v0.25+ to be a dictionary with a NumPy
array for each key. To use the old info style, utilise the :class:`DictInfoToList` wrapper.
Examples:
>>> import gymnasium as gym
>>> envs = gym.make_vec("CartPole-v1", num_envs=3, vectorization_mode="sync", wrappers=(gym.wrappers.TimeAwareObservation,))
>>> envs = gym.wrappers.vector.ClipReward(envs, min_reward=0.2, max_reward=0.8)
>>> envs
<ClipReward, SyncVectorEnv(CartPole-v1, num_envs=3)>
>>> envs.num_envs
3
>>> envs.action_space
MultiDiscrete([2 2 2])
>>> envs.observation_space
Box([[-4.80000019 -inf -0.41887903 -inf 0. ]
[-4.80000019 -inf -0.41887903 -inf 0. ]
[-4.80000019 -inf -0.41887903 -inf 0. ]], [[4.80000019e+00 inf 4.18879032e-01 inf
5.00000000e+02]
[4.80000019e+00 inf 4.18879032e-01 inf
5.00000000e+02]
[4.80000019e+00 inf 4.18879032e-01 inf
5.00000000e+02]], (3, 5), float64)
>>> observations, infos = envs.reset(seed=123)
>>> observations
array([[ 0.01823519, -0.0446179 , -0.02796401, -0.03156282, 0. ],
[ 0.02852531, 0.02858594, 0.0469136 , 0.02480598, 0. ],
[ 0.03517495, -0.000635 , -0.01098382, -0.03203924, 0. ]])
>>> infos
{}
>>> _ = envs.action_space.seed(123)
>>> actions = envs.action_space.sample()
>>> observations, rewards, terminations, truncations, infos = envs.step(actions)
>>> observations
array([[ 0.01734283, 0.15089367, -0.02859527, -0.33293587, 1. ],
[ 0.02909703, -0.16717631, 0.04740972, 0.3319138 , 1. ],
[ 0.03516225, -0.19559774, -0.01162461, 0.25715804, 1. ]])
>>> rewards
array([0.8, 0.8, 0.8])
>>> terminations
array([False, False, False])
>>> truncations
array([False, False, False])
>>> infos
{}
>>> envs.close()
The Vector Environments have the additional attributes for users to understand the implementation
- :attr:`num_envs` - The number of sub-environment in the vector environment
- :attr:`observation_space` - The batched observation space of the vector environment
- :attr:`single_observation_space` - The observation space of a single sub-environment
- :attr:`action_space` - The batched action space of the vector environment
- :attr:`single_action_space` - The action space of a single sub-environment
"""
metadata: dict[str, Any] = {}
spec: EnvSpec | None = None
render_mode: str | None = None
closed: bool = False
observation_space: gym.Space
action_space: gym.Space
single_observation_space: gym.Space
single_action_space: gym.Space
num_envs: int
_np_random: np.random.Generator | None = None
_np_random_seed: int | None = None
[docs]
def reset(
self,
*,
seed: int | None = None,
options: dict[str, Any] | None = None,
) -> tuple[ObsType, dict[str, Any]]: # type: ignore
"""Reset all parallel environments and return a batch of initial observations and info.
Args:
seed: The environment reset seed
options: If to return the options
Returns:
A batch of observations and info from the vectorized environment.
Example:
>>> import gymnasium as gym
>>> envs = gym.make_vec("CartPole-v1", num_envs=3, vectorization_mode="sync")
>>> observations, infos = envs.reset(seed=42)
>>> observations
array([[ 0.0273956 , -0.00611216, 0.03585979, 0.0197368 ],
[ 0.01522993, -0.04562247, -0.04799704, 0.03392126],
[-0.03774345, -0.02418869, -0.00942293, 0.0469184 ]],
dtype=float32)
>>> infos
{}
"""
if seed is not None:
self._np_random, self._np_random_seed = seeding.np_random(seed)
[docs]
def step(
self, actions: ActType
) -> tuple[ObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Take an action for each parallel environment.
Args:
actions: Batch of actions with the :attr:`action_space` shape.
Returns:
Batch of (observations, rewards, terminations, truncations, infos)
Note:
What is returned for a sub-environment whose episode ends depends on the autoreset mode,
``metadata["autoreset_mode"]``, see https://gymnasium.farama.org/api/vector/autoreset/.
- ``AutoresetMode.NEXT_STEP`` (default) - The final observation, reward, termination, truncation and info
are returned as normal. On the next step, the sub-environment's action is ignored and it is reset,
returning the reset observation and info with a reward of zero and no termination or truncation.
- ``AutoresetMode.SAME_STEP`` - The sub-environment is reset within the same step, returning the reset
observation and info with the final reward, termination and truncation. The final observation and info
are in ``infos["final_obs"]`` and ``infos["final_info"]``, masked by ``infos["_final_obs"]`` and
``infos["_final_info"]``.
- ``AutoresetMode.DISABLED`` - The final observation, reward, termination, truncation and info are
returned as normal, and the sub-environment must be reset with ``reset(options={"reset_mask": mask})``
before it is stepped again.
Example:
>>> import gymnasium as gym
>>> import numpy as np
>>> envs = gym.make_vec("CartPole-v1", num_envs=3, vectorization_mode="sync")
>>> _ = envs.reset(seed=42)
>>> actions = np.array([1, 0, 1], dtype=np.int32)
>>> observations, rewards, terminations, truncations, infos = envs.step(actions)
>>> observations
array([[ 0.02727336, 0.18847767, 0.03625453, -0.26141977],
[ 0.01431748, -0.24002443, -0.04731862, 0.3110827 ],
[-0.03822722, 0.1710671 , -0.00848456, -0.2487226 ]],
dtype=float32)
>>> rewards
array([1., 1., 1.])
>>> terminations
array([False, False, False])
>>> truncations
array([False, False, False])
>>> infos
{}
"""
raise NotImplementedError(f"{self.__str__()} step function is not implemented.")
[docs]
def render(self) -> tuple[RenderFrame, ...] | None:
"""Returns the rendered frames from the parallel environments.
Returns:
A tuple of rendered frames from the parallel environments
"""
raise NotImplementedError(
f"{self.__str__()} render function is not implemented."
)
[docs]
def close(self, **kwargs: Any) -> None:
"""Close all parallel environments and release resources.
It also closes all the existing image viewers, then calls :meth:`close_extras` and set
:attr:`closed` as ``True``.
Warnings:
This function itself does not close the environments, it should be handled
in :meth:`close_extras`. This is generic for both synchronous and asynchronous
vectorized environments.
Note:
This will be automatically called when garbage collected or program exited.
Args:
**kwargs: Keyword arguments passed to :meth:`close_extras`
"""
if self.closed:
return
self.close_extras(**kwargs)
self.closed = True
def close_extras(self) -> None:
"""Clean up the extra resources e.g. beyond what's in this base class."""
pass
@property
def np_random(self) -> np.random.Generator:
"""Returns the environment's internal :attr:`_np_random` that if not set will initialise with a random seed.
Returns:
Instances of `np.random.Generator`
"""
if self._np_random is None:
self._np_random, self._np_random_seed = seeding.np_random()
return self._np_random
@np_random.setter
def np_random(self, value: np.random.Generator) -> None:
self._np_random = value
self._np_random_seed = -1
@property
def np_random_seed(self) -> int | None:
"""Returns the environment's internal :attr:`_np_random_seed` that if not set will first initialise with a random int as seed.
If :attr:`np_random_seed` was set directly instead of through :meth:`reset` or :meth:`set_np_random_through_seed`,
the seed will take the value -1.
Returns:
int: the seed of the current `np_random` or -1, if the seed of the rng is unknown
"""
if self._np_random_seed is None:
self._np_random, self._np_random_seed = seeding.np_random()
return self._np_random_seed
@property
def unwrapped(self) -> Self:
"""Return the base environment."""
return self
def _add_info(
self, vector_infos: dict[str, Any], env_info: dict[str, Any], env_num: int
) -> dict[str, Any]:
"""Add env info to the info dictionary of the vectorized environment.
Given the `info` of a single environment add it to the `infos` dictionary
which represents all the infos of the vectorized environment.
Every `key` of `info` is paired with a boolean mask `_key` representing
whether or not the i-indexed environment has this `info`.
Args:
vector_infos (dict): the infos of the vectorized environment
env_info (dict): the info coming from the single environment
env_num (int): the index of the single environment
Returns:
infos (dict): the (updated) infos of the vectorized environment
"""
for key, value in env_info.items():
# It is easier for users to access their `final_obs` in the unbatched array of `obs` objects
if key == "final_obs":
if "final_obs" in vector_infos:
array = vector_infos["final_obs"]
else:
array = np.full(self.num_envs, fill_value=None, dtype=object)
array[env_num] = value
# If value is a dictionary, then we apply the `_add_info` recursively.
elif isinstance(value, dict):
array = self._add_info(vector_infos.get(key, {}), value, env_num)
# Otherwise, we are a base case to group the data
else:
# If the key doesn't exist in the vector infos, then we can create an array of that batch type
if key not in vector_infos:
if type(value) in [int, float, bool] or issubclass(
type(value), np.number
):
array = np.zeros(self.num_envs, dtype=type(value))
elif isinstance(value, np.ndarray):
# We assume that all instances of the np.array info are of the same shape
array = np.zeros(
(self.num_envs, *value.shape), dtype=value.dtype
)
else:
# For unknown objects, we use a Numpy object array
array = np.full(self.num_envs, fill_value=None, dtype=object)
# Otherwise, just use the array that already exists
else:
array = vector_infos[key]
# Assign the data in the `env_num` position
# We only want to run this for the base-case data (not recursive data forcing the ugly function structure)
array[env_num] = value
# Get the array mask and if it doesn't already exist then create a zero bool array
array_mask = vector_infos.get(
f"_{key}", np.zeros(self.num_envs, dtype=np.bool_)
)
array_mask[env_num] = True
# Update the vector info with the updated data and mask information
vector_infos[key], vector_infos[f"_{key}"] = array, array_mask
return vector_infos
def __repr__(self) -> str:
"""Returns a string representation of the vector environment.
Returns:
A string containing the class name, number of environments and environment spec id
"""
if self.spec is None:
return f"{self.__class__.__name__}(num_envs={self.num_envs})"
else:
return (
f"{self.__class__.__name__}({self.spec.id}, num_envs={self.num_envs})"
)
[docs]
class VectorWrapper(
VectorEnv[WrapperObsType, WrapperActType, ArrayType],
Generic[WrapperObsType, WrapperActType, ObsType, ActType, ArrayType],
):
"""Wraps the vectorized environment to allow a modular transformation.
This class is the base class for all wrappers for vectorized environments. The subclass
could override some methods to change the behavior of the original vectorized environment
without touching the original code.
Equivalent to :class:`gymnasium.Wrapper` for vectorized environments, with the same type parameters,
``VectorWrapper[WrapperObsType, WrapperActType, ObsType, ActType, ArrayType]``, where ``WrapperObsType`` and
``WrapperActType`` are the types the wrapper exposes and ``ObsType`` and ``ActType`` are the wrapped environment's.
The reward and termination/truncation ``ArrayType`` is shared between the wrapper and the wrapped environment.
Note:
Don't forget to call ``super().__init__(env)`` if the subclass overrides :meth:`__init__`.
"""
env: VectorEnv[ObsType, ActType, ArrayType]
def __init__(self, env: VectorEnv[ObsType, ActType, ArrayType]) -> None:
"""Initialize the vectorized environment wrapper.
Args:
env: The environment to wrap
"""
self.env = env
if not isinstance(env, VectorEnv):
raise TypeError(
f"Expected env to be a `gymnasium.vector.VectorEnv` but got {type(env)}"
)
self._observation_space: gym.Space | None = None
self._action_space: gym.Space | None = None
self._single_observation_space: gym.Space | None = None
self._single_action_space: gym.Space | None = None
self._metadata: dict[str, Any] | None = None
[docs]
def reset(
self, *, seed: int | None = None, options: dict[str, Any] | None = None
) -> tuple[WrapperObsType, dict[str, Any]]:
"""Reset all environment using seed and options."""
# Only valid if the wrapper doesn't change the observation type, otherwise override `reset`
return cast(
"tuple[WrapperObsType, dict[str, Any]]",
self.env.reset(seed=seed, options=options),
)
[docs]
def step(
self, actions: WrapperActType
) -> tuple[WrapperObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Step through all environments using the actions returning the batched data."""
# Only valid if the wrapper doesn't change the observation or action types, otherwise override `step`
return cast(
"tuple[WrapperObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]",
self.env.step(cast("ActType", actions)),
)
[docs]
def render(self) -> tuple[RenderFrame, ...] | None:
"""Returns the render mode from the base vector environment."""
return self.env.render()
[docs]
def close(self, **kwargs: Any) -> None:
"""Close all environments."""
return self.env.close(**kwargs)
def close_extras(self, **kwargs: Any) -> None:
"""Close all extra resources."""
return self.env.close_extras(**kwargs)
@property
def unwrapped(self) -> VectorEnv:
"""Return the base non-wrapped environment."""
return self.env.unwrapped
def __repr__(self) -> str:
"""Return the string representation of the vectorized environment."""
return f"<{self.__class__.__name__}, {self.env}>"
@property
def observation_space(self) -> gym.Space:
"""Gets the observation space of the vector environment."""
if self._observation_space is None:
return self.env.observation_space
return self._observation_space
@observation_space.setter
def observation_space(self, space: gym.Space) -> None:
"""Sets the observation space of the vector environment."""
self._observation_space = space
@property
def action_space(self) -> gym.Space:
"""Gets the action space of the vector environment."""
if self._action_space is None:
return self.env.action_space
return self._action_space
@action_space.setter
def action_space(self, space: gym.Space) -> None:
"""Sets the action space of the vector environment."""
self._action_space = space
@property
def single_observation_space(self) -> gym.Space:
"""Gets the single observation space of the vector environment."""
if self._single_observation_space is None:
return self.env.single_observation_space
return self._single_observation_space
@single_observation_space.setter
def single_observation_space(self, space: gym.Space) -> None:
"""Sets the single observation space of the vector environment."""
self._single_observation_space = space
@property
def single_action_space(self) -> gym.Space:
"""Gets the single action space of the vector environment."""
if self._single_action_space is None:
return self.env.single_action_space
return self._single_action_space
@single_action_space.setter
def single_action_space(self, space: gym.Space) -> None:
"""Sets the single action space of the vector environment."""
self._single_action_space = space
@property
def num_envs(self) -> int:
"""Gets the wrapped vector environment's num of the sub-environments."""
return self.env.num_envs
@property
def np_random(self) -> np.random.Generator:
"""Returns the environment's internal :attr:`_np_random` that if not set will initialise with a random seed.
Returns:
Instances of `np.random.Generator`
"""
return self.env.np_random
@np_random.setter
def np_random(self, value: np.random.Generator) -> None:
self.env.np_random = value
@property
def np_random_seed(self) -> int | None:
"""The seeds of the vector environment's internal :attr:`_np_random`."""
return self.env.np_random_seed
@property
def metadata(self) -> dict[str, Any]:
"""The metadata of the vector environment."""
if self._metadata is not None:
return self._metadata
return self.env.metadata
@metadata.setter
def metadata(self, value: dict[str, Any]) -> None:
self._metadata = value
@property
def spec(self) -> EnvSpec | None:
"""Gets the specification of the wrapped environment."""
return self.env.spec
@property
def render_mode(self) -> str | None:
"""Returns the `render_mode` from the base environment."""
return self.env.render_mode
@property
def closed(self) -> bool:
"""If the environment has closes."""
return self.env.closed
@closed.setter
def closed(self, value: bool) -> None:
self.env.closed = value
[docs]
class VectorObservationWrapper(
VectorWrapper[WrapperObsType, ActType, ObsType, ActType, ArrayType]
):
"""Wraps the vectorized environment to allow a modular transformation of the observation.
Equivalent to :class:`gymnasium.ObservationWrapper` for vectorized environments.
Same-step autoreset is not supported as ``info["final_obs"]`` is not transformed.
"""
def __init__(self, env: VectorEnv[ObsType, ActType, ArrayType]) -> None:
"""Vector observation wrapper that batch transforms observations.
Args:
env: Vector environment.
"""
super().__init__(env)
if "autoreset_mode" not in env.metadata:
warn(
f"Vector environment ({env}) is missing `autoreset_mode` metadata key."
)
else:
if env.metadata["autoreset_mode"] not in (
AutoresetMode.NEXT_STEP,
AutoresetMode.DISABLED,
):
raise ValueError(
f"Expected autoreset_mode to be NEXT_STEP or DISABLED, got {env.metadata['autoreset_mode']}"
)
def reset(
self, *, seed: int | None = None, options: dict[str, Any] | None = None
) -> tuple[WrapperObsType, dict[str, Any]]:
"""Modifies the observation returned from the environment ``reset`` using the :meth:`observation`."""
observations, infos = self.env.reset(seed=seed, options=options)
return self.observations(observations), infos
def step(
self, actions: ActType
) -> tuple[WrapperObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Modifies the observation returned from the environment ``step`` using the :meth:`observation`."""
observations, rewards, terminations, truncations, infos = self.env.step(actions)
return (
self.observations(observations),
rewards,
terminations,
truncations,
infos,
)
[docs]
def observations(self, observations: ObsType) -> WrapperObsType:
"""Defines the vector observation transformation.
Args:
observations: A vector observation from the environment
Returns:
the transformed observation
"""
raise NotImplementedError
[docs]
class VectorActionWrapper(
VectorWrapper[ObsType, WrapperActType, ObsType, ActType, ArrayType]
):
"""Wraps the vectorized environment to allow a modular transformation of the actions.
Equivalent of :class:`gymnasium.ActionWrapper` for vectorized environments.
"""
def __init__(self, env: VectorEnv[ObsType, ActType, ArrayType]) -> None:
"""Vector action wrapper that batch transforms actions.
Args:
env: Vector environment.
"""
super().__init__(env)
def step(
self, actions: WrapperActType
) -> tuple[ObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Steps through the environment using a modified action by :meth:`action`."""
return self.env.step(self.actions(actions))
[docs]
def actions(self, actions: WrapperActType) -> ActType:
"""Transform the actions before sending them to the environment.
Args:
actions (WrapperActType): the actions to transform
Returns:
ActType: the transformed actions
"""
raise NotImplementedError
[docs]
class VectorRewardWrapper(VectorWrapper[ObsType, ActType, ObsType, ActType, ArrayType]):
"""Wraps the vectorized environment to allow a modular transformation of the reward.
Equivalent of :class:`gymnasium.RewardWrapper` for vectorized environments.
"""
def __init__(self, env: VectorEnv[ObsType, ActType, ArrayType]) -> None:
"""Vector reward wrapper that batch transforms rewards.
Args:
env: Vector environment.
"""
super().__init__(env)
self._autoreset_mode = self.env.metadata.get(
"autoreset_mode", AutoresetMode.NEXT_STEP
)
self._prev_dones = np.zeros((self.num_envs,), dtype=np.bool_)
def reset(
self,
*,
seed: int | list[int] | None = None,
options: dict[str, Any] | None = None,
) -> tuple[ObsType, dict[str, Any]]:
"""Resets the environment and the tracking of which sub-environments will autoreset on the next step."""
if options is not None and "reset_mask" in options:
self._prev_dones[options["reset_mask"]] = False
else:
self._prev_dones = np.zeros((self.num_envs,), dtype=np.bool_)
return self.env.reset(seed=seed, options=options)
def step(
self, actions: ActType
) -> tuple[ObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Steps through the environment returning a reward modified by :meth:`reward`.
For next-step autoreset, the reward of a sub-environment's autoreset step is not an environment reward
therefore, it is returned unmodified, equivalent to applying the wrapper to each sub-environment.
"""
observations, rewards, terminations, truncations, infos = self.env.step(actions)
transformed_rewards = self.rewards(rewards)
if self._autoreset_mode != AutoresetMode.NEXT_STEP:
return observations, transformed_rewards, terminations, truncations, infos
if np.any(self._prev_dones):
transformed_rewards = np.where(
self._prev_dones, rewards, transformed_rewards
)
self._prev_dones = np.logical_or(terminations, truncations)
return observations, transformed_rewards, terminations, truncations, infos
[docs]
def rewards(self, rewards: ArrayType) -> ArrayType:
"""Transform the reward before returning it.
Args:
rewards (array): the reward to transform
Returns:
array: the transformed reward
"""
raise NotImplementedError
