import tensorflow as tf

          num_steps=2,
      ).prefetch(3)
      self._iterator = iter(dataset)
    experience, unused_info = next(self._iterator)
    if self._relabel_type in ["soft", "random"]:
      experience = self._soft_relabel(experience)
    elif self._relabel_type in ["last", "future"]:
      # Reassign the next_states to have the same goal as the current states
      _, tasks = self._task_distribution.split(experience.observation[:, 0])
      next_states, _ = self._task_distribution.split(experience.observation[:,
                                                                            1])
      next_states_and_tasks = self._task_distribution.combine(
          next_states, tasks)
      new_observation = tf.concat(
          [
              experience.observation[:, 0][:, None], next_states_and_tasks[:,
                                                                           None]
          ],
          axis=1,
      )
      assert new_observation.shape == experience.observation.shape
      experience = experience.replace(observation=new_observation)
    if self._relabel_type is not None:
      # Recompute rewards and done flags
      states, tasks = self._task_distribution.split(experience.observation[:,
                                                                           0])
