code: def _get_expected_minibatch(
      all_traj, minibatch_size, current_iteration
  ):
      """Get the `Trajectory` containing the expected minibatch.
  
      Args:
        all_traj: a flattened `Trajectory` without the batch and time dimension.
        minibatch_size: the number of steps included in each minibatch.
        current_iteration: the indx of the current training iteration.
  
      Returns:
        The expected `Trajectory` shaped [minibatch_size, 1, ...].
      """
      expected_traj = tf.nest.map_structure(
          # pylint: disable=g-long-lambda
          lambda x: x[
              minibatch_size
              * current_iteration : minibatch_size
              * (current_iteration + 1)
          ],
          all_traj,
      )
      # Add time dimension to be consistent with the input to agent.train.
      expected_traj = tf.nest.map_structure(
          lambda x: tf.expand_dims(x, 1), expected_traj
      )
      return 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 3 for calling a global function with example usage: _get_expected_minibatch(all_traj, minibatch_size, current_iteration) and returns: expected_traj