code: def test_multi_element_dataset_minibatch(
      self,
      num_epochs,
      num_parallel_environments,
      minibatch_size,
      expected_train_times,
  ):
      num_episodes = 3
      # Create a dataset with three elements. Each element represents an collected
      # episode of length 40.
      get_shape = lambda x: x.shape
      get_dtype = lambda x: tf.as_dtype(x.dtype)
      traj = _create_trajectories(
          n_time_steps=40, batch_size=num_parallel_environments
      )
      unused_info = ()
      shapes = tf.nest.map_structure(get_shape, (traj, unused_info))
      dtypes = tf.nest.map_structure(get_dtype, (traj, unused_info))
  
      def generate_data():
          for _ in range(num_episodes):
              yield (traj, unused_info)
  
      def dataset_fn():
          return tf.data.Dataset.from_generator(
              generate_data,
              dtypes,
              output_shapes=shapes,
          )
  
      fake_agent = FakePPOAgent()
  
      learner = ppo_learner.PPOLearner(
          root_dir=FLAGS.test_tmpdir,
          train_step=tf.Variable(0, dtype=tf.int32),
          agent=fake_agent,
          experience_dataset_fn=dataset_fn,
          normalization_dataset_fn=dataset_fn,
          num_samples=num_episodes,
          num_epochs=num_epochs,
          minibatch_size=minibatch_size,
          # Disable shuffling to have deterministic input into agent.train.
          shuffle_buffer_size=1,
          triggers=None,
      )
      learner.run()
  
      # Check that fake agent was called the expected number of times.
      self.assertEqual(
          fake_agent.train_called_times, expected_train_times
      )
  
      # Check that agent.train() is receiving the expected trajectories.
      if minibatch_size:
          concated_traj = _concat_and_flatten(
              traj, multiplier=num_episodes * num_epochs
          )
  
          for i in range(expected_train_times):
              expected_traj = _get_expected_minibatch(
                  concated_traj, minibatch_size, current_iteration=i
              )
              received_traj = fake_agent.experiences[i]
              tf.nest.map_structure(
                  self.assertAllClose, received_traj, expected_traj
              )
      else:
          for i in range(expected_train_times):
              expected_traj = traj
              received_traj = fake_agent.experiences[i]
              tf.nest.map_structure(
                  self.assertAllClose, received_traj, expected_traj
              )
  
imports: from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.agents.ppo import ppo_actor_network
  from tf_agents.agents.ppo import ppo_agent
  from tf_agents.networks import value_network
  from tf_agents.specs import tensor_spec
  from tf_agents.train import ppo_learner
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import test_utils
label_desc: Use image node_id 4 for calling the PpoLearnerTest obj's underlying member method code with example usage: obj.test_multi_element_dataset_minibatch(num_epochs, num_parallel_environments, minibatch_size, expected_train_times) and returns: tf