docstr: Tests the functionality of agent.begin_episode.
  
  Specifically, the action returned and its effect on the state.
docstr_tok: ["Tests","the","functionality","of","agent.begin_episode",".","Specifically",",","the","action","returned","and","its","effect","on","the","state","."]
code: def testBeginEpisode(self):
      """Tests the functionality of agent.begin_episode.
  
      Specifically, the action returned and its effect on the 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(
          (1,) + self.observation_shape
      )
      onp.array_equal(agent.state, expected_state)
      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(
          (1,) + self.observation_shape, 2
      )
      onp.array_equal(agent.state, expected_state)
      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)
  
file_path: dopamine/tests/dopamine/jax/agents/full_rainbow/full_rainbow_agent_test.py
imports: from absl.testing import absltest
  from dopamine.discrete_domains import atari_lib
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  from dopamine.utils import test_utils
  from flax import linen
  import gin
  import jax.numpy
  import numpy
label_desc: Use image node_id 5 for calling the FullRainbowAgentTest obj's underlying member method code with example usage: obj.testBeginEpisode() without return types