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","."]
code: def calculate_advantages(self, trajectory_batch, shape_only=False):
      (batch_size, seq_len) = trajectory_batch.observation.shape[:2]
      assert trajectory_batch.action.shape[:2] == (batch_size, seq_len)
      assert trajectory_batch.mask.shape == (batch_size, seq_len)
      if shape_only:
          values = np.zeros((batch_size, seq_len))
      else:
          # Compute the value, i.e. baseline in advantage computation.
          values = np.array(self._value_fn(trajectory_batch))
          assert values.shape == (batch_size, seq_len)
      # Compute the advantages using the chosen advantage estimator.
      return self._advantage_estimator(
          rewards=trajectory_batch.reward,
          returns=trajectory_batch.return_,
          dones=trajectory_batch.done,
          values=values,
          discount_mask=trajectory_batch.env_info.discount_mask,
      )
  
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 2 for calling the PolicyTrainTask obj's underlying member method code with example usage: obj.calculate_advantages(trajectory_batch, shape_only) and returns: self