class_docstr: Builder class for MPO agent components.
code: def make_dataset_iterator(
      self, replay_client: reverb.Client
  ) -> Iterator[reverb.ReplaySample]:
      if self.config.num_stacked_observations > 1:
          maybe_stack_observations = functools.partial(
              obs_stacking.stack_reverb_observation,
              stack_size=self.config.num_stacked_observations,
          )
      else:
          maybe_stack_observations = None
  
      dataset = datasets.make_reverb_dataset(
          server_address=replay_client.server_address,
          batch_size=self.config.batch_size // jax.device_count(),
          table={
              adders.DEFAULT_PRIORITY_TABLE: self.config.replay_fraction,
              _QUEUE_TABLE_NAME: 1.0 - self.config.replay_fraction,
          },
          num_parallel_calls=max(16, 4 * jax.local_device_count()),
          max_in_flight_samples_per_worker=(
              2
              * self.sgd_steps_per_learner_step
              * self.config.batch_size
              // jax.device_count()
          ),
          postprocess=maybe_stack_observations,
      )
  
      if self.config.observation_transform:
          # Augment dataset with random translations, simulated by pad-and-crop.
          transform = img_aug.make_transform(
              observation_transform=self.config.observation_transform,
              transform_next_observation=isinstance(
                  self.config.experience_type, mpo_types.FromTransitions
              ),
          )
          dataset = dataset.map(
              transform, num_parallel_calls=16, deterministic=False
          )
  
      # Batch and then flatten to feed multiple SGD steps per learner step.
      if self.sgd_steps_per_learner_step > 1:
          dataset = dataset.batch(
              self.sgd_steps_per_learner_step, drop_remainder=True
          )
          batch_flatten = lambda t: tf.reshape(
              t, [-1] + t.shape[2:].as_list()
          )
          dataset = dataset.map(
              lambda x: tree.map_structure(batch_flatten, x)
          )
  
      return utils.multi_device_put(
          dataset.as_numpy_iterator(), jax.local_devices()
      )
  
imports: import functools
  from typing import Iterator, List, Optional
  from absl import logging
  from acme import core
  from acme import specs
  from acme.adders import base
  from acme.adders import reverb
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import builders
  from acme.agents.jax.mpo import acting
  from acme.agents.jax.mpo import config
  from acme.agents.jax.mpo import learning
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import types
  from acme.datasets import image_augmentation
  from acme.datasets import reverb
  from acme.jax import observation_stacking
  from acme.jax import types
  from acme.jax import utils
  from acme.jax import variable_utils
  from acme.utils import counting
  from acme.utils import loggers
  import chex
  import jax
  import optax
  import reverb
  import tensorflow
  import tree
label_desc: Use image node_id 7 for calling the MPOBuilder obj's underlying member method code with example usage: obj.make_dataset_iterator(replay_client) and returns: utils