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: Initializes a PPOLearner instance.
  
  ```python
  agent = ppo_agent.PPOAgent(...,
    compute_value_and_advantage_in_train=False,
    # Skips updating normalizers in the agent, as it's handled in the learner.
    update_normalizers_in_train=False)
  
  # train_replay_buffer and normalization_replay_buffer point to two Reverb
  # tables that are synchronized. Sampling is done in a FIFO fashion.
  def experience_dataset_fn():
    return train_replay_buffer.as_dataset(sample_batch_size,
      sequence_preprocess_fn=agent.preprocess_sequence)
  def normalization_dataset_fn():
    return normalization_replay_buffer.as_dataset(sample_batch_size,
      sequence_preprocess_fn=agent.preprocess_sequence)
  
  learner = PPOLearner(..., agent, experience_dataset_fn,
    normalization_dataset_fn)
  learner.run()
  ```
  
  Args:
    root_dir: Main directory path where checkpoints, saved_models, and
      summaries will be written to.
    train_step: a scalar tf.int64 `tf.Variable` which will keep track of the
      number of train steps. This is used for artifacts created like
      summaries, or outputs in the root_dir.
    agent: `ppo_agent.PPOAgent` instance to train with. Note that
      update_normalizers_in_train should be set to `False`, otherwise a
      ValueError will be raised. We do not update normalizers in the agent
      again because we already update it in the learner. When mini batching is
      enabled, compute_value_and_advantage_in_train should be set to False,
      and preprocessing should be done as part of the data pipeline as part of
      `replay_buffer.as_dataset`.
    experience_dataset_fn: a function that will create an instance of a
      tf.data.Dataset used to sample experience for training. Each element in
      the dataset is a (Trajectory, SampleInfo) pair. Note that each
      Trajectory in itself might represent an episode (variable length) or a
      collection of transitions depending on how the replay buffer and hence
      the `experience_dataset_fn` was set up.
    normalization_dataset_fn: a function that will create an instance of a
      tf.data.Dataset used for normalization. This dataset is often from a
      separate reverb table that is synchronized with the table used in
      experience_dataset_fn. Each element in the dataset is a (Trajectory,
      SampleInfo) pair. Note that each Trajectory in itself might represent an
      episode (variable length) or a collection of transitions depending on
      how the replay buffer and hence the `experience_dataset_fn` was set up.
    num_samples: The number of samples used for training and normalization. A
      sample here means one reading of the experience_dataset_fn which in turn
      might  have more than 1 Trajectories (For ex. if this is coming from a
      `reverb_replay_buffer` with `sample_batch_size` = 2, then one sample
      means 2 trajectories, and the trainer will pull in num_samples * 2
      trajectories for one iteration of training). If fewer than this amount
      of batches exists in the dataset, the learner will wait for more data to
      be added, or until the reverb timeout is reached.
    num_epochs: The number of iterations to go through the same sequences.
    minibatch_size: The minibatch size. If set, the input data set will be
      flattened, and the dataset used for training is shaped `[minibatch_size,
      1, ...]`, where each element in the dataset represents a single step. If
      None, full sequences will be fed into the agent. Please set this
      parameter to None for RNN networks which requires full sequences.
    shuffle_buffer_size: The buffer size for shuffling the trajectories before
      splitting them into mini batches. Only required when mini batch learning
      is enabled (minibatch_size is set). Otherwise it is ignored. Commonly
      set to a number 1-3x the episode length of your environment.
    after_train_strategy_step_fn: (Optional) callable of the form `fn(sample,
      loss)` which can be used for example to update priorities in a replay
      buffer where sample is pulled from the `experience_iterator` and loss is
      a `LossInfo` named tuple returned from the agent. This is called after
      every train step. It runs using `strategy.run(...)`.
    triggers: List of callables of the form `trigger(train_step)`. After every
      `run` call every trigger is called with the current `train_step` value
      as an np scalar.
    checkpoint_interval: Number of train steps in between checkpoints. Note
      these are placed into triggers and so a check to generate a checkpoint
      only occurs after every `run` call. Set to -1 to disable (this is not
      recommended, because it means that if the pipeline gets preempted, all
      previous progress is lost). This only takes care of the checkpointing
      the training process.  Policies must be explicitly exported through
      triggers.
    summary_interval: Number of train steps in between summaries. Note these
      are placed into triggers and so a check to generate a checkpoint only
      occurs after every `run` call.
    use_kwargs_in_agent_train: If True the experience from the replay buffer
      is passed into the agent as kwargs. This requires samples from the RB to
      be of the form `dict(experience=experience, kwarg1=kwarg1, ...)`. This
      is useful if you have an agent with a custom argspec.
    strategy: (Optional) `tf.distribute.Strategy` to use during training.
  
  Raises:
    ValueError:mini batching is enabled, but shuffle_buffer_size isn't
      provided.
    ValueError: minibatch_size is passed in for RNN networks. RNNs require
      full sequences.
    ValueError:mini batching is enabled, but
      agent._compute_value_and_advantage_in_train is set to `True`.
    ValueError: agent.update_normalizers_in_train or is set to `True`. The
      learner already updates the normalizers, so no need to update again in
      the agent.
docstr_tok: ["Initializes","a","PPOLearner","instance",".","``","`","python","agent","=","ppo_agent.PPOAgent","(","...",",","compute_value_and_advantage_in_train=False",",","#","Skips","updating","normalizers","in","the","agent",",","as","it","'s","handled","in","the","learner",".","update_normalizers_in_train=False",")","#","train_replay_buffer","and","normalization_replay_buffer","point","to","two","Reverb","#","tables","that","are","synchronized",".","Sampling","is","done","in","a","FIFO","fashion",".","def","experience_dataset_fn","(",")",":","return","train_replay_buffer.as_dataset","(","sample_batch_size",",","sequence_preprocess_fn=agent.preprocess_sequence",")","def","normalization_dataset_fn","(",")",":","return","normalization_replay_buffer.as_dataset","(","sample_batch_size",",","sequence_preprocess_fn=agent.preprocess_sequence",")","learner","=","PPOLearner","(","...",",","agent",",","experience_dataset_fn",",","normalization_dataset_fn",")","learner.run","(",")","``","`","Args",":","root_dir",":","Main","directory","path","where","checkpoints",",","saved_models",",","and","summaries","will","be","written","to",".","train_step",":","a","scalar","tf.int64","`","tf.Variable","`","which","will","keep","track","of","the","number","of","train","steps",".","This","is","used","for","artifacts","created","like","summaries",",","or","outputs","in","the","root_dir",".","agent",":","`","ppo_agent.PPOAgent","`","instance","to","train","with",".","Note","that","update_normalizers_in_train","should","be","set","to","`","False","`",",","otherwise","a","ValueError","will","be","raised",".","We","do","not","update","normalizers","in","the","agent","again","because","we","already","update","it","in","the","learner",".","When","mini","batching","is","enabled",",","compute_value_and_advantage_in_train","should","be","set","to","False",",","and","preprocessing","should","be","done","as","part","of","the","data","pipeline","as","part","of","`","replay_buffer.as_dataset","`",".","experience_dataset_fn",":","a","function","that","will","create","an","instance","of","a","tf.data.Dataset","used","to","sample","experience","for","training",".","Each","element","in","the","dataset","is","a","(","Trajectory",",","SampleInfo",")","pair",".","Note","that","each","Trajectory","in","itself","might","represent","an","episode","(","variable","length",")","or","a","collection","of","transitions","depending","on","how","the","replay","buffer","and","hence","the","`","experience_dataset_fn","`","was","set","up",".","normalization_dataset_fn",":","a","function","that","will","create","an","instance","of","a","tf.data.Dataset","used","for","normalization",".","This","dataset","is","often","from","a","separate","reverb","table","that","is","synchronized","with","the","table","used","in","experience_dataset_fn",".","Each","element","in","the","dataset","is","a","(","Trajectory",",","SampleInfo",")","pair",".","Note","that","each","Trajectory","in","itself","might","represent","an","episode","(","variable","length",")","or","a","collection","of","transitions","depending","on","how","the","replay","buffer","and","hence","the","`","experience_dataset_fn","`","was","set","up",".","num_samples",":","The","number","of","samples","used","for","training","and","normalization",".","A","sample","here","means","one","reading","of","the","experience_dataset_fn","which","in","turn","might","have","more","than","1","Trajectories","(","For","ex",".","if","this","is","coming","from","a","`","reverb_replay_buffer","`","with","`","sample_batch_size","`","=","2",",","then","one","sample","means","2","trajectories",",","and","the","trainer","will","pull","in","num_samples","*","2","trajectories","for","one","iteration","of","training",")",".","If","fewer","than","this","amount","of","batches","exists","in","the","dataset",",","the","learner","will","wait","for","more","data","to","be","added",",","or","until","the","reverb","timeout","is","reached",".","num_epochs",":","The","number","of","iterations","to","go","through","the","same","sequences",".","minibatch_size",":","The","minibatch","size",".","If","set",",","the","input","data","set","will","be","flattened",",","and","the","dataset","used","for","training","is","shaped","`","[","minibatch_size",",","1",",","...","]","`",",","where","each","element","in","the","dataset","represents","a","single","step",".","If","None",",","full","sequences","will","be","fed","into","the","agent",".","Please","set","this","parameter","to","None","for","RNN","networks","which","requires","full","sequences",".","shuffle_buffer_size",":","The","buffer","size","for","shuffling","the","trajectories","before","splitting","them","into","mini","batches",".","Only","required","when","mini","batch","learning","is","enabled","(","minibatch_size","is","set",")",".","Otherwise","it","is","ignored",".","Commonly","set","to","a","number","1-3x","the","episode","length","of","your","environment",".","after_train_strategy_step_fn",":","(","Optional",")","callable","of","the","form","`","fn","(","sample",",","loss",")","`","which","can","be","used","for","example","to","update","priorities","in","a","replay","buffer","where","sample","is","pulled","from","the","`","experience_iterator","`","and","loss","is","a","`","LossInfo","`","named","tuple","returned","from","the","agent",".","This","is","called","after","every","train","step",".","It","runs","using","`","strategy.run","(","...",")","`",".","triggers",":","List","of","callables","of","the","form","`","trigger","(","train_step",")","`",".","After","every","`","run","`","call","every","trigger","is","called","with","the","current","`","train_step","`","value","as","an","np","scalar",".","checkpoint_interval",":","Number","of","train","steps","in","between","checkpoints",".","Note","these","are","placed","into","triggers","and","so","a","check","to","generate","a","checkpoint","only","occurs","after","every","`","run","`","call",".","Set","to","-1","to","disable","(","this","is","not","recommended",",","because","it","means","that","if","the","pipeline","gets","preempted",",","all","previous","progress","is","lost",")",".","This","only","takes","care","of","the","checkpointing","the","training","process",".","Policies","must","be","explicitly","exported","through","triggers",".","summary_interval",":","Number","of","train","steps","in","between","summaries",".","Note","these","are","placed","into","triggers","and","so","a","check","to","generate","a","checkpoint","only","occurs","after","every","`","run","`","call",".","use_kwargs_in_agent_train",":","If","True","the","experience","from","the","replay","buffer","is","passed","into","the","agent","as","kwargs",".","This","requires","samples","from","the","RB","to","be","of","the","form","`","dict","(","experience=experience",",","kwarg1=kwarg1",",","...",")","`",".","This","is","useful","if","you","have","an","agent","with","a","custom","argspec",".","strategy",":","(","Optional",")","`","tf.distribute.Strategy","`","to","use","during","training",".","Raises",":","ValueError",":","mini","batching","is","enabled",",","but","shuffle_buffer_size","isn't","provided",".","ValueError",":","minibatch_size","is","passed","in","for","RNN","networks",".","RNNs","require","full","sequences",".","ValueError",":","mini","batching","is","enabled",",","but","agent._compute_value_and_advantage_in_train","is","set","to","`","True","`",".","ValueError",":","agent.update_normalizers_in_train","or","is","set","to","`","True","`",".","The","learner","already","updates","the","normalizers",",","so","no","need","to","update","again","in","the","agent","."]
code: def __init__(
      self,
      root_dir: Text,
      train_step: tf.Variable,
      agent: ppo_agent.PPOAgent,
      experience_dataset_fn: Callable[..., tf.data.Dataset],
      normalization_dataset_fn: Callable[..., tf.data.Dataset],
      num_samples: int,
      num_epochs: int = 1,
      minibatch_size: Optional[int] = None,
      shuffle_buffer_size: Optional[int] = None,
      after_train_strategy_step_fn: Optional[
          Callable[[types.NestedTensor, tf_agent.LossInfo], None]
      ] = None,
      triggers: Optional[Callable[..., None]] = None,
      checkpoint_interval: int = 100000,
      summary_interval: int = 1000,
      use_kwargs_in_agent_train: bool = False,
      strategy: Optional[tf.distribute.Strategy] = None,
  ):
      """Initializes a PPOLearner instance.
  
      ```python
      agent = ppo_agent.PPOAgent(...,
        compute_value_and_advantage_in_train=False,
        # Skips updating normalizers in the agent, as it's handled in the learner.
        update_normalizers_in_train=False)
  
      # train_replay_buffer and normalization_replay_buffer point to two Reverb
      # tables that are synchronized. Sampling is done in a FIFO fashion.
      def experience_dataset_fn():
        return train_replay_buffer.as_dataset(sample_batch_size,
          sequence_preprocess_fn=agent.preprocess_sequence)
      def normalization_dataset_fn():
        return normalization_replay_buffer.as_dataset(sample_batch_size,
          sequence_preprocess_fn=agent.preprocess_sequence)
  
      learner = PPOLearner(..., agent, experience_dataset_fn,
        normalization_dataset_fn)
      learner.run()
      ```
  
      Args:
        root_dir: Main directory path where checkpoints, saved_models, and
          summaries will be written to.
        train_step: a scalar tf.int64 `tf.Variable` which will keep track of the
          number of train steps. This is used for artifacts created like
          summaries, or outputs in the root_dir.
        agent: `ppo_agent.PPOAgent` instance to train with. Note that
          update_normalizers_in_train should be set to `False`, otherwise a
          ValueError will be raised. We do not update normalizers in the agent
          again because we already update it in the learner. When mini batching is
          enabled, compute_value_and_advantage_in_train should be set to False,
          and preprocessing should be done as part of the data pipeline as part of
          `replay_buffer.as_dataset`.
        experience_dataset_fn: a function that will create an instance of a
          tf.data.Dataset used to sample experience for training. Each element in
          the dataset is a (Trajectory, SampleInfo) pair. Note that each
          Trajectory in itself might represent an episode (variable length) or a
          collection of transitions depending on how the replay buffer and hence
          the `experience_dataset_fn` was set up.
        normalization_dataset_fn: a function that will create an instance of a
          tf.data.Dataset used for normalization. This dataset is often from a
          separate reverb table that is synchronized with the table used in
          experience_dataset_fn. Each element in the dataset is a (Trajectory,
          SampleInfo) pair. Note that each Trajectory in itself might represent an
          episode (variable length) or a collection of transitions depending on
          how the replay buffer and hence the `experience_dataset_fn` was set up.
        num_samples: The number of samples used for training and normalization. A
          sample here means one reading of the experience_dataset_fn which in turn
          might  have more than 1 Trajectories (For ex. if this is coming from a
          `reverb_replay_buffer` with `sample_batch_size` = 2, then one sample
          means 2 trajectories, and the trainer will pull in num_samples * 2
          trajectories for one iteration of training). If fewer than this amount
          of batches exists in the dataset, the learner will wait for more data to
          be added, or until the reverb timeout is reached.
        num_epochs: The number of iterations to go through the same sequences.
        minibatch_size: The minibatch size. If set, the input data set will be
          flattened, and the dataset used for training is shaped `[minibatch_size,
          1, ...]`, where each element in the dataset represents a single step. If
          None, full sequences will be fed into the agent. Please set this
          parameter to None for RNN networks which requires full sequences.
        shuffle_buffer_size: The buffer size for shuffling the trajectories before
          splitting them into mini batches. Only required when mini batch learning
          is enabled (minibatch_size is set). Otherwise it is ignored. Commonly
          set to a number 1-3x the episode length of your environment.
        after_train_strategy_step_fn: (Optional) callable of the form `fn(sample,
          loss)` which can be used for example to update priorities in a replay
          buffer where sample is pulled from the `experience_iterator` and loss is
          a `LossInfo` named tuple returned from the agent. This is called after
          every train step. It runs using `strategy.run(...)`.
        triggers: List of callables of the form `trigger(train_step)`. After every
          `run` call every trigger is called with the current `train_step` value
          as an np scalar.
        checkpoint_interval: Number of train steps in between checkpoints. Note
          these are placed into triggers and so a check to generate a checkpoint
          only occurs after every `run` call. Set to -1 to disable (this is not
          recommended, because it means that if the pipeline gets preempted, all
          previous progress is lost). This only takes care of the checkpointing
          the training process.  Policies must be explicitly exported through
          triggers.
        summary_interval: Number of train steps in between summaries. Note these
          are placed into triggers and so a check to generate a checkpoint only
          occurs after every `run` call.
        use_kwargs_in_agent_train: If True the experience from the replay buffer
          is passed into the agent as kwargs. This requires samples from the RB to
          be of the form `dict(experience=experience, kwarg1=kwarg1, ...)`. This
          is useful if you have an agent with a custom argspec.
        strategy: (Optional) `tf.distribute.Strategy` to use during training.
  
      Raises:
        ValueError:mini batching is enabled, but shuffle_buffer_size isn't
          provided.
        ValueError: minibatch_size is passed in for RNN networks. RNNs require
          full sequences.
        ValueError:mini batching is enabled, but
          agent._compute_value_and_advantage_in_train is set to `True`.
        ValueError: agent.update_normalizers_in_train or is set to `True`. The
          learner already updates the normalizers, so no need to update again in
          the agent.
      """
      if minibatch_size and shuffle_buffer_size is None:
          raise ValueError(
              "shuffle_buffer_size must be provided if minibatch_size is not None."
          )
  
      if minibatch_size and (
          agent._actor_net.state_spec or agent._value_net.state_spec
      ):
          raise ValueError(
              "minibatch_size must be set to None for RNN networks."
          )
  
      if minibatch_size and agent._compute_value_and_advantage_in_train:
          raise ValueError(
              "agent.compute_value_and_advantage_in_train should be set to False "
              "when mini batching is used."
          )
  
      if agent.update_normalizers_in_train:
          raise ValueError(
              "agent.update_normalizers_in_train should be set to False when "
              "PPOLearner is used."
          )
  
      strategy = strategy or tf.distribute.get_strategy()
      self._agent = agent
      self._minibatch_size = minibatch_size
      self._shuffle_buffer_size = shuffle_buffer_size
      self._num_epochs = num_epochs
      self._experience_dataset_fn = experience_dataset_fn
      self._normalization_dataset_fn = normalization_dataset_fn
      self._num_samples = num_samples
  
      self._generic_learner = learner.Learner(
          root_dir,
          train_step,
          agent,
          experience_dataset_fn=None,
          after_train_strategy_step_fn=after_train_strategy_step_fn,
          triggers=triggers,
          checkpoint_interval=checkpoint_interval,
          summary_interval=summary_interval,
          use_kwargs_in_agent_train=use_kwargs_in_agent_train,
          strategy=strategy,
      )
  
      self.num_replicas = strategy.num_replicas_in_sync
      self._create_datasets(strategy)
      self.num_frames_for_training = tf.Variable(0, dtype=tf.int32)
  
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 1 to create a new PPOLearner object from inherited base classes: object with example: obj = PPOLearner(root_dir, train_step, agent, experience_dataset_fn, normalization_dataset_fn, num_samples, num_epochs, minibatch_size, shuffle_buffer_size, after_train_strategy_step_fn, triggers, checkpoint_interval, summary_interval, use_kwargs_in_agent_train, strategy)