def policy_inputs(self, trajectory, values):
    """Create inputs to policy model from a TimeStepBatch and values."""
    # How much TD to use is determined by the added policy slice length,
    # as the policy batches need to be this much longer to calculate TD.
    advantages = self._advantage_estimator(
        rewards=trajectory.reward,
        returns=trajectory.return_,
        values=values,
        dones=trajectory.done,
        discount_mask=trajectory.env_info.discount_mask,
    )
    # Observations should be the same length as advantages - so if we are
    # using n_extra_steps, we need to trim the length to match.
    obs = trajectory.observation[:, :advantages.shape[1]]
    act = trajectory.action[:, :advantages.shape[1]]
    mask = trajectory.mask[:, :advantages.shape[1]]  # Mask to zero-out padding.
    if trajectory.dist_inputs is not None:
      dist_inputs = self._get_dist_inputs(trajectory)
      dist_inputs = dist_inputs[:, :advantages.shape[1]]
    else:
      dist_inputs = jnp.zeros(advantages.shape + (self._policy_dist.n_inputs,))
    # Shape checks to help debugging.
    if len(advantages.shape) != 2:
      raise ValueError('Advantages are expected to have shape ' +
                       '[batch_size, length], got: %s' % str(advantages.shape))
    if act.shape[0:2] != advantages.shape:
      raise ValueError('First 2 dimensions of actions should be the same as in '
                       'advantages, %s != %s' % (act.shape[0:2],
                                                 advantages.shape))
    if obs.shape[0:2] != advantages.shape:
      raise ValueError('First 2 dimensions of observations should be the same '
                       'as in advantages, %s != %s' % (obs.shape[0:2],
                                                       advantages.shape))
    if dist_inputs.shape[:2] != advantages.shape:
      raise ValueError('First 2 dimensions of dist_inputs should be the same '
                       'as in advantages, %s != %s' % (dist_inputs.shape[:2],
                                                       advantages.shape))
    if mask.shape != advantages.shape:
      raise ValueError('Mask and advantages shapes should be the same'
                       ', %s != %s' % (mask.shape, advantages.shape))
    return (obs, act, advantages, dist_inputs, mask)
Create inputs to policy model from a TimeStepBatch and values.
trax/rl/actor_critic.py