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 trim_and_mask_batch(self, trajectory_batch, advantages):
      (batch_size, seq_len) = trajectory_batch.observation.shape[:2]
      adv_seq_len = advantages.shape[1]
      # The advantage sequence should be shorter by the margin. Margin is the
      # number of timesteps added to the trajectory slice, to make the advantage
      # estimation more accurate. adv_seq_len determines the length of the target
      # sequence, and is later used to trim the inputs and targets in the training
      # batch. Example for margin 2:
      # observations.shape == (4, 5, 6)
      # rewards.shape == values.shape == (4, 5)
      # advantages.shape == (4, 3)
      assert adv_seq_len <= seq_len
      assert advantages.shape == (batch_size, adv_seq_len)
      # Trim observations, actions and mask to match the target length.
      observations = trajectory_batch.observation[:, :adv_seq_len]
      actions = trajectory_batch.action[:, :adv_seq_len]
      mask = trajectory_batch.mask[:, :adv_seq_len]
      # Apply the control mask, so we only compute policy loss for controllable
      # timesteps.
      mask *= trajectory_batch.env_info.control_mask[:, :adv_seq_len]
      return (observations, actions, 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 4 for calling the PolicyTrainTask obj's underlying member method code with example usage: obj.trim_and_mask_batch(trajectory_batch, advantages) and returns: observations,  actions,  mask