def testBeginEpisode(self):
    """Test the functionality of agent.begin_episode.

    Specifically, the action returned and its effect on state.
    """
    with tf.Session() as sess:
      agent = self._create_test_agent(sess)
      # We fill up the state with 9s. On calling agent.begin_episode the state
      # should be reset to all 0s.
      agent.state_batch.fill(9)
      first_observation = np.ones(
          [self.env_batch_size, self.observation_shape[0],
           self.observation_shape[1], 1])
      self.assertTrue((agent.begin_episode(first_observation) == 0).all())
      # When the all-1s observation is received, it will be placed at the end of
      # the state.
      expected_state = self.zero_state
      expected_state[:, :, :, -1] = np.ones(
          [self.env_batch_size, self.observation_shape[0],
           self.observation_shape[1]])
      self.assertAllEqual(agent.state_batch, expected_state)
      self.assertAllEqual(agent._observation_batch, first_observation[..., 0])
      # No training happens in eval mode.
      self.assertEqual(agent.training_steps, 0)

      # This will now cause training to happen.
      agent.eval_mode = False
      # Having a low replay memory add_count will prevent any of the
      # train/prefetch/sync ops from being called.
      agent._replay.memory.add_count = 0
      second_observation = np.ones(
          [self.env_batch_size, self.observation_shape[0],
           self.observation_shape[1], 1]) * 2
      agent.begin_episode(second_observation)
      # The agent's state will be reset, so we will only be left with the all-2s
      # observation.
      expected_state[:, :, :, -1] = np.full(
          (self.env_batch_size, self.observation_shape[0],
           self.observation_shape[1]), 2
      )
      self.assertAllEqual(agent.state_batch, expected_state)
      self.assertAllEqual(agent._observation_batch,
                          second_observation[:, :, :, 0])
      # training_steps is incremented since we set eval_mode to False.
      self.assertEqual(agent.training_steps, 1)
Test the functionality of agent.begin_episode.

    Specifically, the action returned and its effect on state.
    
tensor2tensor/rl/batch_dqn_agent_test.py