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.
      """
      agent = self._create_test_agent()
      # We fill up the state with 9s. On calling agent.begin_episode the state
      # should be reset to all 0s.
      agent.state.fill(9)
      first_observation = onp.ones(self.observation_shape + (1,))
      self.assertEqual(agent.begin_episode(first_observation), 0)
      # 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] = onp.ones(self.observation_shape)
      self.assertTrue(onp.array_equal(agent.state, expected_state))
      self.assertTrue(
          onp.array_equal(
              agent._observation, 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.add_count = 0
      second_observation = onp.ones(self.observation_shape + (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] = onp.full(self.observation_shape, 2)
      self.assertTrue(onp.array_equal(agent.state, expected_state))
      self.assertTrue(
          onp.array_equal(
              agent._observation, second_observation[:, :, 0]
          )
      )
      # training_steps is incremented since we set eval_mode to False.
      self.assertEqual(agent.training_steps, 1)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import shutil
  from absl import flags
  from absl.testing import absltest
  from dopamine.discrete_domains import atari_lib
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.utils import test_utils
  from flax import linen
  import gin
  import jax.numpy
  import mock
  import numpy
label_desc: Use image node_id 5 for calling the DQNAgentTest obj's underlying member method code with example usage: obj.testBeginEpisode() without return types