[RLlib] Add light-weight Trainer.compute_action() tests for all Algos. (#8356)

.travis.yml

@@ -201,7 +201,8 @@ matrix:
         - travis_wait 60 bazel test --config=ci --build_tests_only --test_tag_filters=quick_train rllib/...
         # Test everything that does not have any of the "main" labels:
         # "learning_tests|quick_train|examples|tests_dir".
-        - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=-learning_tests_tf,-learning_tests_torch,-quick_train,-examples,-tests_dir rllib/...
+        #- ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=-learning_tests_tf,-learning_tests_torch,-quick_train,-examples,-tests_dir rllib/...
+        - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=agents_dir_X rllib/...
 
     # RLlib: Everything in rllib/examples/ directory.
     - os: linux

rllib/BUILD

@@ -115,6 +115,14 @@ py_test(
     srcs = ["agents/ddpg/tests/test_apex_ddpg.py"]
 )
 
+# ARS
+py_test(
+    name = "test_ars",
+    tags = ["agents_dir_X"],
+    size = "medium",
+    srcs = ["agents/ars/tests/test_ars.py"]
+)
+
 # DDPGTrainer
 py_test(
     name = "test_ddpg",
@@ -137,19 +145,27 @@ py_test(
     srcs = ["agents/dqn/tests/test_simple_q.py"]
 )
 
-# IMPALA
+# ES
 py_test(
-    name = "test_vtrace",
+    name = "test_es",
     tags = ["agents_dir"],
-    size = "small",
-    srcs = ["agents/impala/tests/test_vtrace.py"]
+    size = "medium",
+    srcs = ["agents/es/tests/test_es.py"]
 )
+
+# IMPALA
 py_test(
     name = "test_impala",
     tags = ["agents_dir"],
     size = "medium",
     srcs = ["agents/impala/tests/test_impala.py"]
 )
+py_test(
+    name = "test_vtrace",
+    tags = ["agents_dir"],
+    size = "small",
+    srcs = ["agents/impala/tests/test_vtrace.py"]
+)
 
 # MARWILTrainer
 py_test(

rllib/agents/a3c/tests/test_a2c.py

@@ -2,6 +2,7 @@ import unittest
 
 import ray
 from ray.rllib.agents.a3c import A2CTrainer
+from ray.rllib.utils.test_utils import check_compute_action
 
 
 class TestA2C(unittest.TestCase):
@@ -21,6 +22,7 @@ class TestA2C(unittest.TestCase):
                 "use_exec_api": True
             })
         assert isinstance(trainer.train(), dict)
+        check_compute_action(trainer)
 
     def test_a2c_exec_impl_microbatch(ray_start_regular):
         trainer = A2CTrainer(
@@ -31,6 +33,7 @@ class TestA2C(unittest.TestCase):
                 "use_exec_api": True,
             })
         assert isinstance(trainer.train(), dict)
+        check_compute_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ars/ars.py

@@ -313,7 +313,10 @@ class ARSTrainer(Trainer):
 
     @override(Trainer)
     def compute_action(self, observation, *args, **kwargs):
-        return self.policy.compute_actions(observation, update=True)[0]
+        action = self.policy.compute_actions(observation, update=True)[0]
+        if kwargs.get("full_fetch"):
+            return action, [], {}
+        return action
 
     def _collect_results(self, theta_id, min_episodes):
         num_episodes, num_timesteps = 0, 0

rllib/agents/ars/tests/test_ars.py

@@ -0,0 +1,32 @@
+import unittest
+
+import ray
+import ray.rllib.agents.ars as ars
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
+
+
+class TestARS(unittest.TestCase):
+    def test_ars_compilation(self):
+        """Test whether an ARSTrainer can be built on all frameworks."""
+        ray.init(num_cpus=2, local_mode=True)
+        config = ars.DEFAULT_CONFIG.copy()
+        # Keep it simple.
+        config["model"]["fcnet_hiddens"] = [10]
+        config["model"]["fcnet_activation"] = None
+
+        num_iterations = 2
+
+        for _ in framework_iterator(config, ("torch", "tf")):
+            plain_config = config.copy()
+            trainer = ars.ARSTrainer(config=plain_config, env="CartPole-v0")
+            for i in range(num_iterations):
+                results = trainer.train()
+                print(results)
+
+            check_compute_action(trainer)
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/agents/ddpg/tests/test_apex_ddpg.py

@@ -3,7 +3,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.ddpg.apex as apex_ddpg
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 
 class TestApexDDPG(unittest.TestCase):
@@ -40,6 +41,7 @@ class TestApexDDPG(unittest.TestCase):
 
             for _ in range(num_iterations):
                 print(trainer.train())
+            check_compute_action(trainer)
 
             # Test again per-worker scale distribution
             # (should not have changed).

rllib/agents/ddpg/tests/test_ddpg.py

@@ -10,7 +10,8 @@ from ray.rllib.optimizers.async_replay_optimizer import LocalReplayBuffer
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.numpy import fc, huber_loss, l2_loss, relu, sigmoid
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 from ray.rllib.utils.torch_ops import convert_to_torch_tensor
 
 tf = try_import_tf()
@@ -32,6 +33,7 @@ class TestDDPG(unittest.TestCase):
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
+            check_compute_action(trainer)
 
     def test_ddpg_exploration_and_with_random_prerun(self):
         """Tests DDPG's Exploration (w/ random actions for n timesteps)."""

rllib/agents/ddpg/tests/test_td3.py

@@ -3,7 +3,8 @@ import unittest
 
 import ray.rllib.agents.ddpg.td3 as td3
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 tf = try_import_tf()
 
@@ -21,6 +22,7 @@ class TestTD3(unittest.TestCase):
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)
+            check_compute_action(trainer)
 
     def test_td3_exploration_and_with_random_prerun(self):
         """Tests TD3's Exploration (w/ random actions for n timesteps)."""

rllib/agents/dqn/tests/test_apex_dqn.py

@@ -3,7 +3,8 @@ import unittest
 
 import ray
 import ray.rllib.agents.dqn.apex as apex
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 
 class TestApexDQN(unittest.TestCase):
@@ -32,6 +33,8 @@ class TestApexDQN(unittest.TestCase):
             expected = [0.4, 0.016190862, 0.00065536]
             check([i["cur_epsilon"] for i in infos], [0.0] + expected)
 
+            check_compute_action(trainer)
+
             # TODO(ekl) fix iterator metrics bugs w/multiple trainers.
             #            for i in range(1):
             #                results = trainer.train()

rllib/agents/dqn/tests/test_dqn.py

@@ -3,7 +3,8 @@ import unittest
 
 import ray.rllib.agents.dqn as dqn
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 tf = try_import_tf()
 
@@ -23,6 +24,8 @@ class TestDQN(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
+            check_compute_action(trainer)
+
             # Rainbow.
             # TODO(sven): Add torch once DQN-torch supports distributional-Q.
             if fw == "torch":
@@ -38,6 +41,8 @@ class TestDQN(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
+            check_compute_action(trainer)
+
     def test_dqn_exploration_and_soft_q_config(self):
         """Tests, whether a DQN Agent outputs exploration/softmaxed actions."""
         config = dqn.DEFAULT_CONFIG.copy()

rllib/agents/dqn/tests/test_simple_q.py

@@ -8,7 +8,8 @@ from ray.rllib.agents.dqn.simple_q_torch_policy import build_q_losses as \
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.numpy import fc, one_hot, huber_loss
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 tf = try_import_tf()
 
@@ -26,6 +27,8 @@ class TestSimpleQ(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
+            check_compute_action(trainer)
+
     def test_simple_q_loss_function(self):
         """Tests the Simple-Q loss function results on all frameworks."""
         config = dqn.SIMPLE_Q_DEFAULT_CONFIG.copy()

rllib/agents/es/es.py

@@ -295,7 +295,10 @@ class ESTrainer(Trainer):
 
     @override(Trainer)
     def compute_action(self, observation, *args, **kwargs):
-        return self.policy.compute_actions(observation, update=False)[0]
+        action = self.policy.compute_actions(observation, update=False)[0]
+        if kwargs.get("full_fetch"):
+            return action, [], {}
+        return action
 
     @override(Trainer)
     def _stop(self):

rllib/agents/es/es_tf_policy.py

@@ -68,6 +68,7 @@ def make_session(single_threaded):
 
 class ESTFPolicy:
     def __init__(self, obs_space, action_space, config):
+        self.observation_space = obs_space
         self.action_space = action_space
         self.action_space_struct = get_base_struct_from_space(action_space)
         self.action_noise_std = config["action_noise_std"]

rllib/agents/es/tests/test_es.py

@@ -2,10 +2,7 @@ import unittest
 
 import ray
 import ray.rllib.agents.es as es
-from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator
-
-tf = try_import_tf()
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
 
 
 class TestES(unittest.TestCase):
@@ -26,6 +23,8 @@ class TestES(unittest.TestCase):
                 results = trainer.train()
                 print(results)
 
+            check_compute_action(trainer)
+
 
 if __name__ == "__main__":
     import pytest

rllib/agents/impala/tests/test_impala.py

@@ -3,7 +3,7 @@ import unittest
 import ray
 import ray.rllib.agents.impala as impala
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
 
 tf = try_import_tf()
 
@@ -31,6 +31,7 @@ class TestIMPALA(unittest.TestCase):
                 trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
                 for i in range(num_iterations):
                     print(trainer.train())
+                check_compute_action(trainer)
 
                 # Test w/ LSTM.
                 print("w/o LSTM")
@@ -38,6 +39,7 @@ class TestIMPALA(unittest.TestCase):
                 trainer = impala.ImpalaTrainer(config=local_cfg, env=env)
                 for i in range(num_iterations):
                     print(trainer.train())
+                check_compute_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/marwil/tests/test_marwil.py

@@ -3,7 +3,7 @@ import unittest
 import ray
 import ray.rllib.agents.marwil as marwil
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
 
 tf = try_import_tf()
 
@@ -28,6 +28,7 @@ class TestMARWIL(unittest.TestCase):
             trainer = marwil.MARWILTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
+            check_compute_action(trainer, include_prev_action_reward=True)
 
 
 if __name__ == "__main__":

rllib/agents/pg/tests/test_pg.py

@@ -7,7 +7,7 @@ from ray.rllib.evaluation.postprocessing import Postprocessing
 from ray.rllib.models.tf.tf_action_dist import Categorical
 from ray.rllib.models.torch.torch_action_dist import TorchCategorical
 from ray.rllib.policy.sample_batch import SampleBatch
-from ray.rllib.utils import check, fc, framework_iterator
+from ray.rllib.utils import check, fc, framework_iterator, check_compute_action
 
 
 class TestPG(unittest.TestCase):
@@ -25,6 +25,7 @@ class TestPG(unittest.TestCase):
                 "use_exec_api": True
             })
         assert isinstance(trainer.train(), dict)
+        check_compute_action(trainer)
 
     def test_pg_compilation(self):
         """Test whether a PGTrainer can be built with both frameworks."""
@@ -36,6 +37,7 @@ class TestPG(unittest.TestCase):
             trainer = pg.PGTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
+            check_compute_action(trainer, include_prev_action_reward=True)
 
     def test_pg_loss_functions(self):
         """Tests the PG loss function math."""

rllib/agents/ppo/tests/test_appo.py

@@ -3,7 +3,7 @@ import unittest
 import ray
 import ray.rllib.agents.ppo as ppo
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
 
 tf = try_import_tf()
 
@@ -28,12 +28,14 @@ class TestAPPO(unittest.TestCase):
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")
             for i in range(num_iterations):
                 print(trainer.train())
+            check_compute_action(trainer)
 
             _config = config.copy()
             _config["vtrace"] = True
             trainer = ppo.APPOTrainer(config=_config, env="CartPole-v0")
             for i in range(num_iterations):
                 print(trainer.train())
+            check_compute_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ppo/tests/test_ddppo.py

@@ -3,7 +3,7 @@ import unittest
 import ray
 import ray.rllib.agents.ppo as ppo
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.test_utils import framework_iterator
+from ray.rllib.utils.test_utils import framework_iterator, check_compute_action
 
 tf = try_import_tf()
 
@@ -27,6 +27,7 @@ class TestDDPPO(unittest.TestCase):
             trainer = ppo.ddppo.DDPPOTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
+            check_compute_action(trainer)
 
 
 if __name__ == "__main__":

rllib/agents/ppo/tests/test_ppo.py

@@ -14,7 +14,8 @@ from ray.rllib.models.torch.torch_action_dist import TorchCategorical
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.numpy import fc
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 
 tf = try_import_tf()
 
@@ -38,6 +39,7 @@ class TestPPO(unittest.TestCase):
             trainer = ppo.PPOTrainer(config=config, env="CartPole-v0")
             for i in range(num_iterations):
                 trainer.train()
+            check_compute_action(trainer, include_prev_action_reward=True)
 
     def test_ppo_fake_multi_gpu_learning(self):
         """Test whether PPOTrainer can learn CartPole w/ faked multi-GPU."""

rllib/agents/sac/tests/test_sac.py

@@ -13,7 +13,8 @@ from ray.rllib.optimizers.async_replay_optimizer import LocalReplayBuffer
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.numpy import fc, relu
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 from ray.rllib.utils.torch_ops import convert_to_torch_tensor
 
 tf = try_import_tf()
@@ -66,6 +67,7 @@ class TestSAC(unittest.TestCase):
                 for i in range(num_iterations):
                     results = trainer.train()
                     print(results)
+                check_compute_action(trainer)
 
     def test_sac_loss_function(self):
         """Tests SAC loss function results across all frameworks."""

rllib/utils/__init__.py

@@ -13,7 +13,8 @@ from ray.rllib.utils.policy_client import PolicyClient
 from ray.rllib.utils.policy_server import PolicyServer
 from ray.rllib.utils.schedules import LinearSchedule, PiecewiseSchedule, \
     PolynomialSchedule, ExponentialSchedule, ConstantSchedule
-from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.utils.test_utils import check, framework_iterator, \
+    check_compute_action
 from ray.tune.utils import merge_dicts, deep_update
 
 
@@ -70,6 +71,7 @@ def try_import_tree():
 __all__ = [
     "add_mixins",
     "check",
+    "check_compute_action",
     "check_framework",
     "deprecation_warning",
     "fc",

rllib/utils/test_utils.py

@@ -243,3 +243,41 @@ def check_learning_achieved(tune_results, min_reward):
     if tune_results.trials[0].last_result["episode_reward_mean"] < min_reward:
         raise ValueError("`stop-reward` of {} not reached!".format(min_reward))
     print("ok")
+
+
+def check_compute_action(trainer, include_prev_action_reward=False):
+    """Tests different combinations of arguments for trainer.compute_action.
+
+    Args:
+        trainer (Trainer): The trainer object to test.
+        include_prev_action_reward (bool): Whether to include the prev-action
+            and -reward in the `compute_action` call.
+
+    Throws:
+        ValueError: If anything unexpected happens.
+    """
+    try:
+        pol = trainer.get_policy()
+    except AttributeError:
+        pol = trainer.policy
+
+    obs_space = pol.observation_space
+    action_space = pol.action_space
+    for explore in [True, False]:
+        for full_fetch in [True, False]:
+            obs = np.clip(obs_space.sample(), -1.0, 1.0)
+            action_in = action_space.sample() \
+                if include_prev_action_reward else None
+            reward_in = 1.0 if include_prev_action_reward else None
+            action = trainer.compute_action(
+                obs,
+                prev_action=action_in,
+                prev_reward=reward_in,
+                explore=explore,
+                full_fetch=full_fetch)
+            if full_fetch:
+                action, _, _ = action
+            if not action_space.contains(action):
+                raise ValueError(
+                    "Returned action ({}) of trainer {} not in Env's "
+                    "action_space ({})!".format(action, trainer, action_space))