From 2aec77e305c753e36fa588d6128a59de0f552e82 Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Fri, 16 Oct 2020 20:53:30 +0200
Subject: [PATCH] [RLlib] Fix two test cases that only fail on Travis. (#11435)

---
 rllib/evaluation/observation_function.py           |  2 +-
 rllib/tests/test_evaluators.py                     | 11 +++++++++--
 rllib/tests/test_exec_api.py                       |  2 +-
 rllib/utils/exploration/tests/test_explorations.py |  2 +-
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/rllib/evaluation/observation_function.py b/rllib/evaluation/observation_function.py
index 91849a735..90d044ac4 100644
--- a/rllib/evaluation/observation_function.py
+++ b/rllib/evaluation/observation_function.py
@@ -14,7 +14,7 @@ class ObservationFunction:
     in multi-agent scenarios.
 
     Observation functions can be specified in the multi-agent config by
-    specifying ``{"observation_function": your_obs_func}``. Note that
+    specifying ``{"observation_fn": your_obs_func}``. Note that
     ``your_obs_func`` can be a plain Python function.
 
     This API is **experimental**.
diff --git a/rllib/tests/test_evaluators.py b/rllib/tests/test_evaluators.py
index 5b513fd1b..b258dd65c 100644
--- a/rllib/tests/test_evaluators.py
+++ b/rllib/tests/test_evaluators.py
@@ -10,6 +10,14 @@ from ray.tune.registry import register_env
 
 
 class EvalTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init(object_store_memory=1000 * 1024 * 1024, num_cpus=4)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
     def test_dqn_n_step(self):
         obs = [1, 2, 3, 4, 5, 6, 7]
         actions = ["a", "b", "a", "a", "a", "b", "a"]
@@ -32,7 +40,6 @@ class EvalTest(unittest.TestCase):
 
         for agent_cls in agent_classes:
             for fw in framework_iterator(frameworks=("torch", "tf")):
-                ray.init(object_store_memory=1000 * 1024 * 1024)
                 register_env("CartPoleWrapped-v0", env_creator)
                 agent = agent_cls(
                     env="CartPoleWrapped-v0",
@@ -53,6 +60,7 @@ class EvalTest(unittest.TestCase):
                 r1 = agent.train()
                 r2 = agent.train()
                 r3 = agent.train()
+                agent.stop()
 
                 self.assertTrue("evaluation" in r1)
                 self.assertTrue("evaluation" in r3)
@@ -60,7 +68,6 @@ class EvalTest(unittest.TestCase):
                 self.assertFalse("evaluation" in r2)
                 self.assertTrue("episode_reward_mean" in r1["evaluation"])
                 self.assertNotEqual(r1["evaluation"], r3["evaluation"])
-                ray.shutdown()
 
 
 if __name__ == "__main__":
diff --git a/rllib/tests/test_exec_api.py b/rllib/tests/test_exec_api.py
index a4853e9cd..df9a5015e 100644
--- a/rllib/tests/test_exec_api.py
+++ b/rllib/tests/test_exec_api.py
@@ -10,7 +10,7 @@ class TestDistributedExecution(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        ray.init(ignore_reinit_error=True)
+        ray.init(num_cpus=4)
 
     @classmethod
     def tearDownClass(cls):
diff --git a/rllib/utils/exploration/tests/test_explorations.py b/rllib/utils/exploration/tests/test_explorations.py
index 99520b57d..61ec1d3e9 100644
--- a/rllib/utils/exploration/tests/test_explorations.py
+++ b/rllib/utils/exploration/tests/test_explorations.py
@@ -82,7 +82,7 @@ class TestExplorations(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        ray.init(ignore_reinit_error=True)
+        ray.init(num_cpus=4)
 
     @classmethod
     def tearDownClass(cls):