import tensorflow as tf

    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

