import tensorflow as tf

    q_vals, _ = self._critic(critic_input, training=False)
    q_vals_vec = tf.reshape(q_vals, (batch_size, num_tasks))

    rewards, dones = self._task_distribution.evaluate(states_tiled,
                                                      actions_tiled,
                                                      tasks_tiled)
    dones = tf.cast(dones, tf.float32)
    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,
