[RLlib] Add all agents to rllib rollout tests. (#7534)

rllib/BUILD

@@ -553,14 +553,6 @@ py_test(
         ]
 )
 
-py_test(
-    name = "test_impala_rollout",
-    main = "tests/test_rollout.py",
-    data = ["train.py", "rollout.py"],
-    tags = ["quick_train"],
-    srcs = ["tests/test_rollout.py"]
-)
-
 # MARWIL
 
 py_test(
@@ -1073,6 +1065,15 @@ py_test(
     srcs = ["tests/test_reproducibility.py"]
 )
 
+py_test(
+    name = "test_rollout",
+    main = "tests/test_rollout.py",
+    tags = ["tests_dir", "tests_dir_R"],
+    size = "large",
+    data = ["train.py", "rollout.py"],
+    srcs = ["tests/test_rollout.py"]
+)
+
 py_test(
     name = "tests/test_rollout_worker",
     tags = ["tests_dir", "tests_dir_R"],

rllib/tests/test_nested_spaces.py

@@ -1,8 +1,7 @@
-import pickle
-
-import gym
 from gym import spaces
 from gym.envs.registration import EnvSpec
+import gym
+import pickle
 import unittest
 
 import ray

rllib/tests/test_rollout.py

@@ -1,51 +1,86 @@
-# Simple translation of former test_rollout.sh file to be able
-# to run this in bazel test suite.
-
 from pathlib import Path
 import os
+import sys
 import unittest
 
 
+def rollout_test(algo, env="CartPole-v0"):
+    tmp_dir = os.popen("mktemp -d").read()[:-1]
+    if not os.path.exists(tmp_dir):
+        sys.exit(1)
+
+    print("Saving results to {}".format(tmp_dir))
+
+    rllib_dir = str(Path(__file__).parent.parent.absolute())
+    print("RLlib dir = {}\nexists={}".format(rllib_dir,
+                                             os.path.exists(rllib_dir)))
+    os.system("python {}/train.py --local-dir={} --run={} "
+              "--checkpoint-freq=1 ".format(rllib_dir, tmp_dir, algo) +
+              "--config='{\"num_workers\": 1, \"num_gpus\": 0}' "
+              "--stop='{\"training_iteration\": 1}'" + " --env={}".format(env))
+
+    checkpoint_path = os.popen(
+        "ls {}/default/*/checkpoint_1/checkpoint-1".format(tmp_dir)).read()[:
+                                                                            -1]
+    if not os.path.exists(checkpoint_path):
+        sys.exit(1)
+    print("Checkpoint path {} (exists)".format(checkpoint_path))
+
+    # Test rolling out n steps.
+    os.popen("python {}/rollout.py --run={} \"{}\" --steps=25 "
+             "--out=\"{}/rollouts_25steps.pkl\" --no-render".format(
+                 rllib_dir, algo, checkpoint_path, tmp_dir)).read()
+    if not os.path.exists(tmp_dir + "/rollouts_25steps.pkl"):
+        sys.exit(1)
+    print("rollout output (25 steps) exists!".format(checkpoint_path))
+
+    # Test rolling out 1 episode.
+    os.popen("python {}/rollout.py --run={} \"{}\" --episodes=1 "
+             "--out=\"{}/rollouts_1episode.pkl\" --no-render".format(
+                 rllib_dir, algo, checkpoint_path, tmp_dir)).read()
+    if not os.path.exists(tmp_dir + "/rollouts_1episode.pkl"):
+        sys.exit(1)
+    print("rollout output (1 ep) exists!".format(checkpoint_path))
+
+    # Cleanup.
+    os.popen("rm -rf \"{}\"".format(tmp_dir)).read()
+
+
 class TestRollout(unittest.TestCase):
-    def test_rollout(self):
-        tmp_dir = os.popen("mktemp -d").read()[:-1]
-        if not os.path.exists(tmp_dir):
-            sys.exit(1)
+    def test_a2c(self):
+        rollout_test("A2C")
 
-        print("Saving results to {}".format(tmp_dir))
+    def test_a3c(self):
+        rollout_test("A3C")
 
-        rllib_dir = str(Path(__file__).parent.parent.absolute())
-        print("RLlib dir = {}\nexists={}".format(rllib_dir,
-                                                 os.path.exists(rllib_dir)))
-        os.system("python {}/train.py --local-dir={} --run=IMPALA "
-                  "--checkpoint-freq=1 ".format(rllib_dir, tmp_dir) +
-                  "--config='{\"num_workers\": 1, \"num_gpus\": 0}' "
-                  "--env=Pong-ram-v4 --stop='{\"training_iteration\": 1}'")
+    def test_ars(self):
+        rollout_test("ARS")
 
-        checkpoint_path = os.popen(
-            "ls {}/default/*/checkpoint_1/checkpoint-1".format(
-                tmp_dir)).read()[:-1]
-        print("Checkpoint path {}".format(checkpoint_path))
-        if not os.path.exists(checkpoint_path):
-            sys.exit(1)
+    def test_ddpg(self):
+        rollout_test("DDPG", env="Pendulum-v0")
 
-        os.popen("python {}/rollout.py --run=IMPALA \"{}\" --steps=100 "
-                 "--out=\"{}/rollouts_100steps.pkl\" --no-render".format(
-                     rllib_dir, checkpoint_path, tmp_dir)).read()
-        if not os.path.exists(tmp_dir + "/rollouts_100steps.pkl"):
-            sys.exit(1)
+    def test_dqn(self):
+        rollout_test("DQN")
 
-        os.popen("python {}/rollout.py --run=IMPALA \"{}\" --episodes=1 "
-                 "--out=\"{}/rollouts_1episode.pkl\" --no-render".format(
-                     rllib_dir, checkpoint_path, tmp_dir)).read()
-        if not os.path.exists(tmp_dir + "/rollouts_1episode.pkl"):
-            sys.exit(1)
+    def test_es(self):
+        rollout_test("ES")
 
-        # Cleanup.
-        os.popen("rm -rf \"{}\"".format(tmp_dir)).read()
+    def test_impala(self):
+        rollout_test("IMPALA", env="Pong-ram-v4")
+
+    def test_pg(self):
+        rollout_test("PG")
+
+    def test_ppo(self):
+        rollout_test("PPO", env="Pendulum-v0")
+
+    def test_sac(self):
+        rollout_test("SAC", env="Pendulum-v0")
+
+    def test_td3(self):
+        rollout_test("TD3", env="Pendulum-v0")
 
 
 if __name__ == "__main__":
     import pytest
-    import sys
     sys.exit(pytest.main(["-v", __file__]))

rllib/utils/exploration/random.py

@@ -4,7 +4,7 @@ from typing import Union
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
-    tf_function, TensorType
+    TensorType
 from ray.rllib.utils.tuple_actions import TupleActions
 from ray.rllib.models.modelv2 import ModelV2
 
@@ -52,17 +52,23 @@ class Random(Exploration):
         else:
             return self.get_torch_exploration_action(action_dist, explore)
 
-    @tf_function(tf)
     def get_tf_exploration_action_op(self, action_dist, explore):
-        if explore:
+        def true_fn():
             action = tf.py_function(self.action_space.sample, [],
                                     self.dtype_sample)
             # Will be unnecessary, once we support batch/time-aware Spaces.
-            action = tf.expand_dims(tf.cast(action, dtype=self.dtype), 0)
-        else:
-            action = tf.cast(
+            return tf.expand_dims(tf.cast(action, dtype=self.dtype), 0)
+
+        def false_fn():
+            return tf.cast(
                 action_dist.deterministic_sample(), dtype=self.dtype)
 
+        action = tf.cond(
+            pred=tf.constant(explore, dtype=tf.bool)
+            if isinstance(explore, bool) else explore,
+            true_fn=true_fn,
+            false_fn=false_fn)
+
         # TODO(sven): Move into (deterministic_)sample(logp=True|False)
         if isinstance(action, TupleActions):
             batch_size = tf.shape(action[0][0])[0]