class_docstr: Task for policy training.
  
  Trains the policy based on action advantages.
class_docstr_tok: ["Task","for","policy","training",".","Trains","the","policy","based","on","action","advantages","."]
docstr: Initializes PolicyTrainTask.
  
  Args:
    trajectory_batch_stream: Generator of trax.rl.task.TimeStepBatch.
    optimizer: Optimizer for network training.
    lr_schedule: Learning rate schedule for network training.
    policy_distribution: Distribution over actions.
    advantage_estimator: Function
      (rewards, returns, values, dones) -> advantages, created by one of the
      functions from trax.rl.advantages.
    value_fn: Function TimeStepBatch -> array (batch_size, seq_len)
      calculating the baseline for advantage calculation. Can be used to
      implement actor-critic algorithms, by substituting a call to the value
      network as value_fn.
    weight_fn: Function float -> float to apply to advantages. Examples:
      - A2C: weight_fn = id
      - AWR: weight_fn = exp
      - behavioral cloning: weight_fn(_) = 1
    advantage_normalization: Whether to normalize advantages.
    advantage_normalization_epsilon: Epsilon to use then normalizing
      advantages.
    head_selector: Layer to apply to the network output to select the value
      head. Only needed in multitask training. By default, use a no-op layer,
      signified by an empty sequence of layers, ().
docstr_tok: ["Initializes","PolicyTrainTask",".","Args",":","trajectory_batch_stream",":","Generator","of","trax.rl.task.TimeStepBatch",".","optimizer",":","Optimizer","for","network","training",".","lr_schedule",":","Learning","rate","schedule","for","network","training",".","policy_distribution",":","Distribution","over","actions",".","advantage_estimator",":","Function","(","rewards",",","returns",",","values",",","dones",")","-",">","advantages",",","created","by","one","of","the","functions","from","trax.rl.advantages",".","value_fn",":","Function","TimeStepBatch","-",">","array","(","batch_size",",","seq_len",")","calculating","the","baseline","for","advantage","calculation",".","Can","be","used","to","implement","actor-critic","algorithms",",","by","substituting","a","call","to","the","value","network","as","value_fn",".","weight_fn",":","Function","float","-",">","float","to","apply","to","advantages",".","Examples",":","-","A2C",":","weight_fn","=","id","-","AWR",":","weight_fn","=","exp","-","behavioral","cloning",":","weight_fn","(","_",")","=","1","advantage_normalization",":","Whether","to","normalize","advantages",".","advantage_normalization_epsilon",":","Epsilon","to","use","then","normalizing","advantages",".","head_selector",":","Layer","to","apply","to","the","network","output","to","select","the","value","head",".","Only","needed","in","multitask","training",".","By","default",",","use","a","no-op","layer",",","signified","by","an","empty","sequence","of","layers",",","(",")","."]
code: def __init__(
      self,
      trajectory_batch_stream,
      optimizer,
      lr_schedule,
      policy_distribution,
      advantage_estimator,
      value_fn,
      weight_fn=(lambda x: x),
      advantage_normalization=True,
      advantage_normalization_epsilon=1e-5,
      head_selector=(),
  ):
      """Initializes PolicyTrainTask.
  
      Args:
        trajectory_batch_stream: Generator of trax.rl.task.TimeStepBatch.
        optimizer: Optimizer for network training.
        lr_schedule: Learning rate schedule for network training.
        policy_distribution: Distribution over actions.
        advantage_estimator: Function
          (rewards, returns, values, dones) -> advantages, created by one of the
          functions from trax.rl.advantages.
        value_fn: Function TimeStepBatch -> array (batch_size, seq_len)
          calculating the baseline for advantage calculation. Can be used to
          implement actor-critic algorithms, by substituting a call to the value
          network as value_fn.
        weight_fn: Function float -> float to apply to advantages. Examples:
          - A2C: weight_fn = id
          - AWR: weight_fn = exp
          - behavioral cloning: weight_fn(_) = 1
        advantage_normalization: Whether to normalize advantages.
        advantage_normalization_epsilon: Epsilon to use then normalizing
          advantages.
        head_selector: Layer to apply to the network output to select the value
          head. Only needed in multitask training. By default, use a no-op layer,
          signified by an empty sequence of layers, ().
      """
      self.trajectory_batch_stream = trajectory_batch_stream
      self._value_fn = value_fn
      self._advantage_estimator = advantage_estimator
      self._weight_fn = weight_fn
      self._advantage_normalization = advantage_normalization
      self._advantage_normalization_epsilon = (
          advantage_normalization_epsilon
      )
      self.policy_distribution = policy_distribution
  
      labeled_data = map(self.policy_batch, trajectory_batch_stream)
      sample_batch = self.policy_batch(
          next(trajectory_batch_stream), shape_only=True
      )
      loss_layer = distributions.LogLoss(
          distribution=policy_distribution
      )
      loss_layer = tl.Serial(head_selector, loss_layer)
      super().__init__(
          labeled_data,
          loss_layer,
          optimizer,
          sample_batch=sample_batch,
          lr_schedule=lr_schedule,
          loss_name="policy_loss",
      )
  
file_path: trax/trax/rl/policy_tasks.py
imports: import numpy
  from trax import layers
  from trax.fastmath import numpy
  from trax.rl import distributions
  from trax.supervised import training
label_desc: Use image node_id 1 to create a new PolicyTrainTask object from inherited base classes: training with example: obj = PolicyTrainTask(trajectory_batch_stream, optimizer, lr_schedule, policy_distribution, advantage_estimator, value_fn, weight_fn, advantage_normalization, advantage_normalization_epsilon, head_selector)