def value_batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    max_slice_length = self._max_slice_length + self._added_policy_slice_length
    for np_trajectory in self._task.trajectory_batch_stream(
        self._value_batch_size,
        max_slice_length=max_slice_length,
        min_slice_length=(1 + self._added_policy_slice_length),
        margin=self._added_policy_slice_length,
        epochs=self._replay_epochs,
    ):
      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, self._q_value_aggregate, act_log_probs)

      # TODO(pkozakowski): Add some shape assertions and docs.
      # Calculate targets based on the advantages over the target network - this
      # allows TD learning for value networks.
      advantages = self._advantage_estimator(
          rewards=np_trajectory.reward,
          returns=np_trajectory.return_,
          values=values,
          dones=np_trajectory.done,
          discount_mask=np_trajectory.env_info.discount_mask,
      )
      length = advantages.shape[1]
      values = values[:, :length]
      target_returns = values + advantages

      inputs = (np_trajectory.observation[:, :length],)
      if self._q_value:
        inputs += (np_trajectory.action[:, :length],)

      # Insert an extra depth dimension, so the target shape is consistent with
      # the network output shape.
      yield (
          # Inputs: observations and maybe actions.
          *inputs,
          # Targets: computed returns.
          target_returns[:, :, None] / self._value_network_scale,
          # Mask to zero-out padding.
          np_trajectory.mask[:, :length, None],
      )
Use the RLTask self._task to create inputs to the value model.
trax/rl/actor_critic.py