import tensorflow as tf

    dones_vec = tf.reshape(dones, (batch_size, num_tasks))

    relabelled_obs = self._task_distribution.combine(states_tiled, tasks_tiled)
    action_distribution = self._actor(
        relabelled_obs, step_type=(), network_state=())[0]
    log_pi = common.log_probability(action_distribution, actions_tiled,
                                    action_spec)
    log_pi_vec = tf.reshape(log_pi, (batch_size, num_tasks))

    logits_vec = (
        rewards_vec - log_pi_vec + self._gamma * (1.0 - dones_vec) * q_vals_vec)
    if self._relabel_type == "random":
      logits_vec = tf.ones_like(logits_vec)  # Hack to make sampling random

    ## End new version
    if self._normalize_cols:
      logits_vec = logits_vec - tf.math.reduce_logsumexp(
          logits_vec, axis=0)[None]
    relabel_indices = tf.random.categorical(logits=logits_vec, num_samples=1)

    ### Metrics
    global_step = tf.compat.v1.train.get_or_create_global_step()
    orig_indices = tf.range(
        self._sample_batch_size, dtype=relabel_indices.dtype)
