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: Computes a policy training batch based on a trajectory batch.
  
  Args:
    trajectory_batch: trax.rl.task.TimeStepBatch with a batch of trajectory
      slices. Elements should have shape (batch_size, seq_len, ...).
    shape_only: Whether to return dummy zero arrays of correct shape. Useful
      for initializing models.
  
  Returns:
    Triple (observations, actions, weights), where weights are the
    advantage-based weights for the policy loss. Shapes:
    - observations: (batch_size, seq_len) + observation_shape
    - actions: (batch_size, seq_len) + action_shape
    - weights: (batch_size, seq_len)
docstr_tok: ["Computes","a","policy","training","batch","based","on","a","trajectory","batch",".","Args",":","trajectory_batch",":","trax.rl.task.TimeStepBatch","with","a","batch","of","trajectory","slices",".","Elements","should","have","shape","(","batch_size",",","seq_len",",","...",")",".","shape_only",":","Whether","to","return","dummy","zero","arrays","of","correct","shape",".","Useful","for","initializing","models",".","Returns",":","Triple","(","observations",",","actions",",","weights",")",",","where","weights","are","the","advantage-based","weights","for","the","policy","loss",".","Shapes",":","-","observations",":","(","batch_size",",","seq_len",")","+","observation_shape","-","actions",":","(","batch_size",",","seq_len",")","+","action_shape","-","weights",":","(","batch_size",",","seq_len",")"]
code: def policy_batch(self, trajectory_batch, shape_only=False):
      """Computes a policy training batch based on a trajectory batch.
  
      Args:
        trajectory_batch: trax.rl.task.TimeStepBatch with a batch of trajectory
          slices. Elements should have shape (batch_size, seq_len, ...).
        shape_only: Whether to return dummy zero arrays of correct shape. Useful
          for initializing models.
  
      Returns:
        Triple (observations, actions, weights), where weights are the
        advantage-based weights for the policy loss. Shapes:
        - observations: (batch_size, seq_len) + observation_shape
        - actions: (batch_size, seq_len) + action_shape
        - weights: (batch_size, seq_len)
      """
      advantages = self.calculate_advantages(
          trajectory_batch, shape_only=shape_only
      )
      (observations, actions, mask) = self.trim_and_mask_batch(
          trajectory_batch, advantages
      )
      weights = (
          self.calculate_weights(advantages) * mask / jnp.sum(mask)
      )
      return (observations, actions, weights)
  
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 5 for calling the PolicyTrainTask obj's underlying member method code with example usage: obj.policy_batch(trajectory_batch, shape_only) and returns: observations,  actions,  weights