From a5318961dec403c821889e5e510dd13e9a928f8e Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Sat, 26 Dec 2020 20:14:36 -0500 Subject: [PATCH] [RLlib] Preprocessor fixes (multi-discrete) and tests. (#13083) --- rllib/BUILD | 7 +++ rllib/models/modelv2.py | 4 +- rllib/models/preprocessors.py | 26 ++++++-- rllib/models/tests/test_preprocessors.py | 78 ++++++++++++++++++++++++ rllib/policy/torch_policy.py | 3 +- rllib/tests/test_catalog.py | 27 +------- 6 files changed, 110 insertions(+), 35 deletions(-) create mode 100644 rllib/models/tests/test_preprocessors.py diff --git a/rllib/BUILD b/rllib/BUILD index 44a147b6d..c5a109555 100644 --- a/rllib/BUILD +++ b/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/ diff --git a/rllib/models/modelv2.py b/rllib/models/modelv2.py index 85d991294..a6c871d0f 100644 --- a/rllib/models/modelv2.py +++ b/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: diff --git a/rllib/models/preprocessors.py b/rllib/models/preprocessors.py index c1b7803c4..d3cff7a99 100644 --- a/rllib/models/preprocessors.py +++ b/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 diff --git a/rllib/models/tests/test_preprocessors.py b/rllib/models/tests/test_preprocessors.py new file mode 100644 index 000000000..5515b6fea --- /dev/null +++ b/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__])) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 10e875d50..e7a1c69ad 100644 --- a/rllib/policy/torch_policy.py +++ b/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: diff --git a/rllib/tests/test_catalog.py b/rllib/tests/test_catalog.py index d08c6e6b3..b98f7143a 100644 --- a/rllib/tests/test_catalog.py +++ b/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)