import tensorflow as tf

      items = relabelled_items
    super(GoalRelabellingReplayBuffer, self)._add_batch(items)

  @tf.function
  def _soft_relabel(self, experience):
    # experience.observation.shape = [B x T=2 x obs_dim+state_dim]
    states, orig_tasks = self._task_distribution.split(
        experience.observation[:, 0])
    if self._task_distribution.tasks is None:
      tasks = orig_tasks
    else:
      tasks = tf.constant(self._task_distribution.tasks, dtype=tf.float32)
    next_states, _ = self._task_distribution.split(experience.observation[:, 1])
    if self._candidate_task_type == "states":
      candidate_tasks = self._task_distribution.state_to_task(states)
    elif self._candidate_task_type == "next_states":
      candidate_tasks = self._task_distribution.state_to_task(next_states)
    else:
      assert self._candidate_task_type == "tasks"
      candidate_tasks = tasks

    actions = experience.action[:, 0]
