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)
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)
    
trax/rl/policy_tasks.py