test prep

python/ray/rllib/models/preprocessors.py

@@ -148,8 +148,8 @@ def _init(self):
             print("Custom tuple preprocessor shape is", self.shape)
 
         def transform(self, observation):
-            return tuple(
-                [o.transform(p) for (o, p)
-                 in zip(observation, self.preprocessors)])
+            return np.concatenate([
+                p.transform(o)
+                for (o, p) in zip(observation, self.preprocessors)])
 
     return CustomTuplePreprocessor

python/ray/rllib/test/test_catalog.py

@@ -2,14 +2,15 @@
 import numpy as np
 import tensorflow as tf
 import unittest
+from gym.spaces import Box, Discrete, Tuple
 
 import ray
 from ray.tune.registry import get_registry
 
 from ray.rllib.models import ModelCatalog
 from ray.rllib.models.model import Model
 from ray.rllib.models.preprocessors import (
-    NoPreprocessor, OneHotPreprocessor, Preprocessor)
+    NoPreprocessor, OneHotPreprocessor, Preprocessor, make_tuple_preprocessor)
 from ray.rllib.models.fcnet import FullyConnectedNetwork
 from ray.rllib.models.visionnet import VisionNetwork
 
@@ -34,11 +35,21 @@ def tearDown(self):
     def testGymPreprocessors(self):
         p1 = ModelCatalog.get_preprocessor(
             get_registry(), gym.make("CartPole-v0"))
-        assert type(p1) == NoPreprocessor
+        self.assertEqual(type(p1), NoPreprocessor)
 
         p2 = ModelCatalog.get_preprocessor(
             get_registry(), gym.make("FrozenLake-v0"))
-        assert type(p2) == OneHotPreprocessor
+        self.assertEqual(type(p2), OneHotPreprocessor)
+
+    def testTuplePreprocessor(self):
+        cls = make_tuple_preprocessor(
+            [Discrete(5), Box(0, 1, shape=(3,))],
+            [OneHotPreprocessor, NoPreprocessor])
+        preprocessor = cls(Tuple([Discrete(5), Box(0, 1, shape=(3,))]), {})
+        self.assertEqual(preprocessor.shape, (8,))
+        self.assertEqual(
+            list(preprocessor.transform((0, [1, 2, 3]))),
+            [float(x) for x in [1, 0, 0, 0, 0, 1, 2, 3]])
 
     def testCustomPreprocessor(self):
         ray.init()
@@ -47,32 +58,32 @@ def testCustomPreprocessor(self):
         env = gym.make("CartPole-v0")
         p1 = ModelCatalog.get_preprocessor(
             get_registry(), env, {"custom_preprocessor": "foo"})
-        assert type(p1) == CustomPreprocessor
+        self.assertEqual(str(type(p1)), str(CustomPreprocessor))
         p2 = ModelCatalog.get_preprocessor(
             get_registry(), env, {"custom_preprocessor": "bar"})
-        assert type(p2) == CustomPreprocessor2
+        self.assertEqual(str(type(p2)), str(CustomPreprocessor2))
         p3 = ModelCatalog.get_preprocessor(get_registry(), env)
-        assert type(p3) == NoPreprocessor
+        self.assertEqual(type(p3), NoPreprocessor)
 
     def testDefaultModels(self):
         ray.init()
 
         with tf.variable_scope("test1"):
             p1 = ModelCatalog.get_model(
                 get_registry(), np.zeros((10, 3), dtype=np.float32), 5)
-            assert type(p1) == FullyConnectedNetwork
+            self.assertEqual(type(p1), FullyConnectedNetwork)
 
         with tf.variable_scope("test2"):
             p2 = ModelCatalog.get_model(
                 get_registry(), np.zeros((10, 80, 80, 3), dtype=np.float32), 5)
-            assert type(p2) == VisionNetwork
+            self.assertEqual(type(p2), VisionNetwork)
 
     def testCustomModel(self):
         ray.init()
         ModelCatalog.register_custom_model("foo", CustomModel)
         p1 = ModelCatalog.get_model(
             get_registry(), 1, 5, {"custom_model": "foo"})
-        assert type(p1) == CustomModel
+        self.assertEqual(str(type(p1)), str(CustomModel))
 
 
 if __name__ == "__main__":

python/ray/tune/registry.py

@@ -103,6 +103,7 @@ def flush_values_to_object_store(self):
                 obj = ray.put(_to_pinnable(v))
                 self._all_objects[k] = obj
                 self._refs.append(ray.get(obj))
+        assert len(self._refs) == len(self._all_objects)
 
 
 _default_registry = _Registry()