eagerx/core/space.py
import typing as t
import gymnasium.spaces as spaces
import numpy as np
from functools import reduce
import operator as op
class Space(spaces.Space):
"""
A (possibly unbounded) space in R^n. Specifically, a Space represents the
Cartesian product of n closed intervals. Each interval has the form of one
of [a, b], (-oo, b], [a, oo), or (-oo, oo).
There are two common use cases:
* Identical bound for each dimension::
>>> Space(low=-1.0, high=2.0, shape=(3, 4), dtype="float32")
Space(3, 4)
* Independent bound for each dimension::
>>> Space(low=np.array([-1.0, -2.0]), high=np.array([2.0, 4.0]), dtype="float32")
Space(2,)
"""
def __init__(
self,
low: t.Optional[t.Union[t.List, np.ndarray, int, float, bool, t.Tuple]] = None,
high: t.Optional[t.Union[t.List, np.ndarray, int, float, bool, t.Tuple]] = None,
shape: t.Optional[t.Tuple] = None,
dtype: t.Union[np.number, str] = np.float32,
seed: t.Optional[int] = None,
):
"""Constructor for initializing a space.
:param low: Expected lower bound of messages.
:param high: Expected upper bound of messages.
:param shape: Expected shape of messages.
:param dtype: Expected datatype of messages.
:param seed: Seed that is used to draw random samples from the space.
"""
if isinstance(low, (list, tuple)):
low = np.array(low, dtype=dtype)
if isinstance(high, (list, tuple)):
high = np.array(high, dtype=dtype)
# determine shape if it isn't provided directly
if shape is not None:
shape = tuple(shape)
assert low is None or np.isscalar(low) or low.shape == shape, "low.shape doesn't match provided shape"
assert high is None or np.isscalar(high) or high.shape == shape, "high.shape doesn't match provided shape"
elif low is not None and not np.isscalar(low):
shape = low.shape
assert high is not None and np.isscalar(high) or high.shape == shape, "high.shape doesn't match low.shape"
elif high is not None and not np.isscalar(high):
shape = high.shape
assert low is not None and np.isscalar(low) or low.shape == shape, "low.shape doesn't match high.shape"
elif shape is None:
pass
else:
raise ValueError("shape must be provided or inferred from the shapes of low or high")
if np.isscalar(low):
low = np.full(shape, low, dtype=dtype)
if np.isscalar(high):
high = np.full(shape, high, dtype=dtype)
if low is not None and high is not None:
self._space = spaces.Box(low=low, high=high, shape=shape, dtype=dtype, seed=seed)
else:
self._low = low
self._high = high
self._space = None
super().__init__(shape=shape, dtype=dtype, seed=seed)
def __getattr__(self, name):
if self._space is None:
raise AttributeError("Attribute does not exist.")
else:
return getattr(self._space, name)
def contains(self, x) -> bool:
"""
Return boolean specifying if x is a valid
member of this space
:param x: array to check.
"""
if self._space is None:
# todo: convert to ndarray? Maybe not to test that dtype is correct && we are dealing with np array?
# if not isinstance(x, np.ndarray):
# x = np.asarray(x, dtype=self.dtype)
return x.dtype == self.dtype and (self.shape is None or x.shape == self.shape)
else:
return x.dtype == self.dtype and self._space.contains(x) # Check dtype separately
def contains_space(self, space: t.Union["Space", t.Dict]) -> bool:
"""
Return boolean specifying if space is contained
in this space.
Low and high of the space must exactly match (instead of lying within the bounds) to return True
:param space: Space that is to be checked.
"""
if isinstance(space, dict):
space = Space.from_dict(space)
# Check dtype
flag = space.dtype == self.dtype
# Check shape
if flag and space.shape is not None and self.shape is not None:
flag = flag and space.shape == self.shape
# Check bounds
if flag and space._space is not None and self._space is not None:
flag = flag and (space._space.low == self._space.low).all()
flag = flag and (space._space.high == self._space.high).all()
return flag
def sample(self) -> np.ndarray:
"""Randomly sample an element of this space. Can be
uniform or non-uniform sampling based on boundedness of space."""
if not self.is_fully_defined:
raise NotImplementedError("Cannot sample from this incomplete space (i.e. low, high not provided?).")
else:
return self._space.sample()
def to_dict(self) -> t.Dict:
"""Convert the space to a dict representation
:return: Dict representation of the space.
"""
if self._space is None:
if self.shape is None:
return dict(low=None, high=None, shape=None, dtype=self.dtype.name)
else:
return dict(low=None, high=None, shape=list(self.shape), dtype=self.dtype.name)
else:
low = np.nan_to_num(self.low)
high = np.nan_to_num(self.high)
if np.all(low == low.reshape(-1)[0]):
low = low.reshape(-1)[0].item()
if np.all(high == high.reshape(-1)[0]):
high = high.reshape(-1)[0].item()
if isinstance(low, np.ndarray):
low = low.tolist()
if isinstance(high, np.ndarray):
high = high.tolist()
return dict(low=low, high=high, shape=list(self.shape), dtype=self.dtype.name)
@property
def is_np_flattenable(self):
"""Checks whether this space can be flattened."""
return True
@property
def is_fully_defined(self) -> bool:
"""Check if space is fully defined (i.e. low, high, shape and dtype are all provided).
:return: flag
"""
return self._space is not None
@classmethod
def from_dict(cls, d: t.Dict) -> "Space":
"""Create a space from a dict.
:param d: Dict containing the arguments to initialize the space
:return: The space.
"""
return cls(low=d["low"], high=d["high"], shape=d["shape"], dtype=d["dtype"])
def __repr__(self):
if self._space is None:
if self.shape is None:
return f"Space({self.dtype})"
else:
return f"Space({self.shape}, {self.dtype})"
else:
return f"Space({self.low}, {self.high}, {self.shape}, {self.dtype})"
def __eq__(self, other):
if self._space is None:
return (
isinstance(other, Space)
and other._space is None
and (self.shape == other.shape)
and (self.dtype == other.dtype)
)
else:
return (
isinstance(other, spaces.Box)
and (self.shape == other.shape)
and np.allclose(self.low, other.low)
and np.allclose(self.high, other.high)
)
# Register a method to flatten our custom Space using standard gym functions.
@spaces.utils.flatten_space.register(Space)
def flatten_space_box(space):
return spaces.Box(space.low.flatten(), space.high.flatten(), dtype=space.dtype)
# Register a method to unflatten our custom Space using standard gym functions.
@spaces.utils.unflatten.register(Space)
def unflatten_box(space, x):
return np.asarray(x, dtype=space.dtype).reshape(space.shape)
@spaces.utils.flatdim.register(Space)
def flatdim_box_multibinary(space):
return reduce(op.mul, space.shape, 1)
@spaces.utils.flatten.register(Space)
def flatten_box_multibinary(space, x):
return np.asarray(x, dtype=space.dtype).flatten()