code: def testSampleBatches(self, rb_cls):
      self._generate_replay_buffer(rb_cls=rb_cls)
  
      ds = self._replay_buffer.as_dataset(sample_batch_size=5)
      next_trajectory = next_dataset_element(self, ds)
      ds_structure = tf.data.experimental.get_structure(ds)
      self.assertEqual(
          list(ds_structure.observation.shape), [5, 15, 15, 4]
      )
      self.assertEqual(list(ds_structure.action.shape), [5])
      traj = next_trajectory()
      self.assertEqual(traj.observation.shape, (5, 15, 15, 4))
      self.assertEqual(traj.step_type.shape, (5,))
  
imports: from __future__ import division
  from __future__ import unicode_literals
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.replay_buffers import py_hashed_replay_buffer
  from tf_agents.replay_buffers import py_uniform_replay_buffer
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 9 for calling the PyUniformReplayBufferTest obj's underlying member method code with example usage: obj.testSampleBatches(rb_cls) without return types