import tensorflow as tf

              logits[None], num_samples=self._sample_batch_size))
      mask = tf.cast(mask, tf.float32)[:, None]
      relabelled_tasks = mask * orig_tasks + (1 - mask) * relabelled_tasks

    states_and_tasks = self._task_distribution.combine(states, relabelled_tasks)
    next_states_and_tasks = self._task_distribution.combine(
        next_states, relabelled_tasks)
    new_observation = tf.concat(
        [states_and_tasks[:, None], next_states_and_tasks[:, None]], axis=1)
    assert new_observation.shape == experience.observation.shape
    experience = experience.replace(observation=new_observation)
    return experience
