code: def testPolicyStateReset(self):
      num_episodes = 2
      num_expected_steps = 6
  
      env = driver_test_utils.PyEnvironmentMock()
      policy = driver_test_utils.PyPolicyMock(
          env.time_step_spec(), env.action_spec()
      )
      replay_buffer_observer = MockReplayBufferObserver()
      driver = py_driver.PyDriver(
          env,
          policy,
          observers=[replay_buffer_observer],
          max_steps=None,
          max_episodes=num_episodes,
      )
  
      time_step = env.reset()
      policy_state = policy.get_initial_state()
      time_step, policy_state = driver.run(time_step, policy_state)
      trajectories = replay_buffer_observer.gather_all()
      self.assertEqual(
          trajectories, self._trajectories[:num_expected_steps]
      )
      self.assertEqual(
          num_episodes, policy.get_initial_state_call_count
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.drivers import py_driver
  from tf_agents.drivers import test_utils
  from tf_agents.environments import batched_py_environment
  from tf_agents.trajectories import trajectory
label_desc: Use image node_id 7 for calling the PyDriverTest obj's underlying member method code with example usage: obj.testPolicyStateReset() without return types