[RLlib] Preprocessor fixes (multi-discrete) and tests. (#13083)

rllib/BUILD

@@ -1089,6 +1089,13 @@ py_test(
     srcs = ["models/tests/test_distributions.py"]
 )
 
+py_test(
+    name = "test_preprocessors",
+    tags = ["models"],
+    size = "small",
+    srcs = ["models/tests/test_preprocessors.py"]
+)
+
 # --------------------------------------------------------------------
 # Evaluation components
 # rllib/evaluation/

rllib/models/modelv2.py

@@ -438,9 +438,7 @@ def _unpack_obs(obs: TensorType, space: gym.Space,
         tensorlib: The library used to unflatten (reshape) the array/tensor
     """
 
-    if (isinstance(space, gym.spaces.Dict)
-            or isinstance(space, gym.spaces.Tuple)
-            or isinstance(space, Repeated)):
+    if isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple, Repeated)):
         if id(space) in _cache:
             prep = _cache[id(space)]
         else:

rllib/models/preprocessors.py

@@ -79,7 +79,7 @@ class Preprocessor:
     def observation_space(self) -> gym.Space:
         obs_space = gym.spaces.Box(-1., 1., self.shape, dtype=np.float32)
         # Stash the unwrapped space so that we can unwrap dict and tuple spaces
-        # automatically in model.py
+        # automatically in modelv2.py
         classes = (DictFlatteningPreprocessor, OneHotPreprocessor,
                    RepeatedValuesPreprocessor, TupleFlatteningPreprocessor)
         if isinstance(self, classes):
@@ -141,15 +141,31 @@ class AtariRamPreprocessor(Preprocessor):
 
 
 class OneHotPreprocessor(Preprocessor):
+    """One-hot preprocessor for Discrete and MultiDiscrete spaces.
+
+    Examples:
+        >>> self.transform(Discrete(3).sample())
+        ... np.array([0.0, 1.0, 0.0])
+        >>> self.transform(MultiDiscrete([2, 3]).sample())
+        ... np.array([0.0, 1.0, 0.0, 0.0, 1.0])
+    """
+
     @override(Preprocessor)
     def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]:
-        return (self._obs_space.n, )
+        if isinstance(obs_space, gym.spaces.Discrete):
+            return (self._obs_space.n, )
+        else:
+            return (np.sum(self._obs_space.nvec), )
 
     @override(Preprocessor)
     def transform(self, observation: TensorType) -> np.ndarray:
         self.check_shape(observation)
-        arr = np.zeros(self._obs_space.n, dtype=np.float32)
-        arr[observation] = 1
+        arr = np.zeros(self._init_shape(self._obs_space, {}), dtype=np.float32)
+        if isinstance(self._obs_space, gym.spaces.Discrete):
+            arr[observation] = 1
+        else:
+            for i, o in enumerate(observation):
+                arr[np.sum(self._obs_space.nvec[:i]) + o] = 1
         return arr
 
     @override(Preprocessor)
@@ -299,7 +315,7 @@ def get_preprocessor(space: gym.Space) -> type:
     legacy_patch_shapes(space)
     obs_shape = space.shape
 
-    if isinstance(space, gym.spaces.Discrete):
+    if isinstance(space, (gym.spaces.Discrete, gym.spaces.MultiDiscrete)):
         preprocessor = OneHotPreprocessor
     elif obs_shape == ATARI_OBS_SHAPE:
         preprocessor = GenericPixelPreprocessor

rllib/models/tests/test_preprocessors.py

@@ -0,0 +1,78 @@
+import gym
+from gym.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
+import numpy as np
+import unittest
+
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.models.preprocessors import DictFlatteningPreprocessor, \
+    get_preprocessor, NoPreprocessor, TupleFlatteningPreprocessor, \
+    OneHotPreprocessor, AtariRamPreprocessor, GenericPixelPreprocessor
+from ray.rllib.utils.test_utils import check
+
+
+class TestPreprocessors(unittest.TestCase):
+    def test_gym_preprocessors(self):
+        p1 = ModelCatalog.get_preprocessor(gym.make("CartPole-v0"))
+        self.assertEqual(type(p1), NoPreprocessor)
+
+        p2 = ModelCatalog.get_preprocessor(gym.make("FrozenLake-v0"))
+        self.assertEqual(type(p2), OneHotPreprocessor)
+
+        p3 = ModelCatalog.get_preprocessor(gym.make("MsPacman-ram-v0"))
+        self.assertEqual(type(p3), AtariRamPreprocessor)
+
+        p4 = ModelCatalog.get_preprocessor(gym.make("MsPacmanNoFrameskip-v4"))
+        self.assertEqual(type(p4), GenericPixelPreprocessor)
+
+    def test_tuple_preprocessor(self):
+        class TupleEnv:
+            def __init__(self):
+                self.observation_space = Tuple(
+                    [Discrete(5),
+                     Box(0, 5, shape=(3, ), dtype=np.float32)])
+
+        pp = ModelCatalog.get_preprocessor(TupleEnv())
+        self.assertTrue(isinstance(pp, TupleFlatteningPreprocessor))
+        self.assertEqual(pp.shape, (8, ))
+        self.assertEqual(
+            list(pp.transform((0, np.array([1, 2, 3])))),
+            [float(x) for x in [1, 0, 0, 0, 0, 1, 2, 3]])
+
+    def test_dict_flattening_preprocessor(self):
+        space = Dict({
+            "a": Discrete(2),
+            "b": Tuple([Discrete(3), Box(-1.0, 1.0, (4, ))]),
+        })
+        pp = get_preprocessor(space)(space)
+        self.assertTrue(isinstance(pp, DictFlatteningPreprocessor))
+        self.assertEqual(pp.shape, (9, ))
+        check(
+            pp.transform({
+                "a": 1,
+                "b": (1, np.array([0.0, -0.5, 0.1, 0.6]))
+            }), [0.0, 1.0, 0.0, 1.0, 0.0, 0.0, -0.5, 0.1, 0.6])
+
+    def test_one_hot_preprocessor(self):
+        space = Discrete(5)
+        pp = get_preprocessor(space)(space)
+        self.assertTrue(isinstance(pp, OneHotPreprocessor))
+        self.assertTrue(pp.shape == (5, ))
+        check(pp.transform(3), [0.0, 0.0, 0.0, 1.0, 0.0])
+        check(pp.transform(0), [1.0, 0.0, 0.0, 0.0, 0.0])
+
+        space = MultiDiscrete([2, 3, 4])
+        pp = get_preprocessor(space)(space)
+        self.assertTrue(isinstance(pp, OneHotPreprocessor))
+        self.assertTrue(pp.shape == (9, ))
+        check(
+            pp.transform(np.array([1, 2, 0])),
+            [0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0])
+        check(
+            pp.transform(np.array([0, 1, 3])),
+            [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/policy/torch_policy.py

@@ -212,12 +212,11 @@ class TorchPolicy(Policy):
 
         if self.action_sampler_fn:
             action_dist = dist_inputs = None
-            state_out = state_batches
             actions, logp, state_out = self.action_sampler_fn(
                 self,
                 self.model,
                 input_dict,
-                state_out,
+                state_batches,
                 explore=explore,
                 timestep=timestep)
         else:

rllib/tests/test_catalog.py

@@ -1,5 +1,5 @@
 import gym
-from gym.spaces import Box, Discrete, Tuple
+from gym.spaces import Box, Discrete
 import numpy as np
 import unittest
 
@@ -7,8 +7,7 @@ import ray
 from ray.rllib.models import ModelCatalog, MODEL_DEFAULTS, ActionDistribution
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
-from ray.rllib.models.preprocessors import (NoPreprocessor, OneHotPreprocessor,
-                                            Preprocessor)
+from ray.rllib.models.preprocessors import NoPreprocessor, Preprocessor
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.test_utils import framework_iterator
@@ -65,28 +64,6 @@ class TestModelCatalog(unittest.TestCase):
     def tearDown(self):
         ray.shutdown()
 
-    def test_gym_preprocessors(self):
-        p1 = ModelCatalog.get_preprocessor(gym.make("CartPole-v0"))
-        self.assertEqual(type(p1), NoPreprocessor)
-
-        p2 = ModelCatalog.get_preprocessor(gym.make("FrozenLake-v0"))
-        self.assertEqual(type(p2), OneHotPreprocessor)
-
-    def test_tuple_preprocessor(self):
-        ray.init(object_store_memory=1000 * 1024 * 1024)
-
-        class TupleEnv:
-            def __init__(self):
-                self.observation_space = Tuple(
-                    [Discrete(5),
-                     Box(0, 5, shape=(3, ), dtype=np.float32)])
-
-        p1 = ModelCatalog.get_preprocessor(TupleEnv())
-        self.assertEqual(p1.shape, (8, ))
-        self.assertEqual(
-            list(p1.transform((0, np.array([1, 2, 3])))),
-            [float(x) for x in [1, 0, 0, 0, 0, 1, 2, 3]])
-
     def test_custom_preprocessor(self):
         ray.init(object_store_memory=1000 * 1024 * 1024)
         ModelCatalog.register_custom_preprocessor("foo", CustomPreprocessor)