code: def test_train_agent_batch(
      num_envs, max_episode_len, steps, enable_evaluation
  ):
      outdir = tempfile.mkdtemp()
  
      agent = mock.Mock()
      agent.batch_act.side_effect = [[1] * num_envs] * 1000
  
      def make_env():
          env = mock.Mock()
          env.reset.side_effect = [("state", 0)] * 1000
          if max_episode_len is None:
              # Episodic env that terminates after 5 actions
              env.step.side_effect = [
                  (("state", 1), 0, False, {}),
                  (("state", 2), 0, False, {}),
                  (("state", 3), -0.5, False, {}),
                  (("state", 4), 0, False, {}),
                  (("state", 5), 1, True, {}),
              ] * 1000
          else:
              # Continuing env
              env.step.side_effect = [
                  (("state", 1), 0, False, {}),
              ] * 1000
          return env
  
      vec_env = pfrl.envs.SerialVectorEnv(
          [make_env() for _ in range(num_envs)]
      )
  
      hook = mock.Mock()
  
      if enable_evaluation:
          # evaluator.evaluate_if_necessary will be called `ceil(steps / num_envs)` times
          # during training. Here we simulate that eval_interval==steps,
          # i.e., return a float value (= 42 in this case) for the last call only and
          # otherwise return None (= evaluation is not necessary).
          evaluator = mock.Mock()
          n_evaluate_if_necessary_calls = math.ceil(steps / num_envs)
          dummy_eval_score = 42
          side_effect = [None] * (n_evaluate_if_necessary_calls - 1) + [
              dummy_eval_score
          ]
          evaluator.evaluate_if_necessary.side_effect = side_effect
  
          n_logging = 1  # Since all envs will reach to done==True simultaneously.
          n_valid_eval_score_returned = (
              1  # Since we simulated eval_interval==steps.
          )
          # agent.get_statistics will be called for logging & eval_stats_history
          n_get_statistics_calls = (
              n_logging + n_valid_eval_score_returned
          )
          dummy_stats = [
              ("average_q", 3.14),
              ("average_loss", 2.7),
              ("cumulative_steps", 42),
              ("n_updates", 8),
              ("rlen", 1),
          ]
          agent.get_statistics.side_effect = [
              dummy_stats
          ] * n_get_statistics_calls
      else:
          evaluator = None
  
      eval_stats_history = pfrl.experiments.train_agent_batch(
          agent=agent,
          env=vec_env,
          steps=steps,
          outdir=outdir,
          max_episode_len=max_episode_len,
          step_hooks=[hook],
          evaluator=evaluator,
      )
  
      if enable_evaluation:
          expected = [
              dict(**dict(dummy_stats), eval_score=dummy_eval_score)
              for _ in range(n_valid_eval_score_returned)
          ]
      else:
          # No evaluation invoked when evaluator=None is passed to train_agent_batch.
          expected = []
      assert eval_stats_history == expected
  
      iters = math.ceil(steps / num_envs)
      assert agent.batch_act.call_count == iters
      assert agent.batch_observe.call_count == iters
  
      for env in vec_env.envs:
          if max_episode_len is None:
              if num_envs == 1:
                  if steps == 6:
                      # In the beginning and after 5 iterations
                      assert env.reset.call_count == 2
                  else:
                      assert steps == 5
                      # Only in the beginning. While the last state is
                      # terminal, env.reset should not be called because
                      # training is complete.
                      assert env.reset.call_count == 1
              elif num_envs == 2:
                  # Only in the beginning
                  assert env.reset.call_count == 1
              else:
                  assert False
          elif max_episode_len == 2:
              if num_envs == 1:
                  # In the beginning, after 2 and 4 iterations
                  assert env.reset.call_count == 3
              elif num_envs == 2:
                  # In the beginning, after 2 iterations
                  assert env.reset.call_count == 2
              else:
                  assert False
          assert env.step.call_count == iters
  
      if steps % num_envs == 0:
          assert hook.call_count == steps
      else:
          assert hook.call_count == num_envs * iters
  
      # A hook receives (env, agent, step)
      for i, call in enumerate(hook.call_args_list):
          args, kwargs = call
          assert args[0] == vec_env
          assert args[1] == agent
          # step starts with 1
          assert args[2] == i + 1
  
      if enable_evaluation:
          assert (
              evaluator.evaluate_if_necessary.call_count
              == n_evaluate_if_necessary_calls
          )
  
file_path: pfrl/tests/experiments_tests/test_train_agent_batch.py
imports: import math
  import tempfile
  import unittest
  from unittest import mock
  import pytest
  import pfrl
label_desc: Use image node_id 1 for calling a global function with example usage: test_train_agent_batch(num_envs, max_episode_len, steps, enable_evaluation) and returns: env