Source code for akro.dict
"""Cartesian product of multiple named Spaces (also known as a dict of Spaces).
This Space produces samples which are dicts, where the values of those dicts
are drawn from the values of this Space.
"""
import collections
import gym.spaces
import numpy as np
import akro
from akro.requires import requires_tf, requires_theano
from akro.space import Space
[docs]class Dict(gym.spaces.Dict, Space):
"""A dictionary of simpler spaces, e.g. Discrete, Box.
Example usage:
self.observation_space = spaces.Dict({"position": spaces.Discrete(2),
"velocity": spaces.Discrete(3)})
"""
def __init__(self, spaces=None, **kwargs):
super().__init__(spaces, **kwargs)
self.spaces = (collections.OrderedDict([
(k, akro.from_gym(s)) for k, s in self.spaces.items()
]))
@property
def flat_dim(self):
"""Return the length of the flattened vector of the space."""
return sum([space.flat_dim for _, space in self.spaces.items()])
[docs] def flat_dim_with_keys(self, keys):
"""Return a flat dimension of the spaces specified by the keys.
Returns:
sum (int)
"""
return sum([self.spaces[key].flat_dim for key in keys])
[docs] def flatten(self, x):
"""Return an observation of x with collapsed values.
Args:
x (:obj:`Iterable`): The object to flatten.
Returns:
Dict: A Dict where each value is collapsed into a single dimension.
Keys are unchanged.
"""
return np.concatenate(
[space.flatten(x[key]) for key, space in self.spaces.items()],
axis=-1,
)
[docs] def unflatten(self, x):
"""Return an unflattened observation x.
Args:
x (:obj:`Iterable`): The object to unflatten.
Returns:
collections.OrderedDict
"""
dims = np.array([s.flat_dim for s in self.spaces.values()])
flat_x = np.split(x, np.cumsum(dims)[:-1])
return collections.OrderedDict([
(key, self.spaces[key].unflatten(xi))
for key, xi in zip(self.spaces.keys(), flat_x)
])
[docs] def flatten_n(self, xs):
"""Return flattened observations xs.
Args:
xs (:obj:`Iterable`): The object to reshape and flatten
Returns:
np.ndarray: An array of xs in a shape inferred by the size of
its first element.
"""
return np.array([self.flatten(x) for x in xs])
[docs] def unflatten_n(self, xs):
"""Return unflattened observations xs.
Args:
xs (:obj:`Iterable`): The object to reshape and unflatten
Returns:
List[OrderedDict]
"""
return [self.unflatten(x) for x in xs]
[docs] def flatten_with_keys(self, x, keys):
"""Return flattened obs of spaces specified by the keys using x.
Returns:
list
"""
return np.concatenate(
[
space.flatten(x[key])
for key, space in self.spaces.items() if key in keys
],
axis=-1,
)
[docs] def unflatten_with_keys(self, x, keys):
"""Return an unflattened observation.
This is the inverse of `flatten_with_keys`.
Returns:
collections.OrderedDict
"""
dims = np.array([
space.flat_dim for key, space in self.spaces.items() if key in keys
])
flat_x = np.split(x, np.cumsum(dims)[:-1])
return collections.OrderedDict([
(key, space.unflatten(xi))
for (key, space), xi in zip(self.spaces.items(), flat_x)
if key in keys
])
[docs] @requires_tf
def to_tf_placeholder(self, name, batch_dims):
"""Create a tensor placeholder from the Space object.
Args:
name (str): name of the variable
batch_dims (:obj:`list`): batch dimensions to add to the
shape of the object.
Returns:
tf.Tensor: Tensor object with the same properties as
the Dict where the shape is modified by batch_dims.
"""
newdict = Dict()
for key, space in self.spaces.items():
newdict.spaces[key] = space.to_tf_placeholder(name, batch_dims)
return newdict
[docs] @requires_theano
def to_theano_tensor(self, name, batch_dims):
"""Create a theano tensor from the Space object.
Args:
name (str): name of the variable
batch_dims (:obj:`list`): batch dimensions to add to the
shape of the object.
Returns:
theano.tensor.TensorVariable: Tensor object with the
same properties as the Dict where the shape is
modified by batch_dims.
"""
newdict = Dict()
for key, space in self.spaces.items():
newdict.spaces[key] = space.to_theano_tensor(name, batch_dims)
return newdict