Utility functions for vectorisation#

Spaces utility functions#

gymnasium.experimental.vector.utils.batch_space(space: Space[Any], n: int = 1) → Space[Any]#

gymnasium.experimental.vector.utils.batch_space(space: Box, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: Discrete, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: MultiDiscrete, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: MultiBinary, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: Tuple, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: Dict, n: int = 1)

gymnasium.experimental.vector.utils.batch_space(space: Graph | Text | Sequence, n: int = 1)

Create a (batched) space, containing multiple copies of a single space.

Parameters:

space – Space (e.g. the observation space) for a single environment in the vectorized environment.
n – Number of environments in the vectorized environment.

Returns:

Space (e.g. the observation space)

Raises:

ValueError – Cannot batch space does not have a registered function.

Example

>>> from gymnasium.spaces import Box, Dict
>>> import numpy as np
>>> space = Dict({
...     'position': Box(low=0, high=1, shape=(3,), dtype=np.float32),
...     'velocity': Box(low=0, high=1, shape=(2,), dtype=np.float32)
... })
>>> batch_space(space, n=5)
Dict('position': Box(0.0, 1.0, (5, 3), float32), 'velocity': Box(0.0, 1.0, (5, 2), float32))

gymnasium.experimental.vector.utils.concatenate(space: Space, items: Iterable, out: tuple[Any, ...] | dict[str, Any] | np.ndarray) → tuple[Any, ...] | dict[str, Any] | np.ndarray#

gymnasium.experimental.vector.utils.concatenate(space: Box | Discrete | MultiDiscrete | MultiBinary, items: Iterable, out: np.ndarray) → np.ndarray

gymnasium.experimental.vector.utils.concatenate(space: Tuple, items: Iterable, out: tuple[Any, ...]) → tuple[Any, ...]

gymnasium.experimental.vector.utils.concatenate(space: Dict, items: Iterable, out: dict[str, Any]) → dict[str, Any]

gymnasium.experimental.vector.utils.concatenate(space: Space, items: Iterable, out: None) → tuple[Any, ...]

Concatenate multiple samples from space into a single object.

Parameters:

space – Observation space of a single environment in the vectorized environment.
items – Samples to be concatenated.
out – The output object. This object is a (possibly nested) numpy array.

Returns:

The output object. This object is a (possibly nested)

Raises:

ValueError – Space

Example

>>> from gymnasium.spaces import Box
>>> import numpy as np
>>> space = Box(low=0, high=1, shape=(3,), seed=42, dtype=np.float32)
>>> out = np.zeros((2, 3), dtype=np.float32)
>>> items = [space.sample() for _ in range(2)]
>>> concatenate(space, items, out)
array([[0.77395606, 0.43887845, 0.85859793],
       [0.697368  , 0.09417735, 0.97562236]], dtype=float32)

gymnasium.experimental.vector.utils.iterate(space: Space[T_cov], items: Iterable[T_cov]) → Iterator#

gymnasium.experimental.vector.utils.iterate(space: Discrete, items: Iterable)

gymnasium.experimental.vector.utils.iterate(space: Box | MultiDiscrete | MultiBinary, items: np.ndarray)

gymnasium.experimental.vector.utils.iterate(space: Tuple, items: tuple[Any, ...])

gymnasium.experimental.vector.utils.iterate(space: Dict, items: dict[str, Any])

Iterate over the elements of a (batched) space.

Parameters:

space – Observation space of a single environment in the vectorized environment.
items – Samples to be concatenated.
out – The output object. This object is a (possibly nested) numpy array.

Returns:

The output object. This object is a (possibly nested)

Raises:

ValueError – Space is not an instance of gym.Space

Example

>>> from gymnasium.spaces import Box, Dict
>>> import numpy as np
>>> space = Dict({
... 'position': Box(low=0, high=1, shape=(2, 3), seed=42, dtype=np.float32),
... 'velocity': Box(low=0, high=1, shape=(2, 2), seed=42, dtype=np.float32)})
>>> items = space.sample()
>>> it = iterate(space, items)
>>> next(it)
OrderedDict([('position', array([0.77395606, 0.43887845, 0.85859793], dtype=float32)), ('velocity', array([0.77395606, 0.43887845], dtype=float32))])
>>> next(it)
OrderedDict([('position', array([0.697368  , 0.09417735, 0.97562236], dtype=float32)), ('velocity', array([0.85859793, 0.697368  ], dtype=float32))])
>>> next(it)
Traceback (most recent call last):
    ...
StopIteration

gymnasium.experimental.vector.utils.create_empty_array(space: Space, n: int = 1, fn: callable = np.zeros) → tuple[Any, ...] | dict[str, Any] | np.ndarray#

gymnasium.experimental.vector.utils.create_empty_array(space: Box, n: int = 1, fn=np.zeros) → ndarray

gymnasium.experimental.vector.utils.create_empty_array(space: Tuple, n: int = 1, fn=np.zeros) → tuple[Any, ...]

gymnasium.experimental.vector.utils.create_empty_array(space: Dict, n: int = 1, fn=np.zeros) → dict[str, Any]

gymnasium.experimental.vector.utils.create_empty_array(space: Graph, n: int = 1, fn=np.zeros) → tuple[gymnasium.spaces.graph.GraphInstance, ...]

gymnasium.experimental.vector.utils.create_empty_array(space: Text, n: int = 1, fn=np.zeros) → tuple[str, ...]

gymnasium.experimental.vector.utils.create_empty_array(space: Sequence, n: int = 1, fn=np.zeros) → tuple[Any, ...]

gymnasium.experimental.vector.utils.create_empty_array(space: ~gymnasium.spaces.space.Space, n=1, fn=<built-in function zeros>)

Create an empty (possibly nested) (normally numpy-based) array, used in conjunction with concatenate(..., out=array).

In most cases, the array will be contained within the batched space, however, this is not guaranteed.

Parameters:

space – Observation space of a single environment in the vectorized environment.
n – Number of environments in the vectorized environment. If None, creates an empty sample from space.
fn – Function to apply when creating the empty numpy array. Examples of such functions are np.empty or np.zeros.

Returns:

The output object. This object is a (possibly nested)

Raises:

ValueError – Space is not a valid gym.Space instance

Example

>>> from gymnasium.spaces import Box, Dict
>>> import numpy as np
>>> space = Dict({
... 'position': Box(low=0, high=1, shape=(3,), dtype=np.float32),
... 'velocity': Box(low=0, high=1, shape=(2,), dtype=np.float32)})
>>> create_empty_array(space, n=2, fn=np.zeros)
OrderedDict([('position', array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)), ('velocity', array([[0., 0.],
       [0., 0.]], dtype=float32))])

Shared memory functions#

gymnasium.experimental.vector.utils.create_shared_memory(space: Space[Any], n: int = 1, ctx=mp) → dict[str, Any] | tuple[Any, ...] | mp.Array#

gymnasium.experimental.vector.utils.create_shared_memory(space: Box | Discrete | MultiDiscrete | MultiBinary, n: int = 1, ctx=mp)

gymnasium.experimental.vector.utils.create_shared_memory(space: Tuple, n: int = 1, ctx=mp)

gymnasium.experimental.vector.utils.create_shared_memory(space: Dict, n: int = 1, ctx=mp)

gymnasium.experimental.vector.utils.create_shared_memory(space: Text, n: int = 1, ctx=mp)

gymnasium.experimental.vector.utils.create_shared_memory(space: Graph | Sequence, n: int = 1, ctx=mp)

Create a shared memory object, to be shared across processes.

This eventually contains the observations from the vectorized environment.

Parameters:

space – Observation space of a single environment in the vectorized environment.
n – Number of environments in the vectorized environment (i.e. the number of processes).
ctx – The multiprocess module

Returns:

shared_memory for the shared object across processes.

Raises:

CustomSpaceError – Space is not a valid gymnasium.Space instance

gymnasium.experimental.vector.utils.read_from_shared_memory(space: Space, shared_memory: dict | tuple | mp.Array, n: int = 1) → dict[str, Any] | tuple[Any, ...] | np.ndarray#

gymnasium.experimental.vector.utils.read_from_shared_memory(space: Box | Discrete | MultiDiscrete | MultiBinary, shared_memory, n: int = 1)

gymnasium.experimental.vector.utils.read_from_shared_memory(space: Tuple, shared_memory, n: int = 1)

gymnasium.experimental.vector.utils.read_from_shared_memory(space: Dict, shared_memory, n: int = 1)

gymnasium.experimental.vector.utils.read_from_shared_memory(space: Text, shared_memory, n: int = 1) → tuple[str]

Read the batch of observations from shared memory as a numpy array.

..notes::: The numpy array objects returned by read_from_shared_memory shares the memory of shared_memory. Any changes to shared_memory are forwarded to observations, and vice-versa. To avoid any side-effect, use np.copy.

Parameters:

space – Observation space of a single environment in the vectorized environment.
shared_memory – Shared object across processes. This contains the observations from the vectorized environment. This object is created with create_shared_memory.
n – Number of environments in the vectorized environment (i.e. the number of processes).

Returns:

Batch of observations as a (possibly nested)

Raises:

CustomSpaceError – Space is not a valid gymnasium.Space instance

gymnasium.experimental.vector.utils.write_to_shared_memory(space: Space, index: int, value: np.ndarray, shared_memory: dict[str, Any] | tuple[Any, ...] | mp.Array)#

gymnasium.experimental.vector.utils.write_to_shared_memory(space: Box | Discrete | MultiDiscrete | MultiBinary, index: int, value, shared_memory)

gymnasium.experimental.vector.utils.write_to_shared_memory(space: Tuple, index: int, values: tuple[Any, ...], shared_memory)

gymnasium.experimental.vector.utils.write_to_shared_memory(space: Dict, index: int, values: dict[str, Any], shared_memory)

gymnasium.experimental.vector.utils.write_to_shared_memory(space: Text, index: int, values: str, shared_memory)

Write the observation of a single environment into shared memory.

Parameters:

space – Observation space of a single environment in the vectorized environment.
index – Index of the environment (must be in [0, num_envs)).
value – Observation of the single environment to write to shared memory.
shared_memory – Shared object across processes. This contains the observations from the vectorized environment. This object is created with create_shared_memory.

Raises:

CustomSpaceError – Space is not a valid gymnasium.Space instance

Miscellaneous#

gymnasium.experimental.vector.utils.CloudpickleWrapper(fn: Callable[[], Env])#: Wrapper that uses cloudpickle to pickle and unpickle the result.

gymnasium.experimental.vector.utils.clear_mpi_env_vars()#

Clears the MPI of environment variables.

from mpi4py import MPI will call MPI_Init by default. If the child process has MPI environment variables, MPI will think that the child process is an MPI process just like the parent and do bad things such as hang.

This context manager is a hacky way to clear those environment variables temporarily such as when we are starting multiprocessing Processes.

Yields:: Yields for the context manager