def value_batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    max_slice_length = self._max_slice_length
    min_slice_length = 1
    for np_trajectory in self._task.trajectory_batch_stream(
        self._value_batch_size,
        max_slice_length=max_slice_length,
        min_slice_length=min_slice_length,
        margin=0,
        epochs=self._replay_epochs,
    ):
      values_target = self._run_value_model(
          np_trajectory.observation, use_eval_model=True)
      if self._double_dqn:
        values = self._run_value_model(
            np_trajectory.observation, use_eval_model=False
        )
        index_max = np.argmax(values, axis=-1)
        ind_0, ind_1 = np.indices(index_max.shape)
        values_max = values_target[ind_0, ind_1, index_max]
      else:
        values_max = np.array(jnp.max(values_target, axis=-1))

      # The advantage_estimator returns
      #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards
      #        - values_max(s_i)
      # hence we have to add values_max(s_i) in order to get n-step returns:
      #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards
      # Notice, that in DQN the tensor values_max[:, :-self._margin]
      # is the same as values_max[:, :-1].
      n_step_returns = values_max[:, :-self._margin] + \
          self._advantage_estimator(
              rewards=np_trajectory.reward,
              returns=np_trajectory.return_,
              values=values_max,
              dones=np_trajectory.done,
              discount_mask=np_trajectory.env_info.discount_mask,
          )

      length = n_step_returns.shape[1]
      target_returns = n_step_returns[:, :length]
      inputs = np_trajectory.observation[:, :length, :]

      yield (
          # Inputs are observations
          # (batch, length, obs)
          inputs,
          # the max indices will be needed to compute the loss
          np_trajectory.action[:, :length],  # index_max,
          # Targets: computed returns.
          # target_returns, we expect here shapes such as
          # (batch, length, num_actions)
          target_returns / self._value_network_scale,
          # TODO(henrykm): mask has the shape (batch, max_slice_length)
          # that is it misses the action dimension; the preferred format
          # would be np_trajectory.mask[:, :length, :] but for now we pass:
          np.ones(shape=target_returns.shape),
      )
Use the RLTask self._task to create inputs to the value model.
trax/rl/training.py