import tensorflow as tf

    actions = experience.action[:, 0]
    num_tasks = tasks.shape[0]
    batch_size = states.shape[0]
    task_dim = tasks.shape[1]
    obs_dim = states.shape[1]
    action_dim = actions.shape[1]
    action_spec = self._actor.output_tensor_spec

    states_tiled = tf.tile(states[:, None], [1, num_tasks, 1])  # B x B x D
    states_tiled = tf.reshape(states_tiled,
                              [batch_size * num_tasks, obs_dim])  # B*B x D
    actions_tiled = tf.tile(actions[:, None], [1, num_tasks, 1])  # B x B x D
    actions_tiled = tf.reshape(actions_tiled,
                               [batch_size * num_tasks, action_dim])  # B*B x D
    tasks_tiled = tf.tile(tasks[None], [batch_size, 1, 1])  # B x B x D
    tasks_tiled = tf.reshape(tasks_tiled,
                             [batch_size * num_tasks, task_dim])  # B*B x D

    next_states_tiled = tf.tile(next_states[:, None], [1, num_tasks, 1])
    next_states_tiled = tf.reshape(next_states_tiled,
                                   [batch_size * num_tasks, obs_dim])  # B*B x D
    next_relabelled_obs = self._task_distribution.combine(
        next_states_tiled, tasks_tiled)

    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))

