Question: The approach was evaluated on only one dataset. It is difficult to draw conclusions on the generality of the approach, right?
Evidence: In our experimental setup, an ideal context encoder should capture the generalizable information for task inference, namely the difference between reward/dynamics functions across a distribution of tasks. However, as discussed in Section 1, there are two major challenges that impede conventional COMRL algorithms from learning robust representations:  MDP ambiguity arises due to COMRL algorithms' sensitivity to fixed dataset distributions. Take Sparse-Point-Robot for example, as in Figure, for tasks with a goal on the semicircle, the state-action distribution exhibits specific pattern which may reflect task identity. Given D = {(s, a, s , r)} as input, the context encoder may learn a spurious correlation between state-action distributions and task identity, which causes performance degradation under distribution shifts (Table).
Sparse reward in meta-environments could exacerbate MDP ambiguity by making a considerable portion of transitions uninformative for task inference, such as the samples outside any goals in Figure. Attention mechanism, especially the batch-wise channel attention, helps the context encoder attend to the informative portion of the input transitions, and therefore significantly improve the robustness of the learned policies.
To demonstrate the robustness of FOCAL++ in presence of the two challenges above, we tested it against distribution shift by using datasets of various qualities: expert, medium, random and mixed which combines all three. Shown in Table, we observe that overall the performance drop due to distribution shift is significantly lower when attention and contrastive learning are applied.
Moreover, we are aware that even mixing of datasets generated by different behavior policies cannot fully eliminate the risk of MDP ambiguity since the state-action distributions for each task still do not completely overlap. To show that the attention modules introduced by FOCAL++ indeed works as intended by capturing the reward-task dependency, we create a new dataset on Sparse-Point-Robot by merging the state-action support across all tasks and relabelling the sparse reward according to the task-specific reward functions. In principle, this fully prevents information leakage from the state-action distributions, forcing the context encoder to learn to distinguish the reward functions between tasks while minimizing the contrastive loss. Shown in Figure, we experimented with 3 attention variants of FOCAL++ on the relabeled dataset, and found that batch-wise attention significantly improves the performance as intended. Additionally, we visualize the density distribution of batch-wise attention weights assigned to samples in Figure. We see a clear tendency for the batch-attention module to assign zero weight to samples with zero rewards (the absolutely sparse data points which lie outside all goal circles in Figure) and maximum weights to the non-zero-reward transitions, with binary classification AUC = 0.969, which is clear evidence of FOCAL++ learning the correct correlation for task inference by attending to the informative context.
We evaluate FOCAL++ on 6 continuous control meta-environments of robotic locomotion (Todorov et al., 2012) adopted from FOCAL. 4 (Sparse-Point-Robot, Sparse-Cheetah-Vel, Sparse-Cheetah-FwdBack, Sparse-Ant-Fwd-Back) and 2 (Point-Robot-Wind, Walker-2D-Params) environments require adaptation by reward and transition functions respectively. For inference, FOCAL++ aggregates context from a fixed test set to infer task embedding, and is subsequently evaluated online. Besides FOCAL, three other baselines are compared: an offline variant of the PEARL algorithm (Rakelly et al., 2019) (Batch PEARL), a context-based offline BCQ algorithm (Fujimoto et al., 2019) (Contextual BCQ) and a two-stage COMRL algorithm with reward/dynamics relabelling (Li et al., 2019) (MBML).****Figure 5: Result on the relabeled Sparse-Point-Robot dataset. (a) State distributions of the expert datasets for 20 distinct tasks, with goals uniformly distributed on a semicircle. (b) On mixed dataset, FOCAL completely fails in this scenario whereas FOCAL++ variants with batch-wise attention are able to learn. (c) Probability distribution of the batch-wise attention weight of samples with absolutely zero and non-zero reward. Binary classification AUC = 0.969.
Answer: We thoroughly tested FOCAL++ on 6 MuJoCo datasets with sparse rewards, all of which are common meta-RL benchmarks and aligned with well-known papers like PEARL [1] and MQL [2], with one additional relabeled dataset for discussion of MDP ambiguity (Figure 5). As some reviewers suggested, we also added more experiments with other COMRL baselines on non-sparse environments. Please check our General Response 2. Overall, we are really grateful for your comments and advice, which are all valuable in helping us improve the work. Hope our explanation and updated material address your concerns. Please let us know if you have any further comments or concerns or ways in which we can further improve our paper.Our proposed method is firmly grounded with mathematical motivation and proofs.