class_docstr: Manages all the learning details needed when training an PPO.
  
  These include:
    * Using distribution strategies correctly
    * Summaries
    * Checkpoints
    * Minimizing entering/exiting TF context:
        Especially in the case of TPUs scheduling a single TPU program to
        perform multiple train steps is critical for performance.
    * Generalizes the train call to be done correctly across CPU, GPU, or TPU
      executions managed by DistributionStrategies. This uses `strategy.run` and
      then makes sure to do a reduce operation over the `LossInfo` returned by
      the agent.
class_docstr_tok: ["Manages","all","the","learning","details","needed","when","training","an","PPO",".","These","include",":","*","Using","distribution","strategies","correctly","*","Summaries","*","Checkpoints","*","Minimizing","entering\/exiting","TF","context",":","Especially","in","the","case","of","TPUs","scheduling","a","single","TPU","program","to","perform","multiple","train","steps","is","critical","for","performance",".","*","Generalizes","the","train","call","to","be","done","correctly","across","CPU",",","GPU",",","or","TPU","executions","managed","by","DistributionStrategies",".","This","uses","`","strategy.run","`","and","then","makes","sure","to","do","a","reduce","operation","over","the","`","LossInfo","`","returned","by","the","agent","."]
docstr: Create the training dataset and iterator.
docstr_tok: ["Create","the","training","dataset","and","iterator","."]
code: def _create_datasets(self, strategy):
      """Create the training dataset and iterator."""
  
      def _make_dataset(_):
          train_dataset = self._experience_dataset_fn().take(
              self._num_samples
          )
  
          # We take the current batches, repeat for `num_epochs` times and exhaust
          # this data in the current learner run. The next time learner runs, new
          # batches of data will be sampled, cached and repeated.
          # This is enabled by the `Counter().flat_map()` trick below.
          train_dataset = train_dataset.cache().repeat(self._num_epochs)
  
          if self._minibatch_size:
  
              def squash_dataset_element(sequence, info):
                  return tf.nest.map_structure(
                      utils.BatchSquash(2).flatten, (sequence, info)
                  )
  
              # We unbatch the dataset shaped [B, T, ...] to a new dataset that
              # contains individual elements.
              # Note that we unbatch across the time dimension, which could result
              # in mini batches that contain subsets from more than one sequences.
              # PPO agent can handle mini batches across episode boundaries.
              train_dataset = train_dataset.map(
                  squash_dataset_element
              ).unbatch()
              train_dataset = train_dataset.shuffle(
                  self._shuffle_buffer_size
              )
              train_dataset = train_dataset.batch(
                  1, drop_remainder=True
              )
              train_dataset = train_dataset.batch(
                  self._minibatch_size, drop_remainder=True
              )
  
          return train_dataset
  
      def make_dataset(_):
          return tf.data.experimental.Counter().flat_map(_make_dataset)
  
      def _make_normalization_dataset(_):
          return self._normalization_dataset_fn()
  
      self._normalization_dataset = (
          tf.data.experimental.Counter().flat_map(
              _make_normalization_dataset
          )
      )
      self._normalization_iterator = iter(self._normalization_dataset)
  
      with strategy.scope():
          if strategy.num_replicas_in_sync > 1:
              self._train_dataset = (
                  strategy.distribute_datasets_from_function(
                      make_dataset
                  )
              )
          else:
              self._train_dataset = make_dataset(0)
          self._train_iterator = iter(self._train_dataset)
  
file_path: agents/tf_agents/train/ppo_learner.py
imports: from typing import Callable, Optional, Text
  import gin
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.agents.ppo import ppo_agent
  from tf_agents.networks import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.train import learner
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 2 for calling the PPOLearner obj's underlying member method code with example usage: obj._create_datasets(strategy) and returns: train_dataset, tf, self, tf