class_docstr: PPO Builder.
docstr: Creates a dataset.
  
  The iterator batch size is computed as follows:
  
  Let:
    B := learner batch size (config.batch_size)
    H := number of hosts (jax.process_count())
    D := number of local devices per host
  
  The Reverb iterator will load batches of size B // (H * D). After wrapping
  the iterator with utils.multi_device_put, this will result in an iterable
  that provides B // H samples per item, with B // (H * D) samples placed on
  each local device. In a multi-host setup, each host has its own learner
  node and builds its own instance of the iterator. This will result
  in a total batch size of H * (B // H) == B being consumed per learner
  step (since the learner is pmapped across all devices). Note that
  jax.device_count() returns the total number of devices across hosts,
  i.e. H * D.
  
  Args:
    replay_client: the reverb replay client
  
  Returns:
    A replay buffer iterator to be used by the local devices.
docstr_tok: ["Creates","a","dataset",".","The","iterator","batch","size","is","computed","as","follows",":","Let",":","B",":","=","learner","batch","size","(","config.batch_size",")","H",":","=","number","of","hosts","(","jax.process_count","(",")",")","D",":","=","number","of","local","devices","per","host","The","Reverb","iterator","will","load","batches","of","size","B","\/\/","(","H","*","D",")",".","After","wrapping","the","iterator","with","utils.multi_device_put",",","this","will","result","in","an","iterable","that","provides","B","\/\/","H","samples","per","item",",","with","B","\/\/","(","H","*","D",")","samples","placed","on","each","local","device",".","In","a","multi-host","setup",",","each","host","has","its","own","learner","node","and","builds","its","own","instance","of","the","iterator",".","This","will","result","in","a","total","batch","size","of","H","*","(","B","\/\/","H",")","==","B","being","consumed","per","learner","step","(","since","the","learner","is","pmapped","across","all","devices",")",".","Note","that","jax.device_count","(",")","returns","the","total","number","of","devices","across","hosts",",","i.e",".","H","*","D.","Args",":","replay_client",":","the","reverb","replay","client","Returns",":","A","replay","buffer","iterator","to","be","used","by","the","local","devices","."]
code: def make_dataset_iterator(
      self, replay_client: reverb.Client
  ) -> Iterator[reverb.ReplaySample]:
      """Creates a dataset.
  
      The iterator batch size is computed as follows:
  
      Let:
        B := learner batch size (config.batch_size)
        H := number of hosts (jax.process_count())
        D := number of local devices per host
  
      The Reverb iterator will load batches of size B // (H * D). After wrapping
      the iterator with utils.multi_device_put, this will result in an iterable
      that provides B // H samples per item, with B // (H * D) samples placed on
      each local device. In a multi-host setup, each host has its own learner
      node and builds its own instance of the iterator. This will result
      in a total batch size of H * (B // H) == B being consumed per learner
      step (since the learner is pmapped across all devices). Note that
      jax.device_count() returns the total number of devices across hosts,
      i.e. H * D.
  
      Args:
        replay_client: the reverb replay client
  
      Returns:
        A replay buffer iterator to be used by the local devices.
      """
      iterator_batch_size, ragged = divmod(
          self._config.batch_size, jax.device_count()
      )
      if ragged:
          raise ValueError(
              "Learner batch size must be divisible by total number of devices!"
          )
  
      # We don't use datasets.make_reverb_dataset() here to avoid interleaving
      # and prefetching, that doesn't work well with can_sample() check on update.
      # NOTE: Value for max_in_flight_samples_per_worker comes from a
      # recommendation here: https://git.io/JYzXB
      dataset = reverb.TrajectoryDataset.from_table_signature(
          server_address=replay_client.server_address,
          table=self._config.replay_table_name,
          max_in_flight_samples_per_worker=(
              2 * self._config.batch_size // jax.process_count()
          ),
      )
      dataset = dataset.batch(iterator_batch_size, drop_remainder=True)
      dataset = dataset.as_numpy_iterator()
      return utils.multi_device_put(
          iterable=dataset, devices=jax.local_devices()
      )
  
imports: from typing import Iterator, List, Optional
  from acme import adders
  from acme import core
  from acme import specs
  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.ppo import config
  from acme.agents.jax.ppo import learning
  from acme.agents.jax.ppo import networks
  from acme.agents.jax.ppo import normalization
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax import variable_utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import numpy
  import optax
  import reverb
label_desc: Use image node_id 3 for calling the PPOBuilder obj's underlying member method code with example usage: obj.make_dataset_iterator(replay_client) and returns: utils