Implementation Details. All MLPs defined in Equation REF  have 2 layers followed by a ReLU activation, and the embedding layer outputs features of size 128 in all experiments. For training, we use ADAM [1]} optimizer with a learning rate of \(10^{-5}\) . We use Soft Actor-Critic (SAC) [2]} as backbone RL algorithm for all methods. For experiments on X-MAGICAL, we follow [3]} and learn a state-based policy; RL training is performed for 500k steps for all embodiments. For robotic manipulation experiments, we learn a multi-view image-based SAC policy [4]}. We train RL agent for 300k, 800k and 700k steps for Reach, Push and Peg in Box respectively. For fair comparison, we only change the learned reward function across methods and keep the RL setup identical. The success rates presented for all our experiments are averaged over 50 episodes. Refer to Appendix  for further implementation details.
