adds batch tests and does some renaming in order to pass flakes

tests/experiments_tests/test_evaluator.py

@@ -13,6 +13,7 @@
 import mock
 
 import chainerrl
+from chainerrl.experiments import evaluator
 
 
 @testing.parameterize(
@@ -34,7 +35,7 @@ def test_evaluate_if_necessary(self):
         env.reset.return_value = 'obs'
         env.step.return_value = ('obs', 0, True, {})
 
-        evaluator = chainerrl.experiments.evaluator.Evaluator(
+        agent_evaluator = evaluator.Evaluator(
             agent=agent,
             env=env,
             n_steps=None,
@@ -46,14 +47,14 @@ def test_evaluate_if_necessary(self):
             save_best_so_far_agent=self.save_best_so_far_agent,
         )
 
-        evaluator.evaluate_if_necessary(t=1, episodes=1)
+        agent_evaluator.evaluate_if_necessary(t=1, episodes=1)
         self.assertEqual(agent.act.call_count, 0)
 
-        evaluator.evaluate_if_necessary(t=2, episodes=2)
+        agent_evaluator.evaluate_if_necessary(t=2, episodes=2)
         self.assertEqual(agent.act.call_count, 0)
 
         # First evaluation
-        evaluator.evaluate_if_necessary(t=3, episodes=3)
+        agent_evaluator.evaluate_if_necessary(t=3, episodes=3)
         self.assertEqual(agent.act.call_count, self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, self.n_runs)
         if self.save_best_so_far_agent:
@@ -62,7 +63,7 @@ def test_evaluate_if_necessary(self):
             self.assertEqual(agent.save.call_count, 0)
 
         # Second evaluation with the same score
-        evaluator.evaluate_if_necessary(t=6, episodes=6)
+        agent_evaluator.evaluate_if_necessary(t=6, episodes=6)
         self.assertEqual(agent.act.call_count, 2 * self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, 2 * self.n_runs)
         if self.save_best_so_far_agent:
@@ -72,7 +73,7 @@ def test_evaluate_if_necessary(self):
 
         # Third evaluation with a better score
         env.step.return_value = ('obs', 1, True, {})
-        evaluator.evaluate_if_necessary(t=9, episodes=9)
+        agent_evaluator.evaluate_if_necessary(t=9, episodes=9)
         self.assertEqual(agent.act.call_count, 3 * self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, 3 * self.n_runs)
         if self.save_best_so_far_agent:
@@ -100,7 +101,7 @@ def test_evaluate_if_necessary(self):
         env.reset.return_value = 'obs'
         env.step.return_value = ('obs', 0, True, {})
 
-        evaluator = chainerrl.experiments.evaluator.AsyncEvaluator(
+        agent_evaluator = evaluator.AsyncEvaluator(
             n_runs=self.n_runs,
             eval_interval=3,
             outdir=outdir,
@@ -109,14 +110,17 @@ def test_evaluate_if_necessary(self):
             save_best_so_far_agent=self.save_best_so_far_agent,
         )
 
-        evaluator.evaluate_if_necessary(t=1, episodes=1, env=env, agent=agent)
+        agent_evaluator.evaluate_if_necessary(
+            t=1, episodes=1, env=env, agent=agent)
         self.assertEqual(agent.act.call_count, 0)
 
-        evaluator.evaluate_if_necessary(t=2, episodes=2, env=env, agent=agent)
+        agent_evaluator.evaluate_if_necessary(
+            t=2, episodes=2, env=env, agent=agent)
         self.assertEqual(agent.act.call_count, 0)
 
         # First evaluation
-        evaluator.evaluate_if_necessary(t=3, episodes=3, env=env, agent=agent)
+        agent_evaluator.evaluate_if_necessary(
+            t=3, episodes=3, env=env, agent=agent)
         self.assertEqual(agent.act.call_count, self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, self.n_runs)
         if self.save_best_so_far_agent:
@@ -125,7 +129,8 @@ def test_evaluate_if_necessary(self):
             self.assertEqual(agent.save.call_count, 0)
 
         # Second evaluation with the same score
-        evaluator.evaluate_if_necessary(t=6, episodes=6, env=env, agent=agent)
+        agent_evaluator.evaluate_if_necessary(
+            t=6, episodes=6, env=env, agent=agent)
         self.assertEqual(agent.act.call_count, 2 * self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, 2 * self.n_runs)
         if self.save_best_so_far_agent:
@@ -135,14 +140,16 @@ def test_evaluate_if_necessary(self):
 
         # Third evaluation with a better score
         env.step.return_value = ('obs', 1, True, {})
-        evaluator.evaluate_if_necessary(t=9, episodes=9, env=env, agent=agent)
+        agent_evaluator.evaluate_if_necessary(
+            t=9, episodes=9, env=env, agent=agent)
         self.assertEqual(agent.act.call_count, 3 * self.n_runs)
         self.assertEqual(agent.stop_episode.call_count, 3 * self.n_runs)
         if self.save_best_so_far_agent:
             self.assertEqual(agent.save.call_count, 2)
         else:
             self.assertEqual(agent.save.call_count, 0)
 
+
 @testing.parameterize(
     *testing.product({
         'n_episodes': [None, 1],
@@ -167,13 +174,13 @@ def test_timesteps(self):
         ]
 
         if self.n_episodes:
-            with self.assertRaises(AssertionError) as _:
-                    scores = chainerrl.experiments.evaluator.run_evaluation_episodes(
-                    env, agent,
-                    n_steps=self.n_timesteps,
-                    n_episodes=self.n_episodes)
+            with self.assertRaises(AssertionError):
+                    scores = evaluator.run_evaluation_episodes(
+                        env, agent,
+                        n_steps=self.n_timesteps,
+                        n_episodes=self.n_episodes)
         else:
-            scores = chainerrl.experiments.evaluator.run_evaluation_episodes(
+            scores = evaluator.run_evaluation_episodes(
                 env, agent,
                 n_steps=self.n_timesteps,
                 n_episodes=self.n_episodes)
@@ -188,7 +195,6 @@ def test_timesteps(self):
                 self.assertAlmostEqual(scores[0], 0.6)
                 self.assertAlmostEqual(scores[1], 0.5)
 
-
     def test_needs_reset(self):
         agent = mock.Mock()
         env = mock.Mock()
@@ -203,15 +209,83 @@ def test_needs_reset(self):
             (('state', 6), 0, False, {}),
             (('state', 7), 1, True, {}),
         ]
-        scores = chainerrl.experiments.evaluator.run_evaluation_episodes(
+        scores = evaluator.run_evaluation_episodes(
             env, agent, n_steps=None, n_episodes=2)
         self.assertAlmostEqual(len(scores), 2)
         self.assertAlmostEqual(scores[0], 0)
         self.assertAlmostEqual(scores[1], 0.5)
 
 
+@testing.parameterize(
+    *testing.product({
+        'n_episodes': [None, 1],
+        'n_timesteps': [2, 5, 6],
+    })
+)
 class TestBatchRunEvaluationEpisode(unittest.TestCase):
 
+    def test_timesteps(self):
+        agent = mock.Mock()
+        agent.batch_act.side_effect = [[1, 1]] * 5
+
+        def make_env(idx):
+            env = mock.Mock()
+            if idx == 0:
+                # First episode: 0 -> 1 -> 2 -> 3 (reset)
+                # Second episode: 4 -> 5 -> 6 -> 7 (done)
+                env.reset.side_effect = [('state', 0), ('state', 4)]
+                env.step.side_effect = [
+                    (('state', 1), 0, False, {}),
+                    (('state', 2), 0.1, False, {}),
+                    (('state', 3), 0.2, False, {'needs_reset': True}),
+                    (('state', 5), -0.5, False, {}),
+                    (('state', 6), 0, False, {}),
+                    (('state', 7), 1, True, {}),
+                ]
+            else:
+                # First episode: 0 -> 1 (reset)
+                # Second episode: 2 -> 3 (reset)
+                # Third episode: 4 -> 5 -> 6 -> 7 (done)
+                env.reset.side_effect = [
+                    ('state', 0), ('state', 2), ('state', 4)]
+                env.step.side_effect = [
+                    (('state', 1), 2, False, {'needs_reset': True}),
+                    (('state', 3), 3, False, {'needs_reset': True}),
+                    (('state', 5), -0.6, False, {}),
+                    (('state', 6), 0, False, {}),
+                    (('state', 7), 1, True, {}),
+                ]
+            return env
+
+        vec_env = chainerrl.envs.SerialVectorEnv(
+            [make_env(i) for i in range(2)])
+        if self.n_episodes:
+            with self.assertRaises(AssertionError):
+                scores = evaluator.batch_run_evaluation_episodes(
+                    vec_env, agent,
+                    n_steps=self.n_timesteps,
+                    n_episodes=self.n_episodes)
+        else:
+            # First Env:  [1   2   (3_a)  5  6   (7_a)]
+            # Second Env: [(1)(3_b) 5     6 (7_b)]
+            scores = evaluator.batch_run_evaluation_episodes(
+                    vec_env, agent,
+                    n_steps=self.n_timesteps,
+                    n_episodes=self.n_episodes)
+            if self.n_timesteps == 2:
+                self.assertAlmostEqual(len(scores), 1)
+                self.assertAlmostEqual(scores[0], 0.1)
+            elif self.n_timesteps == 5:
+                self.assertAlmostEqual(len(scores), 3)
+                self.assertAlmostEqual(scores[0], 0.3)
+                self.assertAlmostEqual(scores[1], 2.0)
+                self.assertAlmostEqual(scores[2], 3.0)
+            else:
+                self.assertAlmostEqual(len(scores), 3)
+                self.assertAlmostEqual(scores[0], 0.3)
+                self.assertAlmostEqual(scores[1], 2.0)
+                self.assertAlmostEqual(scores[2], 3.0)
+
     def test_needs_reset(self):
         agent = mock.Mock()
         agent.batch_act.side_effect = [[1, 1]] * 5
@@ -251,7 +325,7 @@ def make_env(idx):
         # First Env: [1 2 (3_a) 5 6 (7_a)]
         # Second Env: [(1) (3_b) 5 6 (7_b)]
         # Results: (1), (3a), (3b), (7b)
-        scores = chainerrl.experiments.evaluator.batch_run_evaluation_episodes(
+        scores = evaluator.batch_run_evaluation_episodes(
             vec_env, agent, n_steps=None, n_episodes=4)
         self.assertAlmostEqual(len(scores), 4)
         self.assertAlmostEqual(scores[0], 0)