mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-06 16:01:45 +00:00
344 lines
10 KiB
Python
344 lines
10 KiB
Python
"""Implementation of utility functions that can be applied to spaces.
|
|
|
|
These functions mostly take care of flattening and unflattening elements of spaces
|
|
to facilitate their usage in learning code.
|
|
"""
|
|
import operator as op
|
|
from collections import OrderedDict
|
|
from functools import reduce, singledispatch
|
|
from typing import TypeVar, Union, cast
|
|
|
|
import numpy as np
|
|
|
|
from gym.spaces import (
|
|
Box,
|
|
Dict,
|
|
Discrete,
|
|
Graph,
|
|
GraphInstance,
|
|
MultiBinary,
|
|
MultiDiscrete,
|
|
Space,
|
|
Tuple,
|
|
)
|
|
|
|
|
|
@singledispatch
|
|
def flatdim(space: Space) -> int:
|
|
"""Return the number of dimensions a flattened equivalent of this space would have.
|
|
|
|
Example usage::
|
|
|
|
>>> from gym.spaces import Discrete
|
|
>>> space = Dict({"position": Discrete(2), "velocity": Discrete(3)})
|
|
>>> flatdim(space)
|
|
5
|
|
|
|
Args:
|
|
space: The space to return the number of dimensions of the flattened spaces
|
|
|
|
Returns:
|
|
The number of dimensions for the flattened spaces
|
|
|
|
Raises:
|
|
NotImplementedError: if the space is not defined in ``gym.spaces``.
|
|
"""
|
|
raise NotImplementedError(f"Unknown space: `{space}`")
|
|
|
|
|
|
@flatdim.register(Box)
|
|
@flatdim.register(MultiBinary)
|
|
def _flatdim_box_multibinary(space: Union[Box, MultiBinary]) -> int:
|
|
return reduce(op.mul, space.shape, 1)
|
|
|
|
|
|
@flatdim.register(Discrete)
|
|
def _flatdim_discrete(space: Discrete) -> int:
|
|
return int(space.n)
|
|
|
|
|
|
@flatdim.register(MultiDiscrete)
|
|
def _flatdim_multidiscrete(space: MultiDiscrete) -> int:
|
|
return int(np.sum(space.nvec))
|
|
|
|
|
|
@flatdim.register(Tuple)
|
|
def _flatdim_tuple(space: Tuple) -> int:
|
|
return sum(flatdim(s) for s in space.spaces)
|
|
|
|
|
|
@flatdim.register(Dict)
|
|
def _flatdim_dict(space: Dict) -> int:
|
|
return sum(flatdim(s) for s in space.spaces.values())
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@singledispatch
|
|
def flatten(space: Space[T], x: T) -> np.ndarray:
|
|
"""Flatten a data point from a space.
|
|
|
|
This is useful when e.g. points from spaces must be passed to a neural
|
|
network, which only understands flat arrays of floats.
|
|
|
|
Args:
|
|
space: The space that ``x`` is flattened by
|
|
x: The value to flatten
|
|
|
|
Returns:
|
|
- The flattened ``x``, always returns a 1D array for non-graph spaces.
|
|
- For graph spaces, returns `GraphInstance` where:
|
|
- `nodes` are n x k arrays
|
|
- `edges` are either:
|
|
- m x k arrays
|
|
- None
|
|
- `edge_links` are either:
|
|
- m x 2 arrays
|
|
- None
|
|
|
|
Raises:
|
|
NotImplementedError: If the space is not defined in ``gym.spaces``.
|
|
"""
|
|
raise NotImplementedError(f"Unknown space: `{space}`")
|
|
|
|
|
|
@flatten.register(Box)
|
|
@flatten.register(MultiBinary)
|
|
def _flatten_box_multibinary(space, x) -> np.ndarray:
|
|
return np.asarray(x, dtype=space.dtype).flatten()
|
|
|
|
|
|
@flatten.register(Discrete)
|
|
def _flatten_discrete(space, x) -> np.ndarray:
|
|
onehot = np.zeros(space.n, dtype=space.dtype)
|
|
onehot[x - space.start] = 1
|
|
return onehot
|
|
|
|
|
|
@flatten.register(MultiDiscrete)
|
|
def _flatten_multidiscrete(space, x) -> np.ndarray:
|
|
offsets = np.zeros((space.nvec.size + 1,), dtype=space.dtype)
|
|
offsets[1:] = np.cumsum(space.nvec.flatten())
|
|
|
|
onehot = np.zeros((offsets[-1],), dtype=space.dtype)
|
|
onehot[offsets[:-1] + x.flatten()] = 1
|
|
return onehot
|
|
|
|
|
|
@flatten.register(Tuple)
|
|
def _flatten_tuple(space, x) -> np.ndarray:
|
|
return np.concatenate([flatten(s, x_part) for x_part, s in zip(x, space.spaces)])
|
|
|
|
|
|
@flatten.register(Dict)
|
|
def _flatten_dict(space, x) -> np.ndarray:
|
|
return np.concatenate([flatten(s, x[key]) for key, s in space.spaces.items()])
|
|
|
|
|
|
@flatten.register(Graph)
|
|
def _flatten_graph(space, x) -> np.ndarray:
|
|
"""We're not using `.unflatten() for :class:`Box` and :class:`Discrete` because a graph is not a homogeneous space, see `.flatten` docstring."""
|
|
|
|
def _graph_unflatten(space, x):
|
|
ret = None
|
|
if space is not None and x is not None:
|
|
if isinstance(space, Box):
|
|
ret = x.reshape(x.shape[0], -1)
|
|
elif isinstance(space, Discrete):
|
|
ret = np.zeros((x.shape[0], space.n - space.start), dtype=space.dtype)
|
|
ret[np.arange(x.shape[0]), x - space.start] = 1
|
|
return ret
|
|
|
|
nodes = _graph_unflatten(space.node_space, x.nodes)
|
|
edges = _graph_unflatten(space.edge_space, x.edges)
|
|
|
|
return GraphInstance(nodes, edges, x.edge_links)
|
|
|
|
|
|
@singledispatch
|
|
def unflatten(space: Space[T], x: np.ndarray) -> T:
|
|
"""Unflatten a data point from a space.
|
|
|
|
This reverses the transformation applied by :func:`flatten`. You must ensure
|
|
that the ``space`` argument is the same as for the :func:`flatten` call.
|
|
|
|
Args:
|
|
space: The space used to unflatten ``x``
|
|
x: The array to unflatten
|
|
|
|
Returns:
|
|
A point with a structure that matches the space.
|
|
|
|
Raises:
|
|
NotImplementedError: if the space is not defined in ``gym.spaces``.
|
|
"""
|
|
raise NotImplementedError(f"Unknown space: `{space}`")
|
|
|
|
|
|
@unflatten.register(Box)
|
|
@unflatten.register(MultiBinary)
|
|
def _unflatten_box_multibinary(
|
|
space: Union[Box, MultiBinary], x: np.ndarray
|
|
) -> np.ndarray:
|
|
return np.asarray(x, dtype=space.dtype).reshape(space.shape)
|
|
|
|
|
|
@unflatten.register(Discrete)
|
|
def _unflatten_discrete(space: Discrete, x: np.ndarray) -> int:
|
|
return int(space.start + np.nonzero(x)[0][0])
|
|
|
|
|
|
@unflatten.register(MultiDiscrete)
|
|
def _unflatten_multidiscrete(space: MultiDiscrete, x: np.ndarray) -> np.ndarray:
|
|
offsets = np.zeros((space.nvec.size + 1,), dtype=space.dtype)
|
|
offsets[1:] = np.cumsum(space.nvec.flatten())
|
|
|
|
(indices,) = cast(type(offsets[:-1]), np.nonzero(x))
|
|
return np.asarray(indices - offsets[:-1], dtype=space.dtype).reshape(space.shape)
|
|
|
|
|
|
@unflatten.register(Tuple)
|
|
def _unflatten_tuple(space: Tuple, x: np.ndarray) -> tuple:
|
|
dims = np.asarray([flatdim(s) for s in space.spaces], dtype=np.int_)
|
|
list_flattened = np.split(x, np.cumsum(dims[:-1]))
|
|
return tuple(
|
|
unflatten(s, flattened) for flattened, s in zip(list_flattened, space.spaces)
|
|
)
|
|
|
|
|
|
@unflatten.register(Dict)
|
|
def _unflatten_dict(space: Dict, x: np.ndarray) -> dict:
|
|
dims = np.asarray([flatdim(s) for s in space.spaces.values()], dtype=np.int_)
|
|
list_flattened = np.split(x, np.cumsum(dims[:-1]))
|
|
return OrderedDict(
|
|
[
|
|
(key, unflatten(s, flattened))
|
|
for flattened, (key, s) in zip(list_flattened, space.spaces.items())
|
|
]
|
|
)
|
|
|
|
|
|
@unflatten.register(Graph)
|
|
def _unflatten_graph(space: Graph, x: GraphInstance) -> GraphInstance:
|
|
"""We're not using `.unflatten() for :class:`Box` and :class:`Discrete` because a graph is not a homogeneous space.
|
|
|
|
The size of the outcome is actually not fixed, but determined based on the number of
|
|
nodes and edges in the graph.
|
|
"""
|
|
|
|
def _graph_unflatten(space, x):
|
|
ret = None
|
|
if space is not None and x is not None:
|
|
if isinstance(space, Box):
|
|
ret = x.reshape(-1, *space.shape)
|
|
elif isinstance(space, Discrete):
|
|
ret = np.asarray(np.nonzero(x))[-1, :]
|
|
return ret
|
|
|
|
nodes = _graph_unflatten(space.node_space, x.nodes)
|
|
edges = _graph_unflatten(space.edge_space, x.edges)
|
|
|
|
return GraphInstance(nodes, edges, x.edge_links)
|
|
|
|
|
|
@singledispatch
|
|
def flatten_space(space: Space) -> Box:
|
|
"""Flatten a space into a single ``Box``.
|
|
|
|
This is equivalent to :func:`flatten`, but operates on the space itself. The
|
|
result for non-graph spaces is always a `Box` with flat boundaries. While
|
|
the result for graph spaces is always a `Graph` with `node_space` being a `Box`
|
|
with flat boundaries and `edge_space` being a `Box` with flat boundaries or
|
|
`None`. The box has exactly :func:`flatdim` dimensions. Flattening a sample
|
|
of the original space has the same effect as taking a sample of the flattenend
|
|
space.
|
|
|
|
Example::
|
|
|
|
>>> box = Box(0.0, 1.0, shape=(3, 4, 5))
|
|
>>> box
|
|
Box(3, 4, 5)
|
|
>>> flatten_space(box)
|
|
Box(60,)
|
|
>>> flatten(box, box.sample()) in flatten_space(box)
|
|
True
|
|
|
|
Example that flattens a discrete space::
|
|
|
|
>>> discrete = Discrete(5)
|
|
>>> flatten_space(discrete)
|
|
Box(5,)
|
|
>>> flatten(box, box.sample()) in flatten_space(box)
|
|
True
|
|
|
|
Example that recursively flattens a dict::
|
|
|
|
>>> space = Dict({"position": Discrete(2), "velocity": Box(0, 1, shape=(2, 2))})
|
|
>>> flatten_space(space)
|
|
Box(6,)
|
|
>>> flatten(space, space.sample()) in flatten_space(space)
|
|
True
|
|
|
|
|
|
Example that flattens a graph::
|
|
|
|
>>> space = Graph(node_space=Box(low=-100, high=100, shape=(3, 4)), edge_space=Discrete(5))
|
|
>>> flatten_space(space)
|
|
Graph(Box(-100.0, 100.0, (12,), float32), Box(0, 1, (5,), int64))
|
|
>>> flatten(space, space.sample()) in flatten_space(space)
|
|
True
|
|
|
|
Args:
|
|
space: The space to flatten
|
|
|
|
Returns:
|
|
A flattened Box
|
|
|
|
Raises:
|
|
NotImplementedError: if the space is not defined in ``gym.spaces``.
|
|
"""
|
|
raise NotImplementedError(f"Unknown space: `{space}`")
|
|
|
|
|
|
@flatten_space.register(Box)
|
|
def _flatten_space_box(space: Box) -> Box:
|
|
return Box(space.low.flatten(), space.high.flatten(), dtype=space.dtype)
|
|
|
|
|
|
@flatten_space.register(Discrete)
|
|
@flatten_space.register(MultiBinary)
|
|
@flatten_space.register(MultiDiscrete)
|
|
def _flatten_space_binary(space: Union[Discrete, MultiBinary, MultiDiscrete]) -> Box:
|
|
return Box(low=0, high=1, shape=(flatdim(space),), dtype=space.dtype)
|
|
|
|
|
|
@flatten_space.register(Tuple)
|
|
def _flatten_space_tuple(space: Tuple) -> Box:
|
|
space_list = [flatten_space(s) for s in space.spaces]
|
|
return Box(
|
|
low=np.concatenate([s.low for s in space_list]),
|
|
high=np.concatenate([s.high for s in space_list]),
|
|
dtype=np.result_type(*[s.dtype for s in space_list]),
|
|
)
|
|
|
|
|
|
@flatten_space.register(Dict)
|
|
def _flatten_space_dict(space: Dict) -> Box:
|
|
space_list = [flatten_space(s) for s in space.spaces.values()]
|
|
return Box(
|
|
low=np.concatenate([s.low for s in space_list]),
|
|
high=np.concatenate([s.high for s in space_list]),
|
|
dtype=np.result_type(*[s.dtype for s in space_list]),
|
|
)
|
|
|
|
|
|
@flatten_space.register(Graph)
|
|
def _flatten_space_graph(space: Graph) -> Graph:
|
|
return Graph(
|
|
node_space=flatten_space(space.node_space),
|
|
edge_space=flatten_space(space.edge_space)
|
|
if space.edge_space is not None
|
|
else None,
|
|
)
|