def batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    for np_trajectory in self._task.trajectory_batch_stream(
        self._batch_size, max_slice_length=self._max_slice_length, epochs=[-1]):
      if np_trajectory.dist_inputs is not None:
        old_dist_inputs = np_trajectory.dist_inputs
      else:
        old_dist_inputs = jnp.zeros(
            np_trajectory.reward.shape + (self._policy_dist.n_inputs,)
        )
      old_log_probs = self._policy_dist.log_prob(
          old_dist_inputs, np_trajectory.action
      )
      # Insert an extra depth dimension, so the target shape is consistent with
      # the network output shape.
      yield (np_trajectory.observation,         # Inputs to the value model.
             np_trajectory.return_[:, :, None],
             np_trajectory.done[:, :, None],
             np_trajectory.reward[:, :, None],
             np_trajectory.action,
             old_log_probs,
             np_trajectory.mask)
Use the RLTask self._task to create inputs to the value model.
trax/rl/actor_critic_joint.py