[rllib] test all combinations of {obs_space} x {action_space} (#1449)

python/ray/rllib/dqn/dqn_evaluator.py

@@ -2,10 +2,12 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from gym.spaces import Discrete
 import numpy as np
 import tensorflow as tf
 
 import ray
+from ray.rllib.utils.error import UnsupportedSpaceException
 from ray.rllib.dqn import models
 from ray.rllib.dqn.common.wrappers import wrap_dqn
 from ray.rllib.dqn.common.schedules import LinearSchedule
@@ -51,6 +53,11 @@ class DQNEvaluator(TFMultiGPUSupport):
         self.env = env
         self.config = config
 
+        if not isinstance(env.action_space, Discrete):
+            raise UnsupportedSpaceException(
+                "Action space {} is not supported for DQN.".format(
+                    env.action_space))
+
         tf_config = tf.ConfigProto(**config["tf_session_args"])
         self.sess = tf.Session(config=tf_config)
         self.dqn_graph = models.DQNGraph(registry, env, config, logdir)

python/ray/rllib/es/es.py

@@ -38,14 +38,14 @@ DEFAULT_CONFIG = dict(
     num_workers=10,
     stepsize=0.01,
     observation_filter="MeanStdFilter",
+    noise_size=250000000,
     env_config={})
 
 
 @ray.remote
-def create_shared_noise():
+def create_shared_noise(count):
     """Create a large array of noise to be shared by all workers."""
     seed = 123
-    count = 250000000
     noise = np.random.RandomState(seed).randn(count).astype(np.float32)
     return noise
 
@@ -154,7 +154,7 @@ class ESAgent(Agent):
 
         # Create the shared noise table.
         print("Creating shared noise table.")
-        noise_id = create_shared_noise.remote()
+        noise_id = create_shared_noise.remote(self.config["noise_size"])
         self.noise = SharedNoiseTable(ray.get(noise_id))
 
         # Create the actors.

python/ray/rllib/es/policies.py

@@ -27,7 +27,7 @@ def rollout(policy, env, timestep_limit=None, add_noise=False):
     rews = []
     t = 0
     observation = env.reset()
-    for _ in range(timestep_limit):
+    for _ in range(timestep_limit or 999999):
         ac = policy.compute(observation, add_noise=add_noise)[0]
         observation, rew, done, _ = env.step(ac)
         rews.append(rew)

python/ray/rllib/models/action_dist.py

@@ -120,7 +120,7 @@ class MultiActionDistribution(ActionDistribution):
     """
     def __init__(self, inputs, action_space, child_distributions):
         # you actually have to instantiate the child distributions
-        self.reshaper = Reshaper(action_space)
+        self.reshaper = Reshaper(action_space.spaces)
         split_inputs = self.reshaper.split_tensor(inputs)
         child_list = []
         for i, distribution in enumerate(child_distributions):

python/ray/rllib/models/catalog.py

@@ -63,6 +63,10 @@ class ModelCatalog(object):
             dist_dim (int): The size of the input vector to the distribution.
         """
 
+        # TODO(ekl) are list spaces valid?
+        if isinstance(action_space, list):
+            action_space = gym.spaces.Tuple(action_space)
+
         if isinstance(action_space, gym.spaces.Box):
             if dist_type is None:
                 return DiagGaussian, action_space.shape[0] * 2
@@ -70,10 +74,10 @@ class ModelCatalog(object):
                 return Deterministic, action_space.shape[0]
         elif isinstance(action_space, gym.spaces.Discrete):
             return Categorical, action_space.n
-        elif isinstance(action_space, list):
+        elif isinstance(action_space, gym.spaces.Tuple):
             size = 0
             child_dist = []
-            for action in action_space:
+            for action in action_space.spaces:
                 dist, action_size = ModelCatalog.get_action_dist(action)
                 child_dist.append(dist)
                 size += action_size
@@ -94,21 +98,26 @@ class ModelCatalog(object):
             action_placeholder (Tensor): A placeholder for the actions
         """
 
+        # TODO(ekl) are list spaces valid?
+        if isinstance(action_space, list):
+            action_space = gym.spaces.Tuple(action_space)
+
         if isinstance(action_space, gym.spaces.Box):
             return tf.placeholder(
                 tf.float32, shape=(None, action_space.shape[0]))
         elif isinstance(action_space, gym.spaces.Discrete):
             return tf.placeholder(tf.int64, shape=(None,))
-        elif isinstance(action_space, list):
+        elif isinstance(action_space, gym.spaces.Tuple):
             size = 0
-            for i in range(len(action_space)):
-                size += np.product(action_space[i].shape)
-            # TODO(ev) this obviously won't work for mixed spaces
-            if isinstance(action_space[0], gym.spaces.Discrete):
-                return tf.placeholder(tf.int64, shape=(None,
-                                                       len(action_space)))
-            elif isinstance(action_space[0], gym.spaces.Box):
-                return tf.placeholder(tf.float32, shape=(None, size))
+            all_discrete = True
+            for i in range(len(action_space.spaces)):
+                if isinstance(action_space.spaces[i], gym.spaces.Discrete):
+                    size += 1
+                else:
+                    all_discrete = False
+                    size += np.product(action_space.spaces[i].shape)
+            return tf.placeholder(
+                tf.int64 if all_discrete else tf.float32, shape=(None, size))
         else:
             raise NotImplementedError("action space {}"
                                       " not supported".format(action_space))

python/ray/rllib/test/test_supported_spaces.py

@@ -0,0 +1,140 @@
+import unittest
+import traceback
+
+import gym
+from gym.spaces import Box, Discrete, Tuple
+from gym.envs.registration import EnvSpec
+
+import ray
+from ray.rllib.agent import get_agent_class
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.tune.registry import register_env
+
+ACTION_SPACES_TO_TEST = {
+    "discrete": Discrete(5),
+    "vector": Box(0.0, 1.0, (5,)),
+    "simple_tuple": Tuple([Box(0.0, 1.0, (5,)), Box(0.0, 1.0, (5,))]),
+    "implicit_tuple": [Box(0.0, 1.0, (5,)), Box(0.0, 1.0, (5,))],
+}
+
+OBSERVATION_SPACES_TO_TEST = {
+    "discrete": Discrete(5),
+    "vector": Box(0.0, 1.0, (5,)),
+    "image": Box(0.0, 1.0, (80, 80, 1)),
+    "atari": Box(0.0, 1.0, (210, 160, 3)),
+    "atari_ram": Box(0.0, 1.0, (128,)),
+    "simple_tuple": Tuple([Box(0.0, 1.0, (5,)), Box(0.0, 1.0, (5,))]),
+    "mixed_tuple": Tuple([Discrete(10), Box(0.0, 1.0, (5,))]),
+}
+
+# (alg, action_space, obs_space)
+KNOWN_FAILURES = [
+    # TODO(ekl) multiagent support for a3c
+    ("A3C", "implicit_tuple", "atari"),
+    ("A3C", "implicit_tuple", "atari_ram"),
+    ("A3C", "implicit_tuple", "discrete"),
+    ("A3C", "implicit_tuple", "image"),
+    ("A3C", "implicit_tuple", "mixed_tuple"),
+    ("A3C", "implicit_tuple", "simple_tuple"),
+    ("A3C", "implicit_tuple", "vector"),
+    ("A3C", "mixed_tuple", "atari"),
+    ("A3C", "mixed_tuple", "atari_ram"),
+    ("A3C", "mixed_tuple", "discrete"),
+    ("A3C", "mixed_tuple", "image"),
+    ("A3C", "mixed_tuple", "mixed_tuple"),
+    ("A3C", "mixed_tuple", "simple_tuple"),
+    ("A3C", "mixed_tuple", "vector"),
+    ("A3C", "simple_tuple", "atari"),
+    ("A3C", "simple_tuple", "atari_ram"),
+    ("A3C", "simple_tuple", "discrete"),
+    ("A3C", "simple_tuple", "image"),
+    ("A3C", "simple_tuple", "mixed_tuple"),
+    ("A3C", "simple_tuple", "simple_tuple"),
+    ("A3C", "simple_tuple", "vector"),
+]
+
+
+def make_stub_env(action_space, obs_space):
+    class StubEnv(gym.Env):
+        def __init__(self):
+            self.action_space = action_space
+            self.observation_space = obs_space
+            self._spec = EnvSpec("StubEnv-v0")
+
+        def reset(self):
+            sample = self.observation_space.sample()
+            return sample
+
+        def step(self, action):
+            return self.observation_space.sample(), 1, True, {}
+
+    return StubEnv
+
+
+def check_support(alg, config, stats):
+    for a_name, action_space in ACTION_SPACES_TO_TEST.items():
+        for o_name, obs_space in OBSERVATION_SPACES_TO_TEST.items():
+            print("=== Testing", alg, action_space, obs_space, "===")
+            stub_env = make_stub_env(action_space, obs_space)
+            register_env(
+                "stub_env", lambda c: stub_env())
+            stat = "ok"
+            a = None
+            try:
+                a = get_agent_class(alg)(config=config, env="stub_env")
+                a.train()
+            except UnsupportedSpaceException as e:
+                stat = "unsupported"
+            except Exception as e:
+                stat = "ERROR"
+                print(e)
+                print(traceback.format_exc())
+            finally:
+                if a:
+                    try:
+                        a.stop()
+                    except Exception as e:
+                        print("Ignoring error stopping agent", e)
+                        pass
+            print(stat)
+            print()
+            stats[alg, a_name, o_name] = stat
+
+
+class ModelSupportedSpaces(unittest.TestCase):
+    def testAll(self):
+        ray.init()
+        stats = {}
+        check_support("DQN", {"timesteps_per_iteration": 1}, stats)
+        check_support(
+            "A3C", {"num_workers": 1, "optimizer": {"grads_per_step": 1}},
+            stats)
+        check_support(
+            "PPO",
+            {"num_workers": 1, "num_sgd_iter": 1, "timesteps_per_batch": 1,
+             "devices": ["/cpu:0"], "min_steps_per_task": 1,
+             "sgd_batchsize": 1},
+            stats)
+        check_support(
+            "ES",
+            {"num_workers": 1, "noise_size": 10000000,
+             "episodes_per_batch": 1, "timesteps_per_batch": 1},
+            stats)
+        num_unexpected_errors = 0
+        num_unexpected_success = 0
+        for (alg, a_name, o_name), stat in sorted(stats.items()):
+            if stat in ["ok", "unsupported"]:
+                if (alg, a_name, o_name) in KNOWN_FAILURES:
+                    num_unexpected_success += 1
+            else:
+                if (alg, a_name, o_name) not in KNOWN_FAILURES:
+                    num_unexpected_errors += 1
+            print(
+                alg, "action_space", a_name, "obs_space", o_name,
+                "result", stat)
+        self.assertEqual(num_unexpected_errors, 0)
+        self.assertEqual(num_unexpected_success, 0)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

python/ray/rllib/utils/error.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+class UnsupportedSpaceException(Exception):
+    """Error for an unsupported action or observation space."""
+    pass

python/ray/rllib/utils/reshaper.py

@@ -16,6 +16,8 @@ class Reshaper(object):
                 # Handle both gym arrays and just lists of inputs length
                 if hasattr(space, "shape"):
                     arr_shape = np.asarray(space.shape)
+                elif hasattr(space, "n"):
+                    arr_shape = np.asarray([1])  # discrete space
                 else:
                     arr_shape = space
                 self.shapes.append(arr_shape)

python/ray/rllib/utils/sampler.py

@@ -200,8 +200,8 @@ def _env_runner(env, policy, num_local_steps, horizon, obs_filter):
             "wrapper_config.TimeLimit.max_episode_steps")
     except Exception:
         print("Warning, no horizon specified, assuming infinite")
+    if not horizon:
         horizon = 999999
-    assert horizon > 0
     if hasattr(policy, "get_initial_features"):
         last_features = policy.get_initial_features()
     else: