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):
      # 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],  # Targets: regress to returns.
             np_trajectory.action,              # Policy targets: actions.
             np_trajectory.mask)                 # Padding mask.
Use the RLTask self._task to create inputs to the value model.
trax/rl/actor_critic_joint.py