import tensorflow as tf

          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(
              logits[None], num_samples=self._sample_batch_size))
      mask = tf.cast(mask, tf.float32)[:, None]
      relabelled_tasks = mask * orig_tasks + (1 - mask) * relabelled_tasks

