docstr: Test the functionality of agent.begin_episode.
  
  Specifically, the action returned and its effect on state.
docstr_tok: ["Test","the","functionality","of","agent.begin_episode",".","Specifically",",","the","action","returned","and","its","effect","on","state","."]
code: 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)
  
file_path: tensor2tensor/tensor2tensor/rl/batch_dqn_agent_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import shutil
  from absl import flags
  from dopamine.agents.dqn import dqn_agent
  from dopamine.discrete_domains import atari_lib
  import numpy
  from tensor2tensor.rl import dopamine_connector
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling the BatchDQNAgentTest obj's underlying member method code with example usage: obj.testBeginEpisode() without return types