import tensorflow as tf

            step=global_step,
        )

      # What are the average Q values of the relabelled tasks?
      indices = tf.transpose(
          tf.stack([orig_indices, tf.squeeze(relabel_indices)], axis=0))
      relabel_q_vals = tf.gather_nd(logits_vec, indices)
      tf.compat.v2.summary.scalar(
          name="relabel_q_vals",
          data=tf.reduce_mean(relabel_q_vals),
          step=global_step,
      )

      max_q = tf.reduce_max(logits_vec, axis=1)
      tf.compat.v2.summary.scalar(
          name="max_q", data=tf.reduce_mean(max_q), step=global_step)

    ### End metrics

    # For both state-centric and goal-centric relabelling, the implementation of
    # mixing is the same: we randomly replace some of the indices with the
    # diagonal.
    relabelled_tasks = tf.gather(candidate_tasks, tf.squeeze(relabel_indices))

    if self._relabel_prob == 0:
      relabelled_tasks = orig_tasks
    elif 0 < self._relabel_prob < 1:
      logits = tf.log([1.0 - self._relabel_prob, self._relabel_prob])
      mask = tf.squeeze(
          tf.random.categorical(
