Fix akro.Dict to support existing APIs

Several of the APIs were missing from `akro.Dict` (and untested).
I've copied their implementations from the old `akro.tf.Dict`.

I also rewrote `akro.Dict.unflatten_from_keys` to make more sense.
This shouldn't be a problem, since it isn't used in garage (but is
provided for completeness).

Wrap sub-spaces in Dict and Tuple spaces

Pin gym==0.12.4

We require given gym version, since otherwise gym.spaces.Space is not
defined, which causes akro.Space to fail to import.

Increase version to v0.0.6

Author: krzentner

VERSION

@@ -1 +1 @@
-v0.0.5-dev
+v0.0.6-dev

setup.py

@@ -4,7 +4,7 @@
 # Required dependencies
 required = [
     # Please keep alphabetized
-    'gym',
+    'gym==0.12.4',
     'numpy',
 ]
 

src/akro/dict.py

@@ -3,8 +3,12 @@
 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
 
@@ -18,12 +22,27 @@ class Dict(gym.spaces.Dict, Space):
                                               "velocity": spaces.Discrete(3)})
     """
 
-    @property  # pragma: no cover
+    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."""
-        raise NotImplementedError
+        return sum([space.flat_dim for _, space in self.spaces.items()])
+
+    def flat_dim_with_keys(self, keys):
+        """
+        Return a flat dimension of the spaces specified by the keys.
+
+        Returns:
+            sum (int)
 
-    def flatten(self, x):  # pragma: no cover
+        """
+        return sum([self.spaces[key].flat_dim for key in keys])
+
+    def flatten(self, x):
         """Return an observation of x with collapsed values.
 
         Args:
@@ -34,21 +53,31 @@ def flatten(self, x):  # pragma: no cover
                   Keys are unchanged.
 
         """
-        raise NotImplementedError
-
-    def unflatten(self, x):  # pragma: no cover
+        return np.concatenate(
+            [
+                space.flatten(xi)
+                for space, xi in zip(self.spaces.values(), x.values())
+            ],
+            axis=-1,
+        )
+
+    def unflatten(self, x):
         """Return an unflattened observation x.
 
         Args:
             x (:obj:`Iterable`): The object to unflatten.
 
         Returns:
-            np.ndarray: An array of x in the shape of self.shape.
+            collections.OrderedDict
 
         """
-        raise NotImplementedError
+        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)])
 
-    def flatten_n(self, xs):  # pragma: no cover
+    def flatten_n(self, xs):
         """Return flattened observations xs.
 
         Args:
@@ -59,20 +88,51 @@ def flatten_n(self, xs):  # pragma: no cover
                 its first element.
 
         """
-        raise NotImplementedError
+        return np.array([self.flatten(x) for x in xs])
 
-    def unflatten_n(self, xs):  # pragma: no cover
+    def unflatten_n(self, xs):
         """Return unflattened observations xs.
 
         Args:
             xs (:obj:`Iterable`): The object to reshape and unflatten
 
         Returns:
-            np.ndarray: An array of xs in a shape inferred by the size of
-                its first element and self.shape.
+            List[OrderedDict]
+
+        """
+        return [self.unflatten(x) for x in xs]
+
+    def flatten_with_keys(self, x, keys):
+        """
+        Return flattened obs of spaces specified by the keys using x.
+
+        Returns:
+            list
+
+        """
+        return np.concatenate(
+            [
+                self.spaces[key].flatten(xi)
+                for key, xi in zip(self.spaces.keys(), x.values())
+                if key in keys
+            ],
+            axis=-1,
+        )
+
+    def unflatten_with_keys(self, x, keys):
+        """
+        Return an unflattened observation.
+
+        This is the inverse of `flatten_with_keys`.
+
+        Returns:
+            collections.OrderedDict
 
         """
-        raise NotImplementedError
+        dims = np.array([self.spaces[key].flat_dim for key in keys])
+        flat_x = np.split(x, np.cumsum(dims)[:-1])
+        return collections.OrderedDict([(key, self.spaces[key].unflatten(xi))
+                                        for key, xi in zip(keys, flat_x)])
 
     @requires_tf
     def to_tf_placeholder(self, name, batch_dims):

src/akro/tuple.py

@@ -7,13 +7,17 @@
 import gym.spaces
 import numpy as np
 
+import akro
 from akro.requires import requires_tf, requires_theano
 from akro.space import Space
 
 
 class Tuple(gym.spaces.Tuple, Space):
     """A Tuple of Spaces which produces samples which are Tuples of samples."""
 
+    def __init__(self, spaces):
+        super().__init__([akro.from_gym(space) for space in spaces])
+
     @property
     def flat_dim(self):
         """Return the length of the flattened vector of the space."""

tests/akro/test_dict.py

@@ -1,8 +1,12 @@
+import collections
 import pickle
 import unittest
 
+import numpy as np
+
 from akro import Dict
 from akro import Discrete
+from akro import Box
 from akro import tf
 from akro import theano
 from akro.requires import requires_tf, requires_theano
@@ -23,19 +27,70 @@ def test_pickleable(self):
         assert round_trip.contains(sample)
 
     def test_flat_dim(self):
-        pass
+        d = Dict(collections.OrderedDict(position=Box(0, 10, (2,)),
+                                         velocity=Box(0, 10, (3,))))
+        assert d.flat_dim == 5
+
+    def test_flat_dim_with_keys(self):
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        assert d.flat_dim_with_keys(['position']) == 2
 
     def test_flatten(self):
-        pass
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([1., 2., 3., 4., 5.])
+        s = collections.OrderedDict(position=np.array([1., 2.]),
+                                    velocity=np.array([3., 4., 5.]))
+        assert (d.flatten(s) == f).all()
 
     def test_unflatten(self):
-        pass
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([1., 2., 3., 4., 5.])
+        s = collections.OrderedDict(position=np.array([1., 2.]),
+                                    velocity=np.array([3., 4., 5.]))
+        assert all((s[k] == v).all() for k, v in d.unflatten(f).items())
 
     def test_flatten_n(self):
-        pass
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([[1., 2., 3., 4., 5.],
+                      [6., 7., 8., 9., 0.]])
+        s = [collections.OrderedDict(position=np.array([1., 2.]),
+                                     velocity=np.array([3., 4., 5.])),
+             collections.OrderedDict(position=np.array([6., 7.]),
+                                     velocity=np.array([8., 9., 0.]))]
+        assert (d.flatten_n(s) == f).all()
 
     def test_unflatten_n(self):
-        pass
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([[1., 2., 3., 4., 5.],
+                      [6., 7., 8., 9., 0.]])
+        s = [collections.OrderedDict(position=np.array([1., 2.]),
+                                     velocity=np.array([3., 4., 5.])),
+             collections.OrderedDict(position=np.array([6., 7.]),
+                                     velocity=np.array([8., 9., 0.]))]
+        for i, fi in enumerate(d.unflatten_n(f)):
+            assert all((s[i][k] == v).all() for k, v in fi.items())
+
+    def test_flatten_with_keys(self):
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([3., 4., 5.])
+        s = collections.OrderedDict(position=np.array([1., 2.]),
+                                    velocity=np.array([3., 4., 5.]))
+        assert (d.flatten_with_keys(s, ['velocity']) == f).all()
+
+    def test_unflatten_with_keys(self):
+        d = Dict(collections.OrderedDict([('position', Box(0, 10, (2,))),
+                                          ('velocity', Box(0, 10, (3,)))]))
+        f = np.array([3., 4., 5.])
+        s = collections.OrderedDict(position=np.array([1., 2.]),
+                                    velocity=np.array([3., 4., 5.]))
+        assert all((s[k] == v).all()
+                   for k, v in d.unflatten_with_keys(f, ['velocity']).items())
 
     @requires_tf
     def test_convert_tf(self):