def policy_batches_stream(self):
    """Use the RLTask self._task to create inputs to the policy model."""
    # For now TD-0 estimation of the value. TODO(pkozakowski): Support others?
    for np_trajectory in self._task.trajectory_batch_stream(
        self._policy_batch_size,
        epochs=self._replay_epochs,
        max_slice_length=self._max_slice_length,
    ):
      dist_inputs = self._get_dist_inputs(np_trajectory)
      (q_values, actions, act_log_probs) = self._run_value_model(
          np_trajectory.observation, dist_inputs)
      shapes.assert_same_shape(q_values, act_log_probs)

      # q_values shape: (batch_size, n_samples, length)
      if len(q_values.shape) != 3:
        raise ValueError('Q-values are expected to have shape [batch_size, ' +
                         'n_samples, length], got: %s' % str(q_values.shape))
      if q_values.shape[1] != self._q_value_n_samples:
        raise ValueError('Q-values dimension 1 should = n_samples, %d != %d'
                         % (q_values.shape[1], self._q_value_n_samples))
      if q_values.shape[0] != self._policy_batch_size:
        raise ValueError('Q-values dimension 0 should = policy batch size, ' +
                         '%d!=%d' %(q_values.shape[1], self._policy_batch_size))

      mask = np_trajectory.mask
      mask = np.reshape(mask, [mask.shape[0], 1] + list(mask.shape[1:]))
      mask = jnp.broadcast_to(mask, q_values.shape)
      shapes.assert_same_shape(mask, q_values)
      yield (np_trajectory.observation, actions, q_values, act_log_probs, mask)
Use the RLTask self._task to create inputs to the policy model.
trax/rl/actor_critic.py