import tensorflow as tf

      assert self._candidate_task_type == "tasks"
      candidate_tasks = tasks

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