def policy_batches_stream(self):
    """Use the RLTask self._task to create inputs to the policy model."""
    # Maximum slice length for policy is max_slice_len + the added policy len.
    max_slice_length = self._max_slice_length + self._added_policy_slice_length
    for np_trajectory in self._task.trajectory_batch_stream(
        self._policy_batch_size,
        epochs=self._replay_epochs,
        max_slice_length=max_slice_length,
        margin=self._added_policy_slice_length,
    ):
      dist_inputs = self._get_dist_inputs(np_trajectory)
      (values, _, act_log_probs) = self._run_value_model(
          np_trajectory.observation, dist_inputs)
      values = self._aggregate_values(values, 'mean', act_log_probs)
      if len(values.shape) != 2:
        raise ValueError('Values are expected to have shape ' +
                         '[batch_size, length], got: %s' % str(values.shape))
      if values.shape[0] != self._policy_batch_size:
        raise ValueError('Values first dimension should = policy batch size, ' +
                         '%d != %d' %(values.shape[0], self._policy_batch_size))
      yield self.policy_inputs(np_trajectory, values)
Use the RLTask self._task to create inputs to the policy model.
trax/rl/actor_critic.py