import tensorflow as tf

    sampled_actions_tiled = self._actor(
        next_relabelled_obs, step_type=(), network_state=())[0].sample()
    critic_input = (next_relabelled_obs, sampled_actions_tiled)
    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,
