import tensorflow as tf

    rewards_vec = tf.reshape(rewards, (batch_size, num_tasks))
    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
