[RLlib] Cleanup/unify all test cases. (#7533)

rllib/BUILD

@@ -553,7 +553,6 @@ py_test(
         ]
 )
 
-# From test_rollout.sh (deprecated test file).
 py_test(
     name = "test_impala_rollout",
     main = "tests/test_rollout.py",
@@ -854,6 +853,23 @@ py_test(
     srcs = ["models/tests/test_distributions.py"]
 )
 
+# --------------------------------------------------------------------
+# Optimizers and Memories
+# rllib/optimizers/
+#
+# Tag: optimizers
+# --------------------------------------------------------------------
+
+# This has bugs: See PR https://github.com/ray-project/ray/pull/7534
+# which fixes these and re-adds this test.
+
+# py_test(
+#     name = "test_segment_tree",
+#     tags = ["optimizers"],
+#     size = "small",
+#    srcs = ["optimizers/tests/test_segment_tree.py"]
+# )
+
 # --------------------------------------------------------------------
 # Policies
 # rllib/policy/
@@ -876,11 +892,10 @@ py_test(
 # --------------------------------------------------------------------
 
 py_test(
-    name = "test_framework_agnostic_components",
+    name = "test_explorations",
     tags = ["utils"],
-    size = "small",
-    data = glob(["utils/tests/**"]),
-    srcs = ["utils/tests/test_framework_agnostic_components.py"]
+    size = "large",
+    srcs = ["utils/exploration/tests/test_explorations.py"]
 )
 
 # Schedules
@@ -891,6 +906,14 @@ py_test(
     srcs = ["utils/schedules/tests/test_schedules.py"]
 )
 
+py_test(
+    name = "test_framework_agnostic_components",
+    tags = ["utils"],
+    size = "small",
+    data = glob(["utils/tests/**"]),
+    srcs = ["utils/tests/test_framework_agnostic_components.py"]
+)
+
 # TaskPool
 py_test(
     name = "test_taskpool",
@@ -959,13 +982,6 @@ py_test(
     srcs = ["tests/test_evaluators.py"]
 )
 
-py_test(
-    name = "tests/test_explorations",
-    tags = ["tests_dir", "tests_dir_E", "explorations"],
-    size = "large",
-    srcs = ["tests/test_explorations.py"]
-)
-
 py_test(
     name = "tests/test_external_env",
     tags = ["tests_dir", "tests_dir_E"],

rllib/agents/ddpg/tests/test_ddpg.py

@@ -83,5 +83,6 @@ class TestDDPG(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/agents/ddpg/tests/test_td3.py

@@ -83,5 +83,6 @@ class TestTD3(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/agents/impala/tests/test_vtrace.py

@@ -25,6 +25,7 @@ import numpy as np
 
 from ray.rllib.utils import try_import_tf
 import ray.rllib.agents.impala.vtrace as vtrace
+from ray.rllib.utils.numpy import softmax
 
 tf = try_import_tf()
 
@@ -34,11 +35,6 @@ def _shaped_arange(*shape):
     return np.arange(np.prod(shape), dtype=np.float32).reshape(*shape)
 
 
-def _softmax(logits):
-    """Applies softmax non-linearity on inputs."""
-    return np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
-
-
 def _ground_truth_calculation(discounts, log_rhos, rewards, values,
                               bootstrap_value, clip_rho_threshold,
                               clip_pg_rho_threshold):
@@ -108,7 +104,7 @@ class LogProbsFromLogitsAndActionsTest(tf.test.TestCase,
         # numerically stable. However, in this test we have well-behaved
         # values.
         ground_truth_v = index_with_mask(
-            np.log(_softmax(policy_logits)), action_index_mask)
+            np.log(softmax(policy_logits)), action_index_mask)
 
         with self.test_session() as session:
             self.assertAllClose(ground_truth_v,

rllib/agents/ppo/tests/test.py

@@ -38,4 +38,6 @@ class UtilsTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/agents/ppo/tests/test_ppo.py

@@ -22,8 +22,13 @@ tf = try_import_tf()
 
 
 class TestPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
 
-    ray.init()
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
 
     def test_ppo_compilation(self):
         """Test whether a PPOTrainer can be built with both frameworks."""
@@ -222,5 +227,6 @@ class TestPPO(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/models/tests/test_distributions.py

@@ -195,5 +195,6 @@ class TestDistributions(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/optimizers/tests/test_segment_tree.py

@@ -1,103 +1,99 @@
 import numpy as np
+import unittest
 
 from ray.rllib.optimizers.segment_tree import SumSegmentTree, MinSegmentTree
 
 
-def test_tree_set():
-    tree = SumSegmentTree(4)
+class TestSegmentTree(unittest.TestCase):
+    def test_tree_set(self):
+        tree = SumSegmentTree(4)
 
-    tree[2] = 1.0
-    tree[3] = 3.0
+        tree[2] = 1.0
+        tree[3] = 3.0
 
-    assert np.isclose(tree.sum(), 4.0)
-    assert np.isclose(tree.sum(0, 2), 0.0)
-    assert np.isclose(tree.sum(0, 3), 1.0)
-    assert np.isclose(tree.sum(2, 3), 1.0)
-    assert np.isclose(tree.sum(2, -1), 1.0)
-    assert np.isclose(tree.sum(2, 4), 4.0)
+        assert np.isclose(tree.sum(), 4.0)
+        assert np.isclose(tree.sum(0, 2), 0.0)
+        assert np.isclose(tree.sum(0, 3), 1.0)
+        assert np.isclose(tree.sum(2, 3), 1.0)
+        assert np.isclose(tree.sum(2, -1), 1.0)
+        assert np.isclose(tree.sum(2, 4), 4.0)
 
+    def test_tree_set_overlap(self):
+        tree = SumSegmentTree(4)
 
-def test_tree_set_overlap():
-    tree = SumSegmentTree(4)
+        tree[2] = 1.0
+        tree[2] = 3.0
 
-    tree[2] = 1.0
-    tree[2] = 3.0
+        assert np.isclose(tree.sum(), 3.0)
+        assert np.isclose(tree.sum(2, 3), 3.0)
+        assert np.isclose(tree.sum(2, -1), 3.0)
+        assert np.isclose(tree.sum(2, 4), 3.0)
+        assert np.isclose(tree.sum(1, 2), 0.0)
 
-    assert np.isclose(tree.sum(), 3.0)
-    assert np.isclose(tree.sum(2, 3), 3.0)
-    assert np.isclose(tree.sum(2, -1), 3.0)
-    assert np.isclose(tree.sum(2, 4), 3.0)
-    assert np.isclose(tree.sum(1, 2), 0.0)
+    def test_prefixsum_idx(self):
+        tree = SumSegmentTree(4)
 
+        tree[2] = 1.0
+        tree[3] = 3.0
 
-def test_prefixsum_idx():
-    tree = SumSegmentTree(4)
+        assert tree.find_prefixsum_idx(0.0) == 2
+        assert tree.find_prefixsum_idx(0.5) == 2
+        assert tree.find_prefixsum_idx(0.99) == 2
+        assert tree.find_prefixsum_idx(1.01) == 3
+        assert tree.find_prefixsum_idx(3.00) == 3
+        assert tree.find_prefixsum_idx(4.00) == 3
 
-    tree[2] = 1.0
-    tree[3] = 3.0
+    def test_prefixsum_idx2(self):
+        tree = SumSegmentTree(4)
 
-    assert tree.find_prefixsum_idx(0.0) == 2
-    assert tree.find_prefixsum_idx(0.5) == 2
-    assert tree.find_prefixsum_idx(0.99) == 2
-    assert tree.find_prefixsum_idx(1.01) == 3
-    assert tree.find_prefixsum_idx(3.00) == 3
-    assert tree.find_prefixsum_idx(4.00) == 3
+        tree[0] = 0.5
+        tree[1] = 1.0
+        tree[2] = 1.0
+        tree[3] = 3.0
 
+        assert tree.find_prefixsum_idx(0.00) == 0
+        assert tree.find_prefixsum_idx(0.55) == 1
+        assert tree.find_prefixsum_idx(0.99) == 1
+        assert tree.find_prefixsum_idx(1.51) == 2
+        assert tree.find_prefixsum_idx(3.00) == 3
+        assert tree.find_prefixsum_idx(5.50) == 3
 
-def test_prefixsum_idx2():
-    tree = SumSegmentTree(4)
+    def test_max_interval_tree(self):
+        tree = MinSegmentTree(4)
 
-    tree[0] = 0.5
-    tree[1] = 1.0
-    tree[2] = 1.0
-    tree[3] = 3.0
+        tree[0] = 1.0
+        tree[2] = 0.5
+        tree[3] = 3.0
 
-    assert tree.find_prefixsum_idx(0.00) == 0
-    assert tree.find_prefixsum_idx(0.55) == 1
-    assert tree.find_prefixsum_idx(0.99) == 1
-    assert tree.find_prefixsum_idx(1.51) == 2
-    assert tree.find_prefixsum_idx(3.00) == 3
-    assert tree.find_prefixsum_idx(5.50) == 3
+        assert np.isclose(tree.min(), 0.5)
+        assert np.isclose(tree.min(0, 2), 1.0)
+        assert np.isclose(tree.min(0, 3), 0.5)
+        assert np.isclose(tree.min(0, -1), 0.5)
+        assert np.isclose(tree.min(2, 4), 0.5)
+        assert np.isclose(tree.min(3, 4), 3.0)
 
+        tree[2] = 0.7
 
-def test_max_interval_tree():
-    tree = MinSegmentTree(4)
+        assert np.isclose(tree.min(), 0.7)
+        assert np.isclose(tree.min(0, 2), 1.0)
+        assert np.isclose(tree.min(0, 3), 0.7)
+        assert np.isclose(tree.min(0, -1), 0.7)
+        assert np.isclose(tree.min(2, 4), 0.7)
+        assert np.isclose(tree.min(3, 4), 3.0)
 
-    tree[0] = 1.0
-    tree[2] = 0.5
-    tree[3] = 3.0
+        tree[2] = 4.0
 
-    assert np.isclose(tree.min(), 0.5)
-    assert np.isclose(tree.min(0, 2), 1.0)
-    assert np.isclose(tree.min(0, 3), 0.5)
-    assert np.isclose(tree.min(0, -1), 0.5)
-    assert np.isclose(tree.min(2, 4), 0.5)
-    assert np.isclose(tree.min(3, 4), 3.0)
-
-    tree[2] = 0.7
-
-    assert np.isclose(tree.min(), 0.7)
-    assert np.isclose(tree.min(0, 2), 1.0)
-    assert np.isclose(tree.min(0, 3), 0.7)
-    assert np.isclose(tree.min(0, -1), 0.7)
-    assert np.isclose(tree.min(2, 4), 0.7)
-    assert np.isclose(tree.min(3, 4), 3.0)
-
-    tree[2] = 4.0
-
-    assert np.isclose(tree.min(), 1.0)
-    assert np.isclose(tree.min(0, 2), 1.0)
-    assert np.isclose(tree.min(0, 3), 1.0)
-    assert np.isclose(tree.min(0, -1), 1.0)
-    assert np.isclose(tree.min(2, 4), 3.0)
-    assert np.isclose(tree.min(2, 3), 4.0)
-    assert np.isclose(tree.min(2, -1), 4.0)
-    assert np.isclose(tree.min(3, 4), 3.0)
+        assert np.isclose(tree.min(), 1.0)
+        assert np.isclose(tree.min(0, 2), 1.0)
+        assert np.isclose(tree.min(0, 3), 1.0)
+        assert np.isclose(tree.min(0, -1), 1.0)
+        assert np.isclose(tree.min(2, 4), 3.0)
+        assert np.isclose(tree.min(2, 3), 4.0)
+        assert np.isclose(tree.min(2, -1), 4.0)
+        assert np.isclose(tree.min(3, 4), 3.0)
 
 
 if __name__ == "__main__":
-    test_tree_set()
-    test_tree_set_overlap()
-    test_prefixsum_idx()
-    test_prefixsum_idx2()
-    test_max_interval_tree()
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/policy/tests/test_compute_log_likelihoods.py

@@ -3,6 +3,7 @@ from scipy.stats import norm
 import unittest
 
 import ray.rllib.agents.dqn as dqn
+import ray.rllib.agents.pg as pg
 import ray.rllib.agents.ppo as ppo
 import ray.rllib.agents.sac as sac
 from ray.rllib.utils.framework import try_import_tf
@@ -13,12 +14,12 @@ from ray.rllib.utils.numpy import one_hot, fc, MIN_LOG_NN_OUTPUT, \
 tf = try_import_tf()
 
 
-def test_log_likelihood(run,
-                        config,
-                        prev_a=None,
-                        continuous=False,
-                        layer_key=("fc", (0, 4)),
-                        logp_func=None):
+def do_test_log_likelihood(run,
+                           config,
+                           prev_a=None,
+                           continuous=False,
+                           layer_key=("fc", (0, 4)),
+                           logp_func=None):
     config = config.copy()
     # Run locally.
     config["num_workers"] = 0
@@ -32,10 +33,6 @@ def test_log_likelihood(run,
         obs_batch = np.array([0])
         preprocessed_obs_batch = one_hot(obs_batch, depth=16)
 
-    # Use Soft-Q for DQNs.
-    if run is dqn.DQNTrainer:
-        config["exploration_config"] = {"type": "SoftQ", "temperature": 0.5}
-
     prev_r = None if prev_a is None else np.array(0.0)
 
     # Test against all frameworks.
@@ -43,15 +40,15 @@ def test_log_likelihood(run,
         if run in [dqn.DQNTrainer, sac.SACTrainer] and fw == "torch":
             continue
         print("Testing {} with framework={}".format(run, fw))
-        config["eager"] = True if fw == "eager" else False
-        config["use_pytorch"] = True if fw == "torch" else False
+        config["eager"] = fw == "eager"
+        config["use_pytorch"] = fw == "torch"
 
         trainer = run(config=config, env=env)
         policy = trainer.get_policy()
         vars = policy.get_weights()
         # Sample n actions, then roughly check their logp against their
         # counts.
-        num_actions = 500
+        num_actions = 1000 if not continuous else 50
         actions = []
         for _ in range(num_actions):
             # Single action from single obs.
@@ -62,9 +59,9 @@ def test_log_likelihood(run,
                     prev_reward=prev_r,
                     explore=True))
 
-        # Test 50 actions for their log-likelihoods vs expected values.
+        # Test all taken actions for their log-likelihoods vs expected values.
         if continuous:
-            for idx in range(50):
+            for idx in range(num_actions):
                 a = actions[idx]
                 if fw == "tf" or fw == "eager":
                     if isinstance(vars, list):
@@ -99,19 +96,41 @@ def test_log_likelihood(run,
         else:
             for a in [0, 1, 2, 3]:
                 count = actions.count(a)
-                expected_logp = np.log(count / num_actions)
+                expected_prob = count / num_actions
                 logp = policy.compute_log_likelihoods(
                     np.array([a]),
                     preprocessed_obs_batch,
                     prev_action_batch=np.array([prev_a]),
                     prev_reward_batch=np.array([prev_r]))
-                check(logp, expected_logp, rtol=0.3)
+                check(np.exp(logp), expected_prob, atol=0.2)
 
 
 class TestComputeLogLikelihood(unittest.TestCase):
     def test_dqn(self):
         """Tests, whether DQN correctly computes logp in soft-q mode."""
-        test_log_likelihood(dqn.DQNTrainer, dqn.DEFAULT_CONFIG)
+        config = dqn.DEFAULT_CONFIG.copy()
+        # Soft-Q for DQN.
+        config["exploration_config"] = {"type": "SoftQ", "temperature": 0.5}
+        do_test_log_likelihood(dqn.DQNTrainer, config)
+
+    def test_pg_cont(self):
+        """Tests PG's (cont. actions) compute_log_likelihoods method."""
+        config = pg.DEFAULT_CONFIG.copy()
+        config["model"]["fcnet_hiddens"] = [10]
+        config["model"]["fcnet_activation"] = "linear"
+        prev_a = np.array([0.0])
+        do_test_log_likelihood(
+            pg.PGTrainer,
+            config,
+            prev_a,
+            continuous=True,
+            layer_key=("fc", (0, 2)))
+
+    def test_pg_discr(self):
+        """Tests PG's (cont. actions) compute_log_likelihoods method."""
+        config = pg.DEFAULT_CONFIG.copy()
+        prev_a = np.array(0)
+        do_test_log_likelihood(pg.PGTrainer, config, prev_a)
 
     def test_ppo_cont(self):
         """Tests PPO's (cont. actions) compute_log_likelihoods method."""
@@ -119,20 +138,22 @@ class TestComputeLogLikelihood(unittest.TestCase):
         config["model"]["fcnet_hiddens"] = [10]
         config["model"]["fcnet_activation"] = "linear"
         prev_a = np.array([0.0])
-        test_log_likelihood(ppo.PPOTrainer, config, prev_a, continuous=True)
+        do_test_log_likelihood(ppo.PPOTrainer, config, prev_a, continuous=True)
 
     def test_ppo_discr(self):
         """Tests PPO's (discr. actions) compute_log_likelihoods method."""
         prev_a = np.array(0)
-        test_log_likelihood(ppo.PPOTrainer, ppo.DEFAULT_CONFIG, prev_a)
+        do_test_log_likelihood(ppo.PPOTrainer, ppo.DEFAULT_CONFIG, prev_a)
 
-    def test_sac(self):
-        """Tests SAC's compute_log_likelihoods method."""
+    def test_sac_cont(self):
+        """Tests SAC's (cont. actions) compute_log_likelihoods method."""
         config = sac.DEFAULT_CONFIG.copy()
         config["policy_model"]["hidden_layer_sizes"] = [10]
         config["policy_model"]["hidden_activation"] = "linear"
         prev_a = np.array([0.0])
 
+        # SAC cont uses a squashed normal distribution. Implement it's logp
+        # logic here in numpy for comparing results.
         def logp_func(means, log_stds, values, low=-1.0, high=1.0):
             stds = np.exp(
                 np.clip(log_stds, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT))
@@ -144,10 +165,29 @@ class TestComputeLogLikelihood(unittest.TestCase):
                 np.sum(np.log(1 - np.tanh(unsquashed_values) ** 2),
                        axis=-1)
 
-        test_log_likelihood(
+        do_test_log_likelihood(
             sac.SACTrainer,
             config,
             prev_a,
             continuous=True,
             layer_key=("sequential/action", (0, 2)),
             logp_func=logp_func)
+
+    def test_sac_discr(self):
+        """Tests SAC's (discrete actions) compute_log_likelihoods method."""
+        config = sac.DEFAULT_CONFIG.copy()
+        config["policy_model"]["hidden_layer_sizes"] = [10]
+        config["policy_model"]["hidden_activation"] = "linear"
+        prev_a = np.array(0)
+
+        do_test_log_likelihood(
+            sac.SACTrainer,
+            config,
+            prev_a,
+            layer_key=("sequential/action", (0, 2)))
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/run_regression_tests.py

@@ -6,12 +6,14 @@
 #
 # When using in BAZEL (with py_test), e.g. see in ray/rllib/BUILD:
 # py_test(
-#    name = "run_regression_tests",
-#    main = "tests/run_regression_tests.py",
-#    size = "large",
-#    srcs = ["tests/run_regression_tests.py"],
-#    data = glob(["tuned_examples/regression_tests/**"]),
-#    args = glob(["tuned_examples/regression_tests/**"])
+#     name = "run_regression_tests",
+#     main = "tests/run_regression_tests.py",
+#     tags = ["learning_tests"],
+#     size = "enormous",  # = 60min timeout
+#     srcs = ["tests/run_regression_tests.py"],
+#     data = glob(["tuned_examples/regression_tests/*.yaml"]),
+#     Pass `BAZEL` option and the path to look for yaml regression files.
+#     args = ["BAZEL", "tuned_examples/regression_tests"]
 # )
 
 from pathlib import Path

rllib/tests/test_avail_actions_qmix.py

@@ -1,5 +1,6 @@
-import numpy as np
 from gym.spaces import Tuple, Discrete, Dict, Box
+import numpy as np
+import unittest
 
 import ray
 from ray.tune import register_env
@@ -40,26 +41,33 @@ class AvailActionsTestEnv(MultiAgentEnv):
         return obs, rewards, dones, {}
 
 
-if __name__ == "__main__":
-    grouping = {
-        "group_1": ["agent_1"],  # trivial grouping for testing
-    }
-    obs_space = Tuple([AvailActionsTestEnv.observation_space])
-    act_space = Tuple([AvailActionsTestEnv.action_space])
-    register_env(
-        "action_mask_test",
-        lambda config: AvailActionsTestEnv(config).with_agent_groups(
-            grouping, obs_space=obs_space, act_space=act_space))
+class TestAvailActionsQMix(unittest.TestCase):
+    def test_avail_actions_qmix(self):
+        grouping = {
+            "group_1": ["agent_1"],  # trivial grouping for testing
+        }
+        obs_space = Tuple([AvailActionsTestEnv.observation_space])
+        act_space = Tuple([AvailActionsTestEnv.action_space])
+        register_env(
+            "action_mask_test",
+            lambda config: AvailActionsTestEnv(config).with_agent_groups(
+                grouping, obs_space=obs_space, act_space=act_space))
 
-    ray.init()
-    agent = QMixTrainer(
-        env="action_mask_test",
-        config={
-            "num_envs_per_worker": 5,  # test with vectorization on
-            "env_config": {
-                "avail_action": 3,
-            },
-        })
-    for _ in range(5):
-        agent.train()  # OK if it doesn't trip the action assertion error
-    assert agent.train()["episode_reward_mean"] == 21.0
+        ray.init()
+        agent = QMixTrainer(
+            env="action_mask_test",
+            config={
+                "num_envs_per_worker": 5,  # test with vectorization on
+                "env_config": {
+                    "avail_action": 3,
+                },
+            })
+        for _ in range(5):
+            agent.train()  # OK if it doesn't trip the action assertion error
+        assert agent.train()["episode_reward_mean"] == 21.0
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_catalog.py

@@ -1,10 +1,9 @@
 import gym
+from gym.spaces import Box, Discrete, Tuple
 import numpy as np
 import unittest
-from gym.spaces import Box, Discrete, Tuple
 
 import ray
-
 from ray.rllib.models import ModelCatalog, MODEL_DEFAULTS
 from ray.rllib.models.model import Model
 from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
@@ -55,14 +54,14 @@ class ModelCatalogTest(unittest.TestCase):
     def tearDown(self):
         ray.shutdown()
 
-    def testGymPreprocessors(self):
+    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 testTuplePreprocessor(self):
+    def test_tuple_preprocessor(self):
         ray.init(object_store_memory=1000 * 1024 * 1024)
 
         class TupleEnv:
@@ -77,7 +76,7 @@ class ModelCatalogTest(unittest.TestCase):
             list(p1.transform((0, np.array([1, 2, 3])))),
             [float(x) for x in [1, 0, 0, 0, 0, 1, 2, 3]])
 
-    def testCustomPreprocessor(self):
+    def test_custom_preprocessor(self):
         ray.init(object_store_memory=1000 * 1024 * 1024)
         ModelCatalog.register_custom_preprocessor("foo", CustomPreprocessor)
         ModelCatalog.register_custom_preprocessor("bar", CustomPreprocessor2)
@@ -89,7 +88,7 @@ class ModelCatalogTest(unittest.TestCase):
         p3 = ModelCatalog.get_preprocessor(env)
         self.assertEqual(type(p3), NoPreprocessor)
 
-    def testDefaultModels(self):
+    def test_default_models(self):
         ray.init(object_store_memory=1000 * 1024 * 1024)
 
         with tf.variable_scope("test1"):
@@ -105,7 +104,7 @@ class ModelCatalogTest(unittest.TestCase):
                                         {})
             self.assertEqual(type(p2), VisionNetwork)
 
-    def testCustomModel(self):
+    def test_custom_model(self):
         ray.init(object_store_memory=1000 * 1024 * 1024)
         ModelCatalog.register_custom_model("foo", CustomModel)
         p1 = ModelCatalog.get_model({
@@ -114,7 +113,7 @@ class ModelCatalogTest(unittest.TestCase):
                                     {"custom_model": "foo"})
         self.assertEqual(str(type(p1)), str(CustomModel))
 
-    def testCustomActionDistribution(self):
+    def test_custom_action_distribution(self):
         class Model():
             pass
 
@@ -159,4 +158,6 @@ class ModelCatalogTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_checkpoint_restore.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python
 
-import os
-import shutil
 import gym
 import numpy as np
-import ray
+import os
+import shutil
+import unittest
 
+import ray
 from ray.rllib.agents.registry import get_agent_class
 from ray.tune.trial import ExportFormat
 
@@ -17,8 +18,6 @@ def get_mean_action(alg, obs):
     return np.mean(out)
 
 
-ray.init(num_cpus=10, object_store_memory=1e9)
-
 CONFIGS = {
     "SAC": {
         "explore": False,
@@ -67,15 +66,15 @@ CONFIGS = {
 }
 
 
-def test_ckpt_restore(use_object_store, alg_name, failures):
+def ckpt_restore_test(use_object_store, alg_name, failures):
     cls = get_agent_class(alg_name)
     if "DDPG" in alg_name or "SAC" in alg_name:
-        alg1 = cls(config=CONFIGS[name], env="Pendulum-v0")
-        alg2 = cls(config=CONFIGS[name], env="Pendulum-v0")
+        alg1 = cls(config=CONFIGS[alg_name], env="Pendulum-v0")
+        alg2 = cls(config=CONFIGS[alg_name], env="Pendulum-v0")
         env = gym.make("Pendulum-v0")
     else:
-        alg1 = cls(config=CONFIGS[name], env="CartPole-v0")
-        alg2 = cls(config=CONFIGS[name], env="CartPole-v0")
+        alg1 = cls(config=CONFIGS[alg_name], env="CartPole-v0")
+        alg2 = cls(config=CONFIGS[alg_name], env="CartPole-v0")
         env = gym.make("CartPole-v0")
 
     for _ in range(2):
@@ -106,7 +105,7 @@ def test_ckpt_restore(use_object_store, alg_name, failures):
             failures.append((alg_name, [a1, a2]))
 
 
-def test_export(algo_name, failures):
+def export_test(alg_name, failures):
     def valid_tf_model(model_dir):
         return os.path.exists(os.path.join(model_dir, "saved_model.pb")) \
             and os.listdir(os.path.join(model_dir, "variables"))
@@ -116,52 +115,68 @@ def test_export(algo_name, failures):
             and os.path.exists(os.path.join(checkpoint_dir, "model.index")) \
             and os.path.exists(os.path.join(checkpoint_dir, "checkpoint"))
 
-    cls = get_agent_class(algo_name)
-    if "DDPG" in algo_name or "SAC" in algo_name:
-        algo = cls(config=CONFIGS[name], env="Pendulum-v0")
+    cls = get_agent_class(alg_name)
+    if "DDPG" in alg_name or "SAC" in alg_name:
+        algo = cls(config=CONFIGS[alg_name], env="Pendulum-v0")
     else:
-        algo = cls(config=CONFIGS[name], env="CartPole-v0")
+        algo = cls(config=CONFIGS[alg_name], env="CartPole-v0")
 
     for _ in range(3):
         res = algo.train()
         print("current status: " + str(res))
 
-    export_dir = "/tmp/export_dir_%s" % algo_name
-    print("Exporting model ", algo_name, export_dir)
+    export_dir = "/tmp/export_dir_%s" % alg_name
+    print("Exporting model ", alg_name, export_dir)
     algo.export_policy_model(export_dir)
     if not valid_tf_model(export_dir):
-        failures.append(algo_name)
+        failures.append(alg_name)
     shutil.rmtree(export_dir)
 
-    print("Exporting checkpoint", algo_name, export_dir)
+    print("Exporting checkpoint", alg_name, export_dir)
     algo.export_policy_checkpoint(export_dir)
     if not valid_tf_checkpoint(export_dir):
-        failures.append(algo_name)
+        failures.append(alg_name)
     shutil.rmtree(export_dir)
 
-    print("Exporting default policy", algo_name, export_dir)
+    print("Exporting default policy", alg_name, export_dir)
     algo.export_model([ExportFormat.CHECKPOINT, ExportFormat.MODEL],
                       export_dir)
     if not valid_tf_model(os.path.join(export_dir, ExportFormat.MODEL)) \
             or not valid_tf_checkpoint(os.path.join(export_dir,
                                                     ExportFormat.CHECKPOINT)):
-        failures.append(algo_name)
+        failures.append(alg_name)
     shutil.rmtree(export_dir)
 
 
+class TestCheckpointRestore(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=10, object_store_memory=1e9)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_checkpoint_restore(self):
+        failures = []
+        for use_object_store in [False, True]:
+            for name in [
+                    "SAC", "ES", "DQN", "DDPG", "PPO", "A3C", "APEX_DDPG",
+                    "ARS"
+            ]:
+                ckpt_restore_test(use_object_store, name, failures)
+
+        assert not failures, failures
+        print("All checkpoint restore tests passed!")
+
+        failures = []
+        for name in ["SAC", "DQN", "DDPG", "PPO", "A3C"]:
+            export_test(name, failures)
+        assert not failures, failures
+        print("All export tests passed!")
+
+
 if __name__ == "__main__":
-    failures = []
-    for use_object_store in [False, True]:
-        for name in [
-                "SAC", "ES", "DQN", "DDPG", "PPO", "A3C", "APEX_DDPG", "ARS"
-        ]:
-            test_ckpt_restore(use_object_store, name, failures)
-
-    assert not failures, failures
-    print("All checkpoint restore tests passed!")
-
-    failures = []
-    for name in ["SAC", "DQN", "DDPG", "PPO", "A3C"]:
-        test_export(name, failures)
-    assert not failures, failures
-    print("All export tests passed!")
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_eager_support.py

@@ -7,20 +7,31 @@ from ray.rllib.agents.registry import get_agent_class
 
 def check_support(alg, config, test_trace=True):
     config["eager"] = True
-    if alg in ["APEX_DDPG", "TD3", "DDPG", "SAC"]:
-        config["env"] = "Pendulum-v0"
-    else:
-        config["env"] = "CartPole-v0"
-    a = get_agent_class(alg)
-    config["log_level"] = "ERROR"
 
-    config["eager_tracing"] = False
-    tune.run(a, config=config, stop={"training_iteration": 1})
+    # Test both continuous and discrete actions.
+    for cont in [True, False]:
+        if cont and alg in ["DQN", "APEX", "SimpleQ"]:
+            continue
+        elif not cont and alg in ["DDPG", "APEX_DDPG", "TD3"]:
+            continue
 
-    if test_trace:
-        config["eager_tracing"] = True
+        print("run={} cont. actions={}".format(alg, cont))
+
+        if cont:
+            config["env"] = "Pendulum-v0"
+        else:
+            config["env"] = "CartPole-v0"
+
+        a = get_agent_class(alg)
+        config["log_level"] = "ERROR"
+
+        config["eager_tracing"] = False
         tune.run(a, config=config, stop={"training_iteration": 1})
 
+        if test_trace:
+            config["eager_tracing"] = True
+            tune.run(a, config=config, stop={"training_iteration": 1})
+
 
 class TestEagerSupport(unittest.TestCase):
     def setUp(self):
@@ -29,31 +40,41 @@ class TestEagerSupport(unittest.TestCase):
     def tearDown(self):
         ray.shutdown()
 
-    def testSimpleQ(self):
+    def test_simple_q(self):
         check_support("SimpleQ", {"num_workers": 0, "learning_starts": 0})
 
-    def testDQN(self):
+    def test_dqn(self):
         check_support("DQN", {"num_workers": 0, "learning_starts": 0})
 
-    def testA2C(self):
+    # TODO(sven): Add these once DDPG supports eager.
+    # def test_ddpg(self):
+    #     check_support("DDPG", {"num_workers": 0})
+
+    # def test_apex_ddpg(self):
+    #     check_support("APEX_DDPG", {"num_workers": 1})
+
+    # def test_td3(self):
+    #     check_support("TD3", {"num_workers": 0})
+
+    def test_a2c(self):
         check_support("A2C", {"num_workers": 0})
 
-    def testA3C(self):
+    def test_a3c(self):
         check_support("A3C", {"num_workers": 1})
 
-    def testPG(self):
+    def test_pg(self):
         check_support("PG", {"num_workers": 0})
 
-    def testPPO(self):
+    def test_ppo(self):
         check_support("PPO", {"num_workers": 0})
 
-    def testAPPO(self):
+    def test_appo(self):
         check_support("APPO", {"num_workers": 1, "num_gpus": 0})
 
-    def testIMPALA(self):
+    def test_impala(self):
         check_support("IMPALA", {"num_workers": 1, "num_gpus": 0})
 
-    def testAPEX_DQN(self):
+    def test_apex_dqn(self):
         check_support(
             "APEX", {
                 "num_workers": 2,
@@ -66,7 +87,7 @@ class TestEagerSupport(unittest.TestCase):
                 },
             })
 
-    def testSAC(self):
+    def test_sac(self):
         check_support("SAC", {"num_workers": 0})
 
 

rllib/tests/test_evaluators.py

@@ -1,3 +1,4 @@
+import gym
 import unittest
 
 import ray
@@ -5,11 +6,10 @@ from ray.rllib.agents.dqn import DQNTrainer
 from ray.rllib.agents.a3c import A3CTrainer
 from ray.rllib.agents.dqn.dqn_policy import _adjust_nstep
 from ray.tune.registry import register_env
-import gym
 
 
 class EvalTest(unittest.TestCase):
-    def testDqnNStep(self):
+    def test_dqn_n_step(self):
         obs = [1, 2, 3, 4, 5, 6, 7]
         actions = ["a", "b", "a", "a", "a", "b", "a"]
         rewards = [10.0, 0.0, 100.0, 100.0, 100.0, 100.0, 100.0]
@@ -23,7 +23,7 @@ class EvalTest(unittest.TestCase):
         self.assertEqual(rewards,
                          [91.0, 171.0, 271.0, 271.0, 271.0, 190.0, 100.0])
 
-    def testEvaluationOption(self):
+    def test_evaluation_option(self):
         def env_creator(env_config):
             return gym.make("CartPole-v0")
 
@@ -61,4 +61,6 @@ class EvalTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_external_env.py

@@ -114,7 +114,13 @@ class MultiServing(ExternalEnv):
 
 
 class TestExternalEnv(unittest.TestCase):
-    def testExternalEnvCompleteEpisodes(self):
+    def setUp(self) -> None:
+        ray.init()
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_external_env_complete_episodes(self):
         ev = RolloutWorker(
             env_creator=lambda _: SimpleServing(MockEnv(25)),
             policy=MockPolicy,
@@ -124,7 +130,7 @@ class TestExternalEnv(unittest.TestCase):
             batch = ev.sample()
             self.assertEqual(batch.count, 50)
 
-    def testExternalEnvTruncateEpisodes(self):
+    def test_external_env_truncate_episodes(self):
         ev = RolloutWorker(
             env_creator=lambda _: SimpleServing(MockEnv(25)),
             policy=MockPolicy,
@@ -134,7 +140,7 @@ class TestExternalEnv(unittest.TestCase):
             batch = ev.sample()
             self.assertEqual(batch.count, 40)
 
-    def testExternalEnvOffPolicy(self):
+    def test_external_env_off_policy(self):
         ev = RolloutWorker(
             env_creator=lambda _: SimpleOffPolicyServing(MockEnv(25), 42),
             policy=MockPolicy,
@@ -146,7 +152,7 @@ class TestExternalEnv(unittest.TestCase):
             self.assertEqual(batch["actions"][0], 42)
             self.assertEqual(batch["actions"][-1], 42)
 
-    def testExternalEnvBadActions(self):
+    def test_external_env_bad_actions(self):
         ev = RolloutWorker(
             env_creator=lambda _: SimpleServing(MockEnv(25)),
             policy=BadPolicy,
@@ -155,7 +161,7 @@ class TestExternalEnv(unittest.TestCase):
             batch_mode="truncate_episodes")
         self.assertRaises(Exception, lambda: ev.sample())
 
-    def testTrainCartpoleOffPolicy(self):
+    def test_train_cartpole_off_policy(self):
         register_env(
             "test3", lambda _: PartOffPolicyServing(
                 gym.make("CartPole-v0"), off_pol_frac=0.2))
@@ -172,7 +178,7 @@ class TestExternalEnv(unittest.TestCase):
                 return
         raise Exception("failed to improve reward")
 
-    def testTrainCartpole(self):
+    def test_train_cartpole(self):
         register_env("test", lambda _: SimpleServing(gym.make("CartPole-v0")))
         pg = PGTrainer(env="test", config={"num_workers": 0})
         for i in range(100):
@@ -183,7 +189,7 @@ class TestExternalEnv(unittest.TestCase):
                 return
         raise Exception("failed to improve reward")
 
-    def testTrainCartpoleMulti(self):
+    def test_train_cartpole_multi(self):
         register_env("test2",
                      lambda _: MultiServing(lambda: gym.make("CartPole-v0")))
         pg = PGTrainer(env="test2", config={"num_workers": 0})
@@ -195,7 +201,7 @@ class TestExternalEnv(unittest.TestCase):
                 return
         raise Exception("failed to improve reward")
 
-    def testExternalEnvHorizonNotSupported(self):
+    def test_external_env_horizon_not_supported(self):
         ev = RolloutWorker(
             env_creator=lambda _: SimpleServing(MockEnv(25)),
             policy=MockPolicy,
@@ -206,5 +212,6 @@ class TestExternalEnv(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init()
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_external_multi_agent_env.py

@@ -18,7 +18,13 @@ SimpleMultiServing = make_simple_serving(True, ExternalMultiAgentEnv)
 
 
 class TestExternalMultiAgentEnv(unittest.TestCase):
-    def testExternalMultiAgentEnvCompleteEpisodes(self):
+    def setUp(self) -> None:
+        ray.init()
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_external_multi_agent_env_complete_episodes(self):
         agents = 4
         ev = RolloutWorker(
             env_creator=lambda _: SimpleMultiServing(BasicMultiAgent(agents)),
@@ -30,7 +36,7 @@ class TestExternalMultiAgentEnv(unittest.TestCase):
             self.assertEqual(batch.count, 40)
             self.assertEqual(len(np.unique(batch["agent_index"])), agents)
 
-    def testExternalMultiAgentEnvTruncateEpisodes(self):
+    def test_external_multi_agent_env_truncate_episodes(self):
         agents = 4
         ev = RolloutWorker(
             env_creator=lambda _: SimpleMultiServing(BasicMultiAgent(agents)),
@@ -42,7 +48,7 @@ class TestExternalMultiAgentEnv(unittest.TestCase):
             self.assertEqual(batch.count, 160)
             self.assertEqual(len(np.unique(batch["agent_index"])), agents)
 
-    def testExternalMultiAgentEnvSample(self):
+    def test_external_multi_agent_env_sample(self):
         agents = 2
         act_space = gym.spaces.Discrete(2)
         obs_space = gym.spaces.Discrete(2)
@@ -57,7 +63,7 @@ class TestExternalMultiAgentEnv(unittest.TestCase):
         batch = ev.sample()
         self.assertEqual(batch.count, 50)
 
-    def testTrainExternalMultiCartpoleManyPolicies(self):
+    def test_train_external_multi_cartpole_many_policies(self):
         n = 20
         single_env = gym.make("CartPole-v0")
         act_space = single_env.action_space
@@ -85,5 +91,6 @@ class TestExternalMultiAgentEnv(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init()
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_filters.py

@@ -1,5 +1,5 @@
-import unittest
 import numpy as np
+import unittest
 
 import ray
 from ray.rllib.utils.filter import RunningStat, MeanStdFilter
@@ -76,7 +76,7 @@ class FilterManagerTest(unittest.TestCase):
     def tearDown(self):
         ray.shutdown()
 
-    def testSynchronize(self):
+    def test_synchronize(self):
         """Synchronize applies filter buffer onto own filter"""
         filt1 = MeanStdFilter(())
         for i in range(10):
@@ -103,4 +103,6 @@ class FilterManagerTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_ignore_worker_failure.py

@@ -25,8 +25,8 @@ class FaultInjectEnv(gym.Env):
 
 
 class IgnoresWorkerFailure(unittest.TestCase):
-    def doTest(self, alg, config, fn=None):
-        fn = fn or self._doTestFaultRecover
+    def do_test(self, alg, config, fn=None):
+        fn = fn or self._do_test_fault_recover
         try:
             ray.init(num_cpus=6)
             fn(alg, config)
@@ -34,7 +34,7 @@ class IgnoresWorkerFailure(unittest.TestCase):
             ray.shutdown()
             _register_all()  # re-register the evicted objects
 
-    def _doTestFaultRecover(self, alg, config):
+    def _do_test_fault_recover(self, alg, config):
         register_env("fault_env", lambda c: FaultInjectEnv(c))
         agent_cls = get_agent_class(alg)
 
@@ -47,7 +47,7 @@ class IgnoresWorkerFailure(unittest.TestCase):
         self.assertTrue(result["num_healthy_workers"], 1)
         a.stop()
 
-    def _doTestFaultFatal(self, alg, config):
+    def _do_test_fault_fatal(self, alg, config):
         register_env("fault_env", lambda c: FaultInjectEnv(c))
         agent_cls = get_agent_class(alg)
 
@@ -59,15 +59,15 @@ class IgnoresWorkerFailure(unittest.TestCase):
         self.assertRaises(Exception, lambda: a.train())
         a.stop()
 
-    def testFatal(self):
+    def test_fatal(self):
         # test the case where all workers fail
-        self.doTest("PG", {"optimizer": {}}, fn=self._doTestFaultFatal)
+        self.do_test("PG", {"optimizer": {}}, fn=self._do_test_fault_fatal)
 
-    def testAsyncGrads(self):
-        self.doTest("A3C", {"optimizer": {"grads_per_step": 1}})
+    def test_async_grads(self):
+        self.do_test("A3C", {"optimizer": {"grads_per_step": 1}})
 
-    def testAsyncReplay(self):
-        self.doTest(
+    def test_async_replay(self):
+        self.do_test(
             "APEX", {
                 "timesteps_per_iteration": 1000,
                 "num_gpus": 0,
@@ -80,14 +80,14 @@ class IgnoresWorkerFailure(unittest.TestCase):
                 },
             })
 
-    def testAsyncSamples(self):
-        self.doTest("IMPALA", {"num_gpus": 0})
+    def test_async_samples(self):
+        self.do_test("IMPALA", {"num_gpus": 0})
 
-    def testSyncReplay(self):
-        self.doTest("DQN", {"timesteps_per_iteration": 1})
+    def test_sync_replay(self):
+        self.do_test("DQN", {"timesteps_per_iteration": 1})
 
-    def testMultiGPU(self):
-        self.doTest(
+    def test_multi_g_p_u(self):
+        self.do_test(
             "PPO", {
                 "num_sgd_iter": 1,
                 "train_batch_size": 10,
@@ -95,12 +95,14 @@ class IgnoresWorkerFailure(unittest.TestCase):
                 "sgd_minibatch_size": 1,
             })
 
-    def testSyncSamples(self):
-        self.doTest("PG", {"optimizer": {}})
+    def test_sync_samples(self):
+        self.do_test("PG", {"optimizer": {}})
 
-    def testAsyncSamplingOption(self):
-        self.doTest("PG", {"optimizer": {}, "sample_async": True})
+    def test_async_sampling_option(self):
+        self.do_test("PG", {"optimizer": {}, "sample_async": True})
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_io.py

@@ -203,12 +203,14 @@ class AgentIOTest(unittest.TestCase):
 
 class JsonIOTest(unittest.TestCase):
     def setUp(self):
+        ray.init(num_cpus=1)
         self.test_dir = tempfile.mkdtemp()
 
     def tearDown(self):
         shutil.rmtree(self.test_dir)
+        ray.shutdown()
 
-    def testWriteSimple(self):
+    def test_write_simple(self):
         ioctx = IOContext(self.test_dir, {}, 0, None)
         writer = JsonWriter(
             self.test_dir, ioctx, max_file_size=1000, compress_columns=["obs"])
@@ -217,7 +219,7 @@ class JsonIOTest(unittest.TestCase):
         writer.write(SAMPLES)
         self.assertEqual(len(os.listdir(self.test_dir)), 1)
 
-    def testWriteFileURI(self):
+    def test_write_file_uri(self):
         ioctx = IOContext(self.test_dir, {}, 0, None)
         writer = JsonWriter(
             "file:" + self.test_dir,
@@ -229,7 +231,7 @@ class JsonIOTest(unittest.TestCase):
         writer.write(SAMPLES)
         self.assertEqual(len(os.listdir(self.test_dir)), 1)
 
-    def testWritePaginate(self):
+    def test_write_paginate(self):
         ioctx = IOContext(self.test_dir, {}, 0, None)
         writer = JsonWriter(
             self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"])
@@ -246,7 +248,7 @@ class JsonIOTest(unittest.TestCase):
             "Expected 2|7|12|13 files, but found {} ({})". \
             format(num_files, os.listdir(self.test_dir))
 
-    def testReadWrite(self):
+    def test_read_write(self):
         ioctx = IOContext(self.test_dir, {}, 0, None)
         writer = JsonWriter(
             self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"])
@@ -264,7 +266,7 @@ class JsonIOTest(unittest.TestCase):
         self.assertGreater(len(seen_o), 90)
         self.assertLess(len(seen_o), 101)
 
-    def testSkipsOverEmptyLinesAndFiles(self):
+    def test_skips_over_empty_lines_and_files(self):
         open(self.test_dir + "/empty", "w").close()
         with open(self.test_dir + "/f1", "w") as f:
             f.write("\n")
@@ -284,7 +286,7 @@ class JsonIOTest(unittest.TestCase):
             seen_a.add(batch["actions"][0])
         self.assertEqual(len(seen_a), 2)
 
-    def testSkipsOverCorruptedLines(self):
+    def test_skips_over_corrupted_lines(self):
         with open(self.test_dir + "/f1", "w") as f:
             f.write(_to_json(make_sample_batch(0), []))
             f.write("\n")
@@ -304,7 +306,7 @@ class JsonIOTest(unittest.TestCase):
             seen_a.add(batch["actions"][0])
         self.assertEqual(len(seen_a), 4)
 
-    def testAbortOnAllEmptyInputs(self):
+    def test_abort_on_all_empty_inputs(self):
         open(self.test_dir + "/empty", "w").close()
         reader = JsonReader([
             self.test_dir + "/empty",
@@ -324,5 +326,6 @@ class JsonIOTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=1)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_local.py

@@ -5,12 +5,19 @@ import ray
 
 
 class LocalModeTest(unittest.TestCase):
-    def testLocal(self):
+    def setUp(self) -> None:
         ray.init(local_mode=True)
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_local(self):
         cf = DEFAULT_CONFIG.copy()
         agent = PPOTrainer(cf, "CartPole-v0")
         print(agent.train())
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_lstm.py

@@ -16,8 +16,8 @@ from ray.rllib.utils import try_import_tf
 tf = try_import_tf()
 
 
-class LSTMUtilsTest(unittest.TestCase):
-    def testBasic(self):
+class TestLSTMUtils(unittest.TestCase):
+    def test_basic(self):
         eps_ids = [1, 1, 1, 5, 5, 5, 5, 5]
         agent_ids = [1, 1, 1, 1, 1, 1, 1, 1]
         f = [[101, 102, 103, 201, 202, 203, 204, 205],
@@ -34,7 +34,7 @@ class LSTMUtilsTest(unittest.TestCase):
         self.assertEqual([s.tolist() for s in s_init], [[209, 109, 105]])
         self.assertEqual(seq_lens.tolist(), [3, 4, 1])
 
-    def testMultiDim(self):
+    def test_multi_dim(self):
         eps_ids = [1, 1, 1]
         agent_ids = [1, 1, 1]
         obs = np.ones((84, 84, 4))
@@ -49,7 +49,7 @@ class LSTMUtilsTest(unittest.TestCase):
         self.assertEqual([s.tolist() for s in s_init], [[209]])
         self.assertEqual(seq_lens.tolist(), [3])
 
-    def testBatchId(self):
+    def test_batch_id(self):
         eps_ids = [1, 1, 1, 5, 5, 5, 5, 5]
         batch_ids = [1, 1, 2, 2, 3, 3, 4, 4]
         agent_ids = [1, 1, 1, 1, 1, 1, 1, 1]
@@ -60,7 +60,7 @@ class LSTMUtilsTest(unittest.TestCase):
                                              s, 4)
         self.assertEqual(seq_lens.tolist(), [2, 1, 1, 2, 2])
 
-    def testMultiAgent(self):
+    def test_multi_agent(self):
         eps_ids = [1, 1, 1, 5, 5, 5, 5, 5]
         agent_ids = [1, 1, 2, 1, 1, 2, 2, 3]
         f = [[101, 102, 103, 201, 202, 203, 204, 205],
@@ -78,7 +78,7 @@ class LSTMUtilsTest(unittest.TestCase):
         self.assertEqual(len(f_pad[0]), 20)
         self.assertEqual(len(s_init[0]), 5)
 
-    def testDynamicMaxLen(self):
+    def test_dynamic_max_len(self):
         eps_ids = [5, 2, 2]
         agent_ids = [2, 2, 2]
         f = [[1, 1, 1]]
@@ -184,8 +184,14 @@ class DebugCounterEnv(gym.Env):
         return [self.i], self.i % 3, self.i >= 15, {}
 
 
-class RNNSequencing(unittest.TestCase):
-    def testSimpleOptimizerSequencing(self):
+class TestRNNSequencing(unittest.TestCase):
+    def setUp(self) -> None:
+        ray.init(num_cpus=4)
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_simple_optimizer_sequencing(self):
         ModelCatalog.register_custom_model("rnn", RNNSpyModel)
         register_env("counter", lambda _: DebugCounterEnv())
         ppo = PPOTrainer(
@@ -242,7 +248,7 @@ class RNNSequencing(unittest.TestCase):
         self.assertGreater(abs(np.sum(batch1["state_in"][0][3])), 0)
         self.assertGreater(abs(np.sum(batch1["state_in"][1][3])), 0)
 
-    def testMinibatchSequencing(self):
+    def test_minibatch_sequencing(self):
         ModelCatalog.register_custom_model("rnn", RNNSpyModel)
         register_env("counter", lambda _: DebugCounterEnv())
         ppo = PPOTrainer(
@@ -305,5 +311,6 @@ class RNNSequencing(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=4)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_multi_agent_env.py

@@ -180,7 +180,13 @@ MultiMountainCar = make_multiagent("MountainCarContinuous-v0")
 
 
 class TestMultiAgentEnv(unittest.TestCase):
-    def testBasicMock(self):
+    def setUp(self) -> None:
+        ray.init(num_cpus=4)
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_basic_mock(self):
         env = BasicMultiAgent(4)
         obs = env.reset()
         self.assertEqual(obs, {0: 0, 1: 0, 2: 0, 3: 0})
@@ -204,7 +210,7 @@ class TestMultiAgentEnv(unittest.TestCase):
             "__all__": True
         })
 
-    def testRoundRobinMock(self):
+    def test_round_robin_mock(self):
         env = RoundRobinMultiAgent(2)
         obs = env.reset()
         self.assertEqual(obs, {0: 0})
@@ -218,13 +224,13 @@ class TestMultiAgentEnv(unittest.TestCase):
         obs, rew, done, info = env.step({0: 0})
         self.assertEqual(done["__all__"], True)
 
-    def testNoResetUntilPoll(self):
+    def test_no_reset_until_poll(self):
         env = _MultiAgentEnvToBaseEnv(lambda v: BasicMultiAgent(2), [], 1)
         self.assertFalse(env.get_unwrapped()[0].resetted)
         env.poll()
         self.assertTrue(env.get_unwrapped()[0].resetted)
 
-    def testVectorizeBasic(self):
+    def test_vectorize_basic(self):
         env = _MultiAgentEnvToBaseEnv(lambda v: BasicMultiAgent(2), [], 2)
         obs, rew, dones, _, _ = env.poll()
         self.assertEqual(obs, {0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
@@ -308,7 +314,7 @@ class TestMultiAgentEnv(unittest.TestCase):
                 }
             })
 
-    def testVectorizeRoundRobin(self):
+    def test_vectorize_round_robin(self):
         env = _MultiAgentEnvToBaseEnv(lambda v: RoundRobinMultiAgent(2), [], 2)
         obs, rew, dones, _, _ = env.poll()
         self.assertEqual(obs, {0: {0: 0}, 1: {0: 0}})
@@ -320,7 +326,7 @@ class TestMultiAgentEnv(unittest.TestCase):
         obs, rew, dones, _, _ = env.poll()
         self.assertEqual(obs, {0: {0: 0}, 1: {0: 0}})
 
-    def testMultiAgentSample(self):
+    def test_multi_agent_sample(self):
         act_space = gym.spaces.Discrete(2)
         obs_space = gym.spaces.Discrete(2)
         ev = RolloutWorker(
@@ -338,7 +344,7 @@ class TestMultiAgentEnv(unittest.TestCase):
         self.assertEqual(batch.policy_batches["p0"]["t"].tolist(),
                          list(range(25)) * 6)
 
-    def testMultiAgentSampleSyncRemote(self):
+    def test_multi_agent_sample_sync_remote(self):
         # Allow to be run via Unittest.
         ray.init(num_cpus=4, ignore_reinit_error=True)
         act_space = gym.spaces.Discrete(2)
@@ -357,7 +363,7 @@ class TestMultiAgentEnv(unittest.TestCase):
         batch = ev.sample()
         self.assertEqual(batch.count, 200)
 
-    def testMultiAgentSampleAsyncRemote(self):
+    def test_multi_agent_sample_async_remote(self):
         # Allow to be run via Unittest.
         ray.init(num_cpus=4, ignore_reinit_error=True)
         act_space = gym.spaces.Discrete(2)
@@ -375,7 +381,7 @@ class TestMultiAgentEnv(unittest.TestCase):
         batch = ev.sample()
         self.assertEqual(batch.count, 200)
 
-    def testMultiAgentSampleWithHorizon(self):
+    def test_multi_agent_sample_with_horizon(self):
         act_space = gym.spaces.Discrete(2)
         obs_space = gym.spaces.Discrete(2)
         ev = RolloutWorker(
@@ -390,7 +396,7 @@ class TestMultiAgentEnv(unittest.TestCase):
         batch = ev.sample()
         self.assertEqual(batch.count, 50)
 
-    def testSampleFromEarlyDoneEnv(self):
+    def test_sample_from_early_done_env(self):
         act_space = gym.spaces.Discrete(2)
         obs_space = gym.spaces.Discrete(2)
         ev = RolloutWorker(
@@ -406,7 +412,7 @@ class TestMultiAgentEnv(unittest.TestCase):
                                 ".*don't have a last observation.*",
                                 lambda: ev.sample())
 
-    def testMultiAgentSampleRoundRobin(self):
+    def test_multi_agent_sample_round_robin(self):
         act_space = gym.spaces.Discrete(2)
         obs_space = gym.spaces.Discrete(10)
         ev = RolloutWorker(
@@ -570,7 +576,7 @@ class TestMultiAgentEnv(unittest.TestCase):
             KeyError,
             lambda: pg.compute_action([0, 0, 0, 0], policy_id="policy_3"))
 
-    def _testWithOptimizer(self, optimizer_cls):
+    def _test_with_optimizer(self, optimizer_cls):
         n = 3
         env = gym.make("CartPole-v0")
         act_space = env.action_space
@@ -629,15 +635,15 @@ class TestMultiAgentEnv(unittest.TestCase):
         raise Exception("failed to improve reward")
 
     def test_multi_agent_sync_optimizer(self):
-        self._testWithOptimizer(SyncSamplesOptimizer)
+        self._test_with_optimizer(SyncSamplesOptimizer)
 
     def test_multi_agent_async_gradients_optimizer(self):
         # Allow to be run via Unittest.
         ray.init(num_cpus=4, ignore_reinit_error=True)
-        self._testWithOptimizer(AsyncGradientsOptimizer)
+        self._test_with_optimizer(AsyncGradientsOptimizer)
 
     def test_multi_agent_replay_optimizer(self):
-        self._testWithOptimizer(SyncReplayOptimizer)
+        self._test_with_optimizer(SyncReplayOptimizer)
 
     def test_train_multi_cartpole_many_policies(self):
         n = 20
@@ -668,5 +674,6 @@ class TestMultiAgentEnv(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=4)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_multi_agent_pendulum.py

@@ -1,38 +1,53 @@
 """Integration test: (1) pendulum works, (2) single-agent multi-agent works."""
+import unittest
 
 import ray
 from ray.rllib.tests.test_multi_agent_env import make_multiagent
 from ray.tune import run_experiments
 from ray.tune.registry import register_env
 
-if __name__ == "__main__":
-    ray.init()
-    MultiPendulum = make_multiagent("Pendulum-v0")
-    register_env("multi_pend", lambda _: MultiPendulum(1))
-    trials = run_experiments({
-        "test": {
-            "run": "PPO",
-            "env": "multi_pend",
-            "stop": {
-                "timesteps_total": 500000,
-                "episode_reward_mean": -200,
-            },
-            "config": {
-                "train_batch_size": 2048,
-                "vf_clip_param": 10.0,
-                "num_workers": 0,
-                "num_envs_per_worker": 10,
-                "lambda": 0.1,
-                "gamma": 0.95,
-                "lr": 0.0003,
-                "sgd_minibatch_size": 64,
-                "num_sgd_iter": 10,
-                "model": {
-                    "fcnet_hiddens": [64, 64],
+
+class TestMultiAgentPendulum(unittest.TestCase):
+    def setUp(self) -> None:
+        ray.init()
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_multi_agent_pendulum(self):
+        MultiPendulum = make_multiagent("Pendulum-v0")
+        register_env("multi_pend", lambda _: MultiPendulum(1))
+        trials = run_experiments({
+            "test": {
+                "run": "PPO",
+                "env": "multi_pend",
+                "stop": {
+                    "timesteps_total": 500000,
+                    "episode_reward_mean": -200,
                 },
-                "batch_mode": "complete_episodes",
-            },
-        }
-    })
-    if trials[0].last_result["episode_reward_mean"] < -200:
-        raise ValueError("Did not get to -200 reward", trials[0].last_result)
+                "config": {
+                    "train_batch_size": 2048,
+                    "vf_clip_param": 10.0,
+                    "num_workers": 0,
+                    "num_envs_per_worker": 10,
+                    "lambda": 0.1,
+                    "gamma": 0.95,
+                    "lr": 0.0003,
+                    "sgd_minibatch_size": 64,
+                    "num_sgd_iter": 10,
+                    "model": {
+                        "fcnet_hiddens": [64, 64],
+                    },
+                    "batch_mode": "complete_episodes",
+                },
+            }
+        })
+        if trials[0].last_result["episode_reward_mean"] < -200:
+            raise ValueError("Did not get to -200 reward",
+                             trials[0].last_result)
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_nested_spaces.py

@@ -1,8 +1,8 @@
 import pickle
 
+import gym
 from gym import spaces
 from gym.envs.registration import EnvSpec
-import gym
 import unittest
 
 import ray
@@ -24,7 +24,6 @@ from ray.rllib.utils import try_import_tf, try_import_torch
 tf = try_import_tf()
 _, nn = try_import_torch()
 
-
 DICT_SPACE = spaces.Dict({
     "sensors": spaces.Dict({
         "position": spaces.Box(low=-100, high=100, shape=(3, )),
@@ -218,7 +217,15 @@ class TupleSpyModel(Model):
 
 
 class NestedSpacesTest(unittest.TestCase):
-    def testInvalidModel(self):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=5)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_invalid_model(self):
         ModelCatalog.register_custom_model("invalid", InvalidModel)
         self.assertRaises(ValueError, lambda: PGTrainer(
             env="CartPole-v0", config={
@@ -227,7 +234,7 @@ class NestedSpacesTest(unittest.TestCase):
                 },
             }))
 
-    def testInvalidModel2(self):
+    def test_invalid_model2(self):
         ModelCatalog.register_custom_model("invalid2", InvalidModel2)
         self.assertRaisesRegexp(
             ValueError, "Expected output.*",
@@ -238,7 +245,7 @@ class NestedSpacesTest(unittest.TestCase):
                     },
                 }))
 
-    def doTestNestedDict(self, make_env, test_lstm=False):
+    def do_test_nested_dict(self, make_env, test_lstm=False):
         ModelCatalog.register_custom_model("composite", DictSpyModel)
         register_env("nested", make_env)
         pg = PGTrainer(
@@ -267,7 +274,7 @@ class NestedSpacesTest(unittest.TestCase):
             self.assertEqual(seen[1][0].tolist(), cam_i)
             self.assertEqual(seen[2][0].tolist(), task_i)
 
-    def doTestNestedTuple(self, make_env):
+    def do_test_nested_tuple(self, make_env):
         ModelCatalog.register_custom_model("composite2", TupleSpyModel)
         register_env("nested2", make_env)
         pg = PGTrainer(
@@ -294,36 +301,38 @@ class NestedSpacesTest(unittest.TestCase):
             self.assertEqual(seen[1][0].tolist(), cam_i)
             self.assertEqual(seen[2][0].tolist(), task_i)
 
-    def testNestedDictGym(self):
-        self.doTestNestedDict(lambda _: NestedDictEnv())
+    def test_nested_dict_gym(self):
+        self.do_test_nested_dict(lambda _: NestedDictEnv())
 
-    def testNestedDictGymLSTM(self):
-        self.doTestNestedDict(lambda _: NestedDictEnv(), test_lstm=True)
+    def test_nested_dict_gym_lstm(self):
+        self.do_test_nested_dict(lambda _: NestedDictEnv(), test_lstm=True)
 
-    def testNestedDictVector(self):
-        self.doTestNestedDict(
+    def test_nested_dict_vector(self):
+        self.do_test_nested_dict(
             lambda _: VectorEnv.wrap(lambda i: NestedDictEnv()))
 
-    def testNestedDictServing(self):
-        self.doTestNestedDict(lambda _: SimpleServing(NestedDictEnv()))
+    def test_nested_dict_serving(self):
+        self.do_test_nested_dict(lambda _: SimpleServing(NestedDictEnv()))
 
-    def testNestedDictAsync(self):
-        self.doTestNestedDict(lambda _: BaseEnv.to_base_env(NestedDictEnv()))
+    def test_nested_dict_async(self):
+        self.do_test_nested_dict(
+            lambda _: BaseEnv.to_base_env(NestedDictEnv()))
 
-    def testNestedTupleGym(self):
-        self.doTestNestedTuple(lambda _: NestedTupleEnv())
+    def test_nested_tuple_gym(self):
+        self.do_test_nested_tuple(lambda _: NestedTupleEnv())
 
-    def testNestedTupleVector(self):
-        self.doTestNestedTuple(
+    def test_nested_tuple_vector(self):
+        self.do_test_nested_tuple(
             lambda _: VectorEnv.wrap(lambda i: NestedTupleEnv()))
 
-    def testNestedTupleServing(self):
-        self.doTestNestedTuple(lambda _: SimpleServing(NestedTupleEnv()))
+    def test_nested_tuple_serving(self):
+        self.do_test_nested_tuple(lambda _: SimpleServing(NestedTupleEnv()))
 
-    def testNestedTupleAsync(self):
-        self.doTestNestedTuple(lambda _: BaseEnv.to_base_env(NestedTupleEnv()))
+    def test_nested_tuple_async(self):
+        self.do_test_nested_tuple(
+            lambda _: BaseEnv.to_base_env(NestedTupleEnv()))
 
-    def testMultiAgentComplexSpaces(self):
+    def test_multi_agent_complex_spaces(self):
         ModelCatalog.register_custom_model("dict_spy", DictSpyModel)
         ModelCatalog.register_custom_model("tuple_spy", TupleSpyModel)
         register_env("nested_ma", lambda _: NestedMultiAgentEnv())
@@ -373,7 +382,7 @@ class NestedSpacesTest(unittest.TestCase):
             self.assertEqual(seen[1][0].tolist(), cam_i)
             self.assertEqual(seen[2][0].tolist(), task_i)
 
-    def testRolloutDictSpace(self):
+    def test_rollout_dict_space(self):
         register_env("nested", lambda _: NestedDictEnv())
         agent = PGTrainer(env="nested")
         agent.train()
@@ -388,7 +397,7 @@ class NestedSpacesTest(unittest.TestCase):
         # Test rollout works on restore
         rollout(agent2, "nested", 100)
 
-    def testPyTorchModel(self):
+    def test_py_torch_model(self):
         ModelCatalog.register_custom_model("composite", TorchSpyModel)
         register_env("nested", lambda _: NestedDictEnv())
         a2c = A2CTrainer(
@@ -420,5 +429,6 @@ class NestedSpacesTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=5)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_optimizers.py

@@ -18,11 +18,15 @@ tf = try_import_tf()
 
 
 class LRScheduleTest(unittest.TestCase):
-    def tearDown(self):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=2)
+
+    @classmethod
+    def tearDownClass(cls):
         ray.shutdown()
 
-    def testBasic(self):
-        ray.init(num_cpus=2)
+    def test_basic(self):
         ppo = PPOTrainer(
             env="CartPole-v0",
             config={"lr_schedule": [[0, 1e-5], [1000, 0.0]]})
@@ -32,11 +36,15 @@ class LRScheduleTest(unittest.TestCase):
 
 
 class AsyncOptimizerTest(unittest.TestCase):
-    def tearDown(self):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=4, object_store_memory=1000 * 1024 * 1024)
+
+    @classmethod
+    def tearDownClass(cls):
         ray.shutdown()
 
-    def testBasic(self):
-        ray.init(num_cpus=4, object_store_memory=1000 * 1024 * 1024)
+    def test_basic(self):
         local = _MockWorker()
         remotes = ray.remote(_MockWorker)
         remote_workers = [remotes.remote() for i in range(5)]
@@ -47,12 +55,15 @@ class AsyncOptimizerTest(unittest.TestCase):
 
 
 class PPOCollectTest(unittest.TestCase):
-    def tearDown(self):
-        ray.shutdown()
-
-    def testPPOSampleWaste(self):
+    @classmethod
+    def setUpClass(cls):
         ray.init(num_cpus=4, object_store_memory=1000 * 1024 * 1024)
 
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_ppo_sample_waste(self):
         # Check we at least collect the initial wave of samples
         ppo = PPOTrainer(
             env="CartPole-v0",
@@ -92,7 +103,7 @@ class PPOCollectTest(unittest.TestCase):
 
 
 class SampleBatchTest(unittest.TestCase):
-    def testConcat(self):
+    def test_concat(self):
         b1 = SampleBatch({"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])})
         b2 = SampleBatch({"a": np.array([1]), "b": np.array([4])})
         b3 = SampleBatch({"a": np.array([1]), "b": np.array([5])})
@@ -105,34 +116,34 @@ class SampleBatchTest(unittest.TestCase):
 
 
 class AsyncSamplesOptimizerTest(unittest.TestCase):
-    @classmethod
-    def tearDownClass(cls):
-        ray.shutdown()
-
     @classmethod
     def setUpClass(cls):
         ray.init(num_cpus=8, object_store_memory=1000 * 1024 * 1024)
 
-    def testSimple(self):
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_simple(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(workers)
         self._wait_for(optimizer, 1000, 1000)
 
-    def testMultiGPU(self):
+    def test_multi_gpu(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(workers, num_gpus=1, _fake_gpus=True)
         self._wait_for(optimizer, 1000, 1000)
 
-    def testMultiGPUParallelLoad(self):
+    def test_multi_gpu_parallel_load(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(
             workers, num_gpus=1, num_data_loader_buffers=1, _fake_gpus=True)
         self._wait_for(optimizer, 1000, 1000)
 
-    def testMultiplePasses(self):
+    def test_multiple_passes(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(
@@ -145,7 +156,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
         self.assertLess(optimizer.stats()["num_steps_sampled"], 5000)
         self.assertGreater(optimizer.stats()["num_steps_trained"], 8000)
 
-    def testReplay(self):
+    def test_replay(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(
@@ -162,7 +173,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
         self.assertGreater(replay_ratio, 0.7)
         self.assertLess(stats["num_steps_trained"], stats["num_steps_sampled"])
 
-    def testReplayAndMultiplePasses(self):
+    def test_replay_and_multiple_passes(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(
@@ -181,7 +192,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
         replay_ratio = stats["num_steps_replayed"] / stats["num_steps_sampled"]
         self.assertGreater(replay_ratio, 0.7)
 
-    def testMultiTierAggregationBadConf(self):
+    def test_multi_tier_aggregation_bad_conf(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         aggregators = TreeAggregator.precreate_aggregators(4)
@@ -189,7 +200,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
         self.assertRaises(ValueError,
                           lambda: optimizer.aggregator.init(aggregators))
 
-    def testMultiTierAggregation(self):
+    def test_multi_tier_aggregation(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         aggregators = TreeAggregator.precreate_aggregators(1)
@@ -197,7 +208,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
         optimizer.aggregator.init(aggregators)
         self._wait_for(optimizer, 1000, 1000)
 
-    def testRejectBadConfigs(self):
+    def test_reject_bad_configs(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         self.assertRaises(
@@ -226,7 +237,7 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
             _fake_gpus=True)
         self._wait_for(optimizer, 1000, 1000)
 
-    def testLearnerQueueTimeout(self):
+    def test_learner_queue_timeout(self):
         local, remotes = self._make_envs()
         workers = WorkerSet._from_existing(local, remotes)
         optimizer = AsyncSamplesOptimizer(
@@ -265,4 +276,6 @@ class AsyncSamplesOptimizerTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_perf.py

@@ -8,10 +8,18 @@ from ray.rllib.tests.test_rollout_worker import MockPolicy
 
 
 class TestPerf(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=5)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
     # Tested on Intel(R) Core(TM) i7-4600U CPU @ 2.10GHz
     # 11/23/18: Samples per second 8501.125113727468
     # 03/01/19: Samples per second 8610.164353268685
-    def testBaselinePerformance(self):
+    def test_baseline_performance(self):
         for _ in range(20):
             ev = RolloutWorker(
                 env_creator=lambda _: gym.make("CartPole-v0"),
@@ -28,5 +36,6 @@ class TestPerf(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=5)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_pipeline.py

@@ -7,10 +7,12 @@ from ray.rllib.agents.a3c import A2CTrainer
 class TestPipeline(unittest.TestCase):
     """General tests for the pipeline API."""
 
-    def setUp(self):
+    @classmethod
+    def setUpClass(cls):
         ray.init()
 
-    def tearDown(self):
+    @classmethod
+    def tearDownClass(cls):
         ray.shutdown()
 
     def test_pipeline_stats(ray_start_regular):

rllib/tests/test_reproducibility.py

@@ -8,7 +8,7 @@ import gym
 
 
 class TestReproducibility(unittest.TestCase):
-    def testReproducingTrajectory(self):
+    def test_reproducing_trajectory(self):
         class PickLargest(gym.Env):
             def __init__(self):
                 self.observation_space = gym.spaces.Box(
@@ -61,4 +61,6 @@ class TestReproducibility(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_rollout.py

@@ -3,43 +3,49 @@
 
 from pathlib import Path
 import os
-import sys
+import unittest
+
+
+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)
+
+        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=IMPALA "
+                  "--checkpoint-freq=1 ".format(rllib_dir, tmp_dir) +
+                  "--config='{\"num_workers\": 1, \"num_gpus\": 0}' "
+                  "--env=Pong-ram-v4 --stop='{\"training_iteration\": 1}'")
+
+        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)
+
+        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)
+
+        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)
+
+        # Cleanup.
+        os.popen("rm -rf \"{}\"".format(tmp_dir)).read()
+
 
 if __name__ == "__main__":
-    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=IMPALA --checkpoint-freq=1 ".
-        format(rllib_dir, tmp_dir) +
-        "--config='{\"num_workers\": 1, \"num_gpus\": 0}' --env=Pong-ram-v4 "
-        "--stop='{\"training_iteration\": 1}'")
-
-    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)
-
-    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)
-
-    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)
-
-    # Cleanup.
-    os.popen("rm -rf \"{}\"".format(tmp_dir)).read()
-    print("OK")
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_rollout.sh

@@ -1,31 +0,0 @@
-#!/bin/bash -e
-
-TRAIN=/ray/rllib/train.py
-if [ ! -e "$TRAIN" ]; then
-    TRAIN=../train.py
-fi
-ROLLOUT=/ray/rllib/rollout.py
-if [ ! -e "$ROLLOUT" ]; then
-    ROLLOUT=../rollout.py
-fi
-
-TMP=`mktemp -d`
-echo "Saving results to $TMP"
-
-$TRAIN --local-dir=$TMP --run=IMPALA --checkpoint-freq=1 \
-    --config='{"num_workers": 1, "num_gpus": 0}' --env=Pong-ram-v4 \
-    --stop='{"training_iteration": 1}'
-find $TMP
-
-CHECKPOINT_PATH=`ls $TMP/default/*/checkpoint_1/checkpoint-1`
-echo "Checkpoint path $CHECKPOINT_PATH"
-test -e "$CHECKPOINT_PATH"
-
-$ROLLOUT --run=IMPALA "$CHECKPOINT_PATH" --steps=100 \
-    --out="$TMP/rollouts_100steps.pkl" --no-render
-test -e "$TMP/rollouts_100steps.pkl"
-$ROLLOUT --run=IMPALA "$CHECKPOINT_PATH" --episodes=1 \
-    --out="$TMP/rollouts_1episode.pkl" --no-render
-test -e "$TMP/rollouts_1episode.pkl"
-rm -rf "$TMP"
-echo "OK"

rllib/tests/test_rollout_worker.py

@@ -1,9 +1,9 @@
+from collections import Counter
 import gym
 import numpy as np
 import random
 import time
 import unittest
-from collections import Counter
 
 import ray
 from ray.rllib.agents.pg import PGTrainer
@@ -124,6 +124,14 @@ class MockVectorEnv(VectorEnv):
 
 
 class TestRolloutWorker(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(num_cpus=5)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
     def test_basic(self):
         ev = RolloutWorker(
             env_creator=lambda _: gym.make("CartPole-v0"), policy=MockPolicy)
@@ -452,5 +460,6 @@ class TestRolloutWorker(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    ray.init(num_cpus=5)
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/tests/test_supported_spaces.py

@@ -2,7 +2,6 @@ import gym
 from gym.spaces import Box, Discrete, Tuple, Dict, MultiDiscrete
 from gym.envs.registration import EnvSpec
 import numpy as np
-import sys
 import unittest
 import traceback
 
@@ -17,6 +16,7 @@ from ray.rllib.tests.test_multi_agent_env import MultiCartpole, \
     MultiMountainCar
 from ray.rllib.utils.error import UnsupportedSpaceException
 from ray.tune.registry import register_env
+
 tf = try_import_tf()
 
 ACTION_SPACES_TO_TEST = {
@@ -289,6 +289,9 @@ class ModelSupportedSpaces(unittest.TestCase):
 
 
 if __name__ == "__main__":
+    import pytest
+    import sys
+
     if len(sys.argv) > 1 and sys.argv[1] == "--smoke":
         ACTION_SPACES_TO_TEST = {
             "discrete": Discrete(5),
@@ -297,4 +300,5 @@ if __name__ == "__main__":
             "vector": Box(0.0, 1.0, (5, ), dtype=np.float32),
             "atari": Box(0.0, 1.0, (210, 160, 3), dtype=np.float32),
         }
-    unittest.main(verbosity=2)
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/utils/exploration/tests/test_explorations.py

@@ -14,12 +14,12 @@ import ray.rllib.agents.sac as sac
 from ray.rllib.utils import check
 
 
-def test_explorations(run,
-                      env,
-                      config,
-                      dummy_obs,
-                      prev_a=None,
-                      expected_mean_action=None):
+def do_test_explorations(run,
+                         env,
+                         config,
+                         dummy_obs,
+                         prev_a=None,
+                         expected_mean_action=None):
     """Calls an Agent's `compute_actions` with different `explore` options."""
 
     config = config.copy()
@@ -94,10 +94,17 @@ class TestExplorations(unittest.TestCase):
     Tests all Exploration components and the deterministic flag for
     compute_action calls.
     """
-    ray.init(ignore_reinit_error=True)
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(ignore_reinit_error=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
 
     def test_a2c(self):
-        test_explorations(
+        do_test_explorations(
             a3c.A2CTrainer,
             "CartPole-v0",
             a3c.DEFAULT_CONFIG,
@@ -105,7 +112,7 @@ class TestExplorations(unittest.TestCase):
             prev_a=np.array(1))
 
     def test_a3c(self):
-        test_explorations(
+        do_test_explorations(
             a3c.A3CTrainer,
             "CartPole-v0",
             a3c.DEFAULT_CONFIG,
@@ -113,7 +120,7 @@ class TestExplorations(unittest.TestCase):
             prev_a=np.array(1))
 
     def test_ddpg(self):
-        test_explorations(
+        do_test_explorations(
             ddpg.DDPGTrainer,
             "Pendulum-v0",
             ddpg.DEFAULT_CONFIG,
@@ -121,15 +128,16 @@ class TestExplorations(unittest.TestCase):
             expected_mean_action=0.0)
 
     def test_simple_dqn(self):
-        test_explorations(dqn.SimpleQTrainer, "CartPole-v0",
-                          dqn.DEFAULT_CONFIG, np.array([0.0, 0.1, 0.0, 0.0]))
+        do_test_explorations(dqn.SimpleQTrainer,
+                             "CartPole-v0", dqn.DEFAULT_CONFIG,
+                             np.array([0.0, 0.1, 0.0, 0.0]))
 
     def test_dqn(self):
-        test_explorations(dqn.DQNTrainer, "CartPole-v0", dqn.DEFAULT_CONFIG,
-                          np.array([0.0, 0.1, 0.0, 0.0]))
+        do_test_explorations(dqn.DQNTrainer, "CartPole-v0", dqn.DEFAULT_CONFIG,
+                             np.array([0.0, 0.1, 0.0, 0.0]))
 
     def test_impala(self):
-        test_explorations(
+        do_test_explorations(
             impala.ImpalaTrainer,
             "CartPole-v0",
             impala.DEFAULT_CONFIG,
@@ -137,7 +145,7 @@ class TestExplorations(unittest.TestCase):
             prev_a=np.array(0))
 
     def test_pg(self):
-        test_explorations(
+        do_test_explorations(
             pg.PGTrainer,
             "CartPole-v0",
             pg.DEFAULT_CONFIG,
@@ -145,7 +153,7 @@ class TestExplorations(unittest.TestCase):
             prev_a=np.array(1))
 
     def test_ppo_discr(self):
-        test_explorations(
+        do_test_explorations(
             ppo.PPOTrainer,
             "CartPole-v0",
             ppo.DEFAULT_CONFIG,
@@ -153,7 +161,7 @@ class TestExplorations(unittest.TestCase):
             prev_a=np.array(0))
 
     def test_ppo_cont(self):
-        test_explorations(
+        do_test_explorations(
             ppo.PPOTrainer,
             "Pendulum-v0",
             ppo.DEFAULT_CONFIG,
@@ -162,7 +170,7 @@ class TestExplorations(unittest.TestCase):
             expected_mean_action=0.0)
 
     def test_sac(self):
-        test_explorations(
+        do_test_explorations(
             sac.SACTrainer,
             "Pendulum-v0",
             sac.DEFAULT_CONFIG,
@@ -170,7 +178,7 @@ class TestExplorations(unittest.TestCase):
             expected_mean_action=0.0)
 
     def test_td3(self):
-        test_explorations(
+        do_test_explorations(
             td3.TD3Trainer,
             "Pendulum-v0",
             td3.TD3_DEFAULT_CONFIG,
@@ -179,4 +187,6 @@ class TestExplorations(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/utils/schedules/tests/test_schedules.py

@@ -114,5 +114,6 @@ class TestSchedules(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/utils/tests/test_framework_agnostic_components.py

@@ -15,6 +15,48 @@ tf.enable_eager_execution()
 torch, _ = try_import_torch()
 
 
+class DummyComponent:
+    """A simple class that can be used for testing framework-agnostic logic.
+
+    Implements a simple `add()` method for adding a value to
+    `self.prop_b`.
+    """
+
+    def __init__(self,
+                 prop_a,
+                 prop_b=0.5,
+                 prop_c=None,
+                 framework="tf",
+                 **kwargs):
+        self.framework = framework
+        self.prop_a = prop_a
+        self.prop_b = prop_b
+        self.prop_c = prop_c or "default"
+        self.prop_d = kwargs.pop("prop_d", 4)
+        self.kwargs = kwargs
+
+    def add(self, value):
+        if self.framework == "tf":
+            return self._add_tf(value)
+        return self.prop_b + value
+
+    def _add_tf(self, value):
+        return tf.add(self.prop_b, value)
+
+
+class NonAbstractChildOfDummyComponent(DummyComponent):
+    pass
+
+
+class AbstractDummyComponent(DummyComponent, metaclass=ABCMeta):
+    """Used for testing `from_config()`.
+    """
+
+    @abstractmethod
+    def some_abstract_method(self):
+        raise NotImplementedError
+
+
 class TestFrameWorkAgnosticComponents(unittest.TestCase):
     """
     Tests the Component base class to implement framework-agnostic functional
@@ -94,48 +136,7 @@ class TestFrameWorkAgnosticComponents(unittest.TestCase):
         check(component.add(-5.1), np.array([-6.6]))  # prop_b == -1.5
 
 
-class DummyComponent:
-    """A simple class that can be used for testing framework-agnostic logic.
-
-    Implements a simple `add()` method for adding a value to
-    `self.prop_b`.
-    """
-
-    def __init__(self,
-                 prop_a,
-                 prop_b=0.5,
-                 prop_c=None,
-                 framework="tf",
-                 **kwargs):
-        self.framework = framework
-        self.prop_a = prop_a
-        self.prop_b = prop_b
-        self.prop_c = prop_c or "default"
-        self.prop_d = kwargs.pop("prop_d", 4)
-        self.kwargs = kwargs
-
-    def add(self, value):
-        if self.framework == "tf":
-            return self._add_tf(value)
-        return self.prop_b + value
-
-    def _add_tf(self, value):
-        return tf.add(self.prop_b, value)
-
-
-class NonAbstractChildOfDummyComponent(DummyComponent):
-    pass
-
-
-class AbstractDummyComponent(DummyComponent, metaclass=ABCMeta):
-    """Used for testing `from_config()`.
-    """
-
-    @abstractmethod
-    def some_abstract_method(self):
-        raise NotImplementedError
-
-
 if __name__ == "__main__":
-    import unittest
-    unittest.main(verbosity=1)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/utils/tests/test_taskpool.py

@@ -135,4 +135,6 @@ class TaskPoolTest(unittest.TestCase):
 
 
 if __name__ == "__main__":
-    unittest.main(verbosity=2)
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))