Dear Reviewer,

Thanks for finding our paper novel, empirical rigor and practical relevance. We hope the following statement clears your concern.

WQ1: Emphasize the key points of Definition 1.

A for WQ1:

• In the revised version, we have updated Definition 1 to emphasize predefined task distinction thresholds $\delta$ and $d$.

• We have refined the requirement on the policy from "a specific policy $\pi$" to "any $\epsilon$-optimal policy $\pi$ that $V_{\pi}(s)\geq V^*(s)-\epsilon, \forall s \in \mathcal{S}$".

---

WQ2: A few confounders that are not accounted for.

A for WQ2:

• The sizes of these models in the toy experiment are as follows: CQL-double 56,075, CQL 23,947, DAIL 33,697. DAIL's size is between CQL and CQL-double. When in BabyAI, the numbers of parameters are: CQL 1,481,217, DAIL 1,488,179, where the impact of this parameter scale becomes even less significant.

• As suggested, we run a CQL-half and a CQL-quarter on the toy experiment, and here are the results. Due to space constraints, we have presented only the results for 32, 128, and 512 tasks for reference (in percentage).

Table 1: Average success rates over 100 evaluations for each number of instructions.

Results above indicate that even when the model's parameter scale is halved or reduced to one-quarter, CQL's performance remains nearly identical, which demonstrates that the model size is not the key confounder of the performance differences between CQL and DAIL.

---

WQ3 (WC7-11)

A for WQ3:

We have reconstructed part of the proof and updated the theoretical results accordingly. The new sample efficiency result is as follows:

$n_{\text{value}} \geq \frac{C_{\text{value}} \log(3M^2 / \eta)}{\delta^2}, n_{\text{dist}} \geq \frac{C_{\text{dist}} \log(3M^2 / \eta)}{d^2}$

To ensure clarity and coherence, we first provide a brief summary of the revised proof, followed by our responses to each of the comments.

Proofs:(modifying from line 552)

Following the triangle inequality of the absolute value, we have: (WC9)

$$\left|\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]-\mathbb{E}\_{\pi}[|Q\_i-Q\_j|]\right| \leq \left|\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]-\mathbb{E}\_\mathcal{D}\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]\right|+ \left|\mathbb{E}\_\mathcal{D}\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]-\mathbb{E}\_{\pi}[|Q\_i-Q\_j|]\right|$$

In general, we assume that the offline RL dataset is collected by an $\epsilon$-optimal policy, i.e., there exists an $\epsilon$-optimal policy $\pi$ that $\mathcal{D}$ satisfies

$$\mathbb{E}\_\mathcal{D}\,\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|] = \mathbb{E}\_{\pi}[|\hat{Q}\_i-\hat{Q}\_j|], \quad \mathbb{E}\_\mathcal{D}(\hat{Q}) = Q\_{\pi} = Q$$

Based on the above derivation steps (with some intermediate steps omitted for word limit), we can upper-bound the original expression by following the three expressions, $\left|\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]-\mathbb{E}\_\mathcal{D}\mathbb{E}\_{(s,a)\sim\mathcal{D}}[|\hat{Q}\_i-\hat{Q}\_j|]\right|$ (term 1), $\mathbb{E}\_{\pi}|\hat{Q}\_i-Q\_i|$ (term 2), $\mathbb{E}\_{\pi}|\hat{Q}\_j-Q\_j|$ (term 3).

Similar to the derivation of probabilities in Line 569 of the original text, we can also convert the original probability into the sum of the probabilities that each of the three terms exceeds $d/2$. Moreover, each of these three terms takes the form of an empirical estimate approximating the true expectation. Based on the existing conditions, we can apply Theorem 1 (Hoeffding's inequality) in[1], and obtain that for $\forall l_i,l_j\in \mathcal{L}$, (WC10, 11)

$$\mathrm{Pr}\left(\text{term 1, 2, 3}\geq \delta/2\right)\leq 2\exp\left(-\tfrac{n\delta^2}{C\_{\text{value}}}\right) ...$$

The subsequent proof follows the same reasoning as in the original manuscript. Besides, we modified the proof of distributional RL accordingly, where we upper-bound the original expression by similar three expressions.

Answers:

• A for WC7: Z is a shorthand of $Z^\pi$ and denotes a random variable. However, we often do not have access to the random variable Z but only its empirical version $Z_{\theta}$, which is a parameterized model to be learned.

• A for WC8: Yes, $\hat{\theta}$ denotes the parameters of a target network, a common stabilization technique in deep RL. This network is initialized identically to the learning model but updated periodically (e.g., via polyak averaging $\hat{\theta}\leftarrow \tau \hat{\theta}+(1-\tau)\theta$). Its primary role is to provide stable regression targets for the Bellman equation.

• A for WC9: We have removed the vague assumption in Line 552 and given a more rigorous derivation.

• A for WC10: In the revised version, we have explicitly indicated the theorem from the referenced work. In addition, we use Corollary 5.2 and Remark 5 in [2] for proofs of distributional RL.

• A for WC11: We have modified the proof to eliminate the dependence on arguments related to "sub-Gaussian" and "1-Lipschitz" properties.

---

WQ4: Quantitative measures of sample-efficiency.

A for WQ4:

We have computed the AUC for the learning curves, and present the AUC comparisons of ID learning curves at the key milestones of 10 and 20 epochs below. The complete results are comprehensively detailed in the revised manuscript.

Table 2: AUC of in-distribution BabyAI tasks.

Statistical-significance test. As suggested, we perform statistical significance tests on these AUC results at 20 epochs. The result shows that, at the $\alpha=0.05$ significance level, paired t-tests demonstrate statistically significant performance improvements for both the distributional component (p=0.0342) and the semantic alignment module (p=0.0332) individually. Furthermore, DAIL yields highly significant enhancements (p=0.002), confirming the synergistic efficacy of the complete methodology.

These results of both epochs confirm the significant sample efficiency gains of our method.

---

SQ5: Quantitative measure of clustering

A for SQ5:

As suggested, we measure the clustering quality with the Silhouette score[3]. The labels required to calculate the score are defined as follows: two instructions share the same label if they require the agent to perform the same action on the same kind of object with the same color as the final target object. The distance between language embeddings is measured using cosine distance.

Table 3: Silhouette score of in-distribution BabyAI tasks.

The results demonstrate that in the All-task scenario, the clustering metric of CQL even exhibits negative values, indicating a severe task confusion phenomenon where it fails to effectively partition tasks, while both proposed methods evidently improve clustering performance.

---

WC1-WC6 and WC12: Clarity issues.

A for WC1-WC6 and WC12:

We sincerely thank you for identifying the areas needing clarification or reference in our manuscript. These points have been carefully addressed in the revised version.

WC13, SS2: Test-train split in BabyAI

A for WC13, SS2:

Our experimental design indeed aims to evaluate compositional generalization. We use only the full given instruction (verbatim text) as the unit for splitting. We sampled approximately 60% of the unique instructions, and the training set contains only the trajectories in which the agent attempted to complete these tasks. We only constrain whether an instruction's text is included in the set. In Appendix E.2.1, we describe more details of ID/OOD split.

WC14: Test-train split in ALFRED

A for WC14:

For the ALFRED dataset, we strictly adhere to the original work's train-test split between the Training set and Evaluation set (see Appendix E.2.2 for details). While this setup inherently tests generalization capabilities, it differs from BabyAI’s ID/OOD split: the original ALFRED split allows partial instruction overlap between training and evaluation sets (despite environmental variations).

WC15: Clarify ``w.p.'', $z_i$, and $p_i$ in Equation 11.

A for WC15:

Equation 11 models the discrete value distribution $Z_\theta$ with two key components:

• $z_i$ (termed atoms in our manuscript), which are support points on the quantized value domain.
• $p_i$, representing the probability mass that the state-action value matches the support point $z_i$.

"w.p." here stands for "with probability". Please refer to Appendix C for further explanations. We have revised this paragraph to make it clearer.

---

WO1-WO3, Misc:

A for WO1-WO3:

We sincerely appreciate your detailed suggestions on related works, word choice, paragraph organization, and the logical presentation. They have been comprehensively addressed in the revised manuscript.

Thanks again for your detailed and valuable comments. We sincerely hope our response has cleared your concerns, and we are looking forward to more discussions.

---

References:

[1] Concentration inequalities. 2003.

[2] Convergence and concentration of empirical measures under Wasserstein distance in unbounded functional spaces. 2020.

[3] Silhouettes: A Graphical Aid to the Interpretation and Validation of Cluster Analysis, 1987.

Dear Reviewer,

Thanks for finding our work fresh and accessible. We hope the following statement clears your concern.

---

W1: Inadequately justifying the practical significance

A for W1: Language-conditioned RL holds significant practical importance. There is a growing expectation that a single agent should be able to accomplish various tasks specified by human instructions[1].

We present that task ambiguity is a realistic problem in language-conditioned RL. It is mainly obstructed by the complexity and diversity of language instructions.

To make our methods practical in the real world, we choose ALFRED as one of benchmarks, which includes sequential decision-making tasks involving household activities through language instructions and first-person vision. ALFRED is widely used for testing visual navigation and embodied agents[2], our experiments on ALFRED show high real-world relevance.

---

W2: Offline RL baselines and multi-task/meta-RL methods

A for W2:

As suggested, we conduct extra experiments on state-of-the-art offline RL (MCQ[3], TD3+BC[4]) and meta RL methods (CCM[5], FOCAL[6], CSRO[7], UNICORN[8]). The experimental results of training from the high-quality dataset of BabyAI are as follows:

Table 1: Performance of out-of-distribution BabyAI tasks.

• Offline RL: The focus of offline RL methods is to avoid the phenomenon where policies perform poorly in unknown state-action spaces, which differs significantly from the task ambiguity problem addressed in this paper. Consequently, the performance gains prove ineffective or even detrimental.

• Meta-RL: They focus on encoding trajectories to infer implicit task information. In our setting, however, task information is explicitly provided through instructions, and the real challenge lies in the vast number and diversity of instructions, as well as their complex relationships with tasks. Since the emphasis of meta-RL research does not address these challenges, its improvements are limited, and it is ineffective in resolving task ambiguity.

---

W3.1: Disconnect between the proposed technical solution and core challenges

A for W3.1:

We respectfully disagree with your view that our method fails to address the core issue of task ambiguity, for reasons below:

• Definition: In our work, task ambiguity refers to the phenomenon in language-conditioned tasks where agents confuse similar tasks due to the vast scale and fuzzy boundaries of instructions, leading to performance degradation. To address this, we tackle task ambiguity from both policy learning and representation learning perspectives.
• Policy Learning: We are the first to theoretically demonstrate that distributional RL methods can effectively enhance an agent's ability to distinguish between tasks, and we apply this to language-conditioned problems.
• Representation Learning: We propose a Trajectory-Wise Semantic Alignment, which leverages the correspondence between instructions and trajectories unique in language-conditioned tasks to apply contrastive learning, which directly enhances task discrimination in the language embedding space.
• Experimental Validation: Our comprehensive experimental results demonstrate the effectiveness of the aforementioned methods. Furthermore, visualization experiments (Figure 7) qualitatively validate that our proposed method effectively solves the task ambiguity. Further, we quantitatively measure the improvement in embedding quality using the Silhouette score[9]. The results demonstrate that both proposed methods contribute to better language embeddings.

Table 2: Silhouette score of task embedding on BabyAI.

W3.2: Trajectory-Wise Semantic Alignment lacks explanation and empirical validation

A for W3.2:

As discussed in A for W3.1, we design Trajectory-Wise Semantic Alignment to enhance the distinguishability of different tasks. To address your concern about the design of this component, we elaborate detailed implementations and experimental validations as follows:

• Implementation: We utilize the correspondence between instructions and trajectories to implement contrastive learning. In the dataset, each trajectory is annotated with an instruction describing its behavior. For a given trajectory, we treat its annotated instruction as a positive sample and select any other instruction annotation from the dataset as a negative sample. The trajectory encoding is derived from the output of a sequence model encoding history information (Equation 8), which is then used alongside language embeddings for contrastive learning training. The implementation details are provided in Figure 13 in Appendix F.

• Empirical validation: We further validate the component's efficacy through ablation studies in Sections 5.3 and 5.4: 1) Figure 7 visually demonstrates how alignment sharpens task representations. 2) Table 6 quantitatively shows significant performance gains, particularly on PutNext tasks: CQL (25.6±2.5) $\to$ DAIL w/o Distributional (39.6±0.6), and DAIL w/o Alignment (39.1±1.0) $\to$ DAIL (57.9±0.9). These results confirm the contribution of the proposed improvement. Additional data and analysis are available in Appendix H.

---

Q1: The fundamental mismatch in experimental setup

A for Q1:

The reviewer’s understanding of task ambiguity aligns with ours, and our experiments have been specifically designed to model this issue. We elaborate on this from two perspectives: (i)relationships between tasks and instructions, and (ii)errors caused by task ambiguity.

• The relationship between instructions and tasks is complex. In practice, multiple instructions may correspond to the same task (e.g., in ALFRED, a single trajectory can be described by three different instructions). This makes task ambiguity highly likely to occur, where different instructions are incorrectly treated as corresponding to the same task.

• Our experiments are designed to address task ambiguity. We observe that baseline methods often fail to learn task embeddings correctly, thereby confusing distinct task objectives(e.g., encoding "pick up a red/blue ball" into similar embeddings). Our goal is to separate the representations of inherently different policies.

As a supplement, in our formulation, the definition of a "task" is solely determined by the reward distribution $p(r|s, a, l)$, and is independent of environment configurations such as map layouts or different initial states.

---

Q2: Clustering patterns between DAIL and DAIL w/o Distributional

A for Q2:

• To address this issue, we quantitatively measure the clustering quality with the Silhouette score[9]. We present the scores in Answer for W3.1 (Table 2). Results indicate that distributional indeed improves task embedding.

• Appendix H presents ablation study results on the distributional component, as shown in Table 6. DAIL significantly outperforms DAIL w/o Distributional, demonstrating the effectiveness of distributional component.

• Moreover, in Figure 7 and the ablation experiment, the difference between CQL and DAIL w/o Alignment also stems from the distributional component. We can observe a substantial improvement in both task embeddings and experimental performance.

From these comparisons, we conclude that the distributional design makes a noteworthy contribution to the task embedding.

---

Q3: Ablation studies on $\alpha$

A for Q3:

As suggested, we do the ablation studies on $\alpha$ coefficient on BabyAI. We present the results for PutNext and All tasks here with complete results added to the revised manuscript.

Table 4: Ablation results on $\alpha$.

The ablation results above show that DAIL is not very sensitive to the choice of $\alpha$ between 0.5 and 2. The performance begins to degrade when the influence of the loss is either too small($\alpha=0.2$) or too large($\alpha=5$).

---

References:

[1] Rt-2: Vision-language-action models transfer web knowledge to robotic control. CoRL. PMLR, 2023.

[2] Context-aware planning and environment-aware memory for instruction following embodied agents. ICCV, 2023.

[3] Mildly conservative q-learning for offline reinforcement learning, NeurIPS 2022.

[4] A minimalist approach to offline reinforcement learning, NeurIPS 2021.

[5] Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning, AAAI 2020.

[6] FOCAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularization, ICLR 2021.

[7] Context shift reduction for offline meta-reinforcement learning, NeurIPS 2023.

[8] Towards an information theoretic framework of context-based offline meta-reinforcement learning, NeurIPS 2024.

[9] Silhouettes: A Graphical Aid to the Interpretation and Validation of Cluster Analysis. Journal of computational and applied mathematics, 1987.

Dear Reviewer,

Thank you for your valuable comments. We address the points you raised in the following, and we hope our responses can resolve your concerns.

---

W1, W2, and Q1: Lack of references to and comparisons with existing research on task representations in RL.

A for W1, W2, and Q1:

Thanks for pointing out the relevant papers—these references are indeed valuable. We have studied all and reproduced several of the papers as the reviewer suggested, and our analysis and experimental results are summarized below:

• LCMDP: While language-conditioned MDP can be viewed as a special case of contextual MDP, the key challenge differs a lot. LCMDP is mainly obstructed by the complexity and diversity of language instructions, which makes it difficult for the agent to ground them to the intended task goals correctly.

• Related works: In our original paper, the Related Work section (Appendix A) focused on task representation of goal-conditioned reinforcement learning. We incorporate an extra discussion of the papers and related research that the reviewer mentioned in the revised version.

• Supplementary experiments： Some papers the reviewer mentioned have a different research focus from our paper. They contribute to learning implicit task information from trajectories, while task information is explicitly provided through instructions in language-conditioned RL. As a result, the improvements of these algorithms are minor, and some cannot even be reproduced due to the experimental setting. We provide the extra experimental results and analysis as follows:

Table 1: Performance of in-distribution BabyAI tasks.

Table 2: Performance of out-of-distribution BabyAI tasks.

• CCM[1]: It introduces contrastive learning methods to enhance transition encoding, and we adapt it to language instruction encoding. The results indicate that this method provides only marginal improvement over CQL, primarily because it overlooks the relationship between language instructions and trajectories.

• PSM[2]: This paper focuses on optimal policy transfer, which relies on learning state representations in similar scenarios and is limited to a single or a small number of tasks. In contrast, our work considers a partially observable environment that varies across runs, with diverse and abundant tasks. As a result, it is not applicable to our setting.

• FOCAL[3]: It introduces distance metrics to separate different trajectory embeddings. We adapt it to language embedding, which shows some improvements on CQL. However, it still fails to account for the relationship between language instructions and trajectories, and thus does not fundamentally resolve the problem of task ambiguity.

• [4]: This paper leverages the idea that two tasks are similar if observing an agent’s performance on one task reduces the uncertainty about its performance on the other. We did not implement it because the task similarity estimate needs rollouts from the environment, which is not available in offline settings.

• CSRO[5] and UNICORN[6]: We adapt the task representation methods from two state-of-the-art meta-RL approaches to language-conditioned setting and conduct reproduction experiments. Their performance was comparable to that of FOCAL.

---

W3 and Q2: Definition 1 is stated without sufficient justification.

A for W3 and Q2:

For language-conditioned policy $\pi(a|s,l)$, when task representations are similar, the corresponding policies should also exhibit similar behaviors. This not only helps prevent abrupt behavioral shifts but is also an important condition for designing robust control systems[8]. Conversely, if two different tasks are mapped to nearby positions in the task representation space, the agent may fail to distinguish between them, potentially resulting in incorrect behaviors.

Moreover, given the complexity and vast number of language instructions, estimation errors in the Q-function can easily propagate backward and cause confusion in the task embeddings. To quantify this phenomenon, we introduce Definition 1.

Consequently, in Definition 1, the "similarity" between tasks specifically refers to the similarity of their true Q-values.

---

W4: Theoretical analysis of sample complexity upper bounds.

A for W4:

Although the sample efficiency gains are not in the exponents, our main focus is on the difference between the predefined task discrimination threshold of value-based ($\delta$) and distributional-based RL ($d$). Following the proof in Appendix D from our paper, we obtain that $\delta \leq d$, when the mean difference of the Q function is small but the distribution difference is large, $\delta<<d$.

We also prove that both $C_{\text{value}}, C_{\text{dist}}\in (0, \mathcal{O}(Q_{\max}^2)]$ are nearly identical in terms of magnitude. In particular, since $Q_{\max}=1$ in most language-conditioned scenarios, their impact on the overall sample efficiency becomes even less.

Q3: Theoretical problems

A for Q3:

• Theoretical analysis: The sample complexity upper bounds are quite tight. Our gains are not in constants inside the log sufficient but in the denominator ($\delta$ and $d$). We believe that even others use different concentration inequalities to achieve better bounds, they can only achieve smaller constant coefficients [9], which have little effect while the index of other variables remains the same.

• Experiments: Moreover, we have computed the AUC(Area-Under-the-Curve) for the learning curves and statistical significance tests on these AUC results at 20 epochs. We present the AUC comparisons of ID learning curves at the key milestones of 10 and 20 epochs below.

  Table 3: AUC of in-distribution BabyAI tasks. Each score is evaluated over 3 seeds.

  Algorithm	10(↑)	20(↑)
  CQL	38,385.3±853.5	81,876.0±848.0
  DAIL w/o Alignment	40,374.7±601.3	85,247.1±540.2
  DAIL w/o Distributional	40,752.9±248.7	85,416.5±543.6
  DAIL	42,370.6±404.8	88,450.5±291.6

  The statistical significance tests show that, at the $\alpha=0.05$ significance level, paired t-tests demonstrate statistically significant performance improvements for the distributional component (p=0.0342). Furthermore, DAIL yields highly significant enhancements (p=0.002). The above results indicate that the distributional component indeed contributes to achieving higher sample efficiency.

---

W5: The method appears to be incremental.

A for W5:

We respectfully disagree that our method appears to be incremental for the following reasons:

• Theory. To the best of our knowledge, this work is the first to focus on the task ambiguity problem inherent in language instruction learning. We clearly define this problem within the language-conditioned RL framework and empirically demonstrate its prevalence in existing RL methods through experimental analysis (Figure 2 and the results for IQL and CQL in Table 1).

• Method. To address this problem, we attempt to tackle it from two perspectives: policy learning and representation learning.

  • For policy learning, we prove, for the first time, that distributional reinforcement learning methods can effectively enhance the agent's ability to distinguish tasks, and apply this approach to language-conditioned tasks.
  • For representation learning, we propose Trajectory-Wise Alignment, which leverages the correspondence between language instructions and trajectories unique to language-conditioned tasks. By introducing self-supervised constraints, this approach directly enhances the distinguishability of different tasks in the language embedding space.

• Experiment results. Comprehensive experiments validate the substantial performance gains achieved by DAIL (Table 6 and Figure 20). These significant gains also underscore the critical importance of addressing the task ambiguity problem in LCMDP, as we have claimed.

Thanks again for your timely and valuable comments. We sincerely hope that our response has cleared the concern you had regarding our problem settings and novelty, and we are looking forward to more discussions.

---

Reference:

[1] Fu et al. Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning, AAAI 2020.

[2] Agarwal et al. Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning, ICLR 2021.

[3] Li et al. FOCAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularization, ICLR 2021.

[4] Mahajan et al. Learning Embeddings for Sequential Tasks Using Population of Agents, IJCAI 2024.

[5] Gao, Yunkai, et al. Context shift reduction for offline meta-reinforcement learning, NeurIPS 2023.

[6] Li, Lanqing, et al. Towards an information theoretic framework of context-based offline meta-reinforcement learning, NeurIPS 2024.

[7] Weng T W, Zhang H, Chen P Y, et al. Evaluating the robustness of neural networks: An extreme value theory approach. ICLR 2018.

[8] Boucheron, et al. Concentration inequalities. 2003. 208-240.

I thank the authors for their rebuttal.

W1, W2, and Q1: Thanks for the clarification.

W3 and Q2: I'm not still not convinced with the justification. Distinction is defined w.r.t. a shared policy \pi. Now, this policy could be anything. Consider two tasks that have the same underlying MDP, but different observations: the first one is grayscale, the second is colored. Now, consider a policy \pi that solves the first task, but not the second. The Q values are going to be different, but the semantics match.

W4: Having some gain (in terms of a constant multiplier) inside the log is equivalent to having a constant gain outside of it. I'm still not convinced by the argument.

Based on the above points, I have increased my score to a 3.

Dear Reviewer,

We would like to thank the reviewer for raising the score! We also appreciate your further comments, and hope our following responses can resolve your remaining concerns.

---

A for W3 and Q2:

We address your concerns from the following two perspectives:

First, we agree that relying on a specific policy is not rigorous enough. Following your suggestion, we have revised "a specific policy $\pi$" to "any $\epsilon$-optimal policy $\pi$ such that $V_{\pi}(s)\geq V^*(s)-\epsilon,\ \forall s \in \mathcal{S}$". This modification makes the definition more general and facilitates a cleaner derivation.

Second, in our setting, task embedding focuses solely on differences in reward functions, assuming a fixed observation space. When both colored and grayscale observations are present in the dataset, two "tasks" with the same semantics share identical rest of the conditions, so their optimal policies with coupled observations (and thus the corresponding $Q^*$) should be equivalent. In our revised version, we restrict Definition 1 to any $\epsilon$-optimal policy rather than a specific policy. Any such policy is capable of solving the task within a minor suboptimality. Consequently, the corresponding $Q_\pi$ are expected to be very similar, and we treat them as the same task under the definition.

---

A for W4:

We believe the advantage of the distributional component primarily lies in its effect on the denominator, i.e, the predefined task discrimination threshold of value-based ($\delta$) and distributional-based RL ($d$). In Appendix D, we prove that the absolute difference between the means of two distributions is upper bounded by their Wasserstein distance, i.e, $W_1(Z\_i,Z\_j)\geq |\mathbb{E}\_{Z_i}(Q)-\mathbb{E}\_{Z\_j}(Q)|$ ($Z\_i,Z\_j$ denotes distribution of Q values for task $i$ and $j$). Therefore, under the distributional setting, we can always define a finer (i.e., bigger) discrimination threshold, $\delta \leq d$. When the mean difference of the Q function is small but the distribution difference is large, we can have $\delta<<d$.

Here is an intuitive example. Consider a multi-modal distribution $Z_1$ that takes the value 0.3 with 50% probability and 0.1 with 50% probability, giving it an expected value of 0.2. In contrast, distribution $Z_2$ always takes the fixed value 0.2. For non-distributional methods, the expectations of both distributions are 0.2, making it impossible to distinguish between them, whereas distributional methods can.

The above example is highly relevant to our setting, while significant task ambiguity exists in language-conditioned tasks. Therefore, as with $Z_1$ and $Z_2$ in the example above, non-distributional methods struggle to distinguish between different tasks, which is supported by our results in AUC (measuring sample efficiency) shown in Table 3 in the previous rebuttal, with paired t-tests demonstrate statistically significant performance improvements for the distributional component (p=0.0342) at the $\alpha=0.05$ significance level.

---

Thanks again for your valuable comments. We sincerely hope our response has cleared your remaining concerns.

Best regards,

The Authors

Dear Reviewer,

Thanks for finding our paper easy to understand and empirically comprehensive. The points you raised are explained in the following.

---

W1: Proposed losses can be found origin from the existing literature

A for W1:

We appreciate your perspective that proposed losses can be found originating from existing literature. While we agree with this perspective, we wish to emphasize that our core contribution lies not in refining specific losses but rather in identifying critical challenges inherent to language instruction tasks and proposing targeted solutions to address them.

To the best of our knowledge, this work is the first to focus on the task ambiguity problem inherent in language instruction learning. We clearly define this problem within the language-conditioned RL framework and empirically demonstrate its prevalence in existing RL methods through experimental analysis (Figure 2 and the results for IQL and CQL in Table 1).

To address this problem, we attempt to tackle it from two perspectives: policy learning and representation learning.

• For policy learning, we prove, for the first time, that distributional reinforcement learning methods can effectively enhance the agent's ability to distinguish tasks, and apply this approach to language-conditioned tasks.
• For representation learning, we propose Trajectory-Wise Alignment, which leverages the correspondence between language instructions and trajectories unique to language-conditioned tasks. By introducing self-supervised constraints, this approach directly enhances the distinguishability of different tasks in the language embedding.

Comprehensive experiments validate the substantial performance gains achieved by DAIL (Table 6 and Figure 20). These significant gains also underscore the critical importance of addressing the task ambiguity problem as we have claimed.

---

W2.1: Confusion of the language embedding relates to policy

A for W2.1:

We claim that the confusion of language embeddings is closely related to the policy for the following reasons.

Language embeddings provide the only information to distinguish between different tasks in Language-conditioned MDP settings. Therefore, when two distinct instructions are mapped to similar language embeddings, the agent struggles to differentiate between them, potentially leading to erroneous behaviors deviating from the original instructions.

On the other hand, when language embeddings are similar, policy behaviors should also exhibit similarity in robust control systems (so-called Lipschitz continuity[1]). As a result, when embeddings exhibit ambiguity, this confusion propagates into the agent's behavioral patterns.

Besides, our ablation studies (Table 6 and Figure 20) also support our claim that policies with more discernible language embeddings have better performances.

W2.2: The points can still be different in high dimensions

A for W2.2:

To address your concern, we quantitatively measure the clustering quality with the Silhouette score[2] using language embeddings before dimensionality reduction. We present the scores for all tasks and PutNext sub-tasks below.

Table 1: Silhouette score of in-distribution BabyAI tasks.

The results demonstrate that both proposed improvements contribute to enhancing the performance of language embeddings in high dimensions.

W2.3: Similar behaviors should have similar embeddings

A for W2.3:

Yes, we agree. The results of our visualization also match your intuition. For example, in the third figure of Figure 7, 'Goto the door' is significantly closer to "Open the door" rather than to 'Goto the ball' since the former two share a lot more in behaviors.

W2.4: Does the benefit come from the distributed value itself rather than the latent embedding?

A for W2.4:

Our partial experimental results support that the advantages of distributional RL are not solely due to the distributional value itself:

• In the toy experiment, when the number of tasks is low, the performance gap between the distributional method and CQL is not substantial. However, as the number of tasks increases, its performance advantage becomes pronounced. This indicates that its benefits extend beyond the distributional value itself.
• Comparing the results of CQL and DAIL w/o Alignment shown in Figure 7 (where the latter figure represents results applying the distributional method on top of CQL), the distributional method clearly enhances the discriminative power of the language embeddings. This improvement is particularly evident in tasks within the "Pickup/Goto the box" and "Pickup/Goto the key" groups, where tasks involving key versus box are distinctly separable."

---

Q1: Would pre-trained components solve the problem directly?

A for Q1:

While pre-trained components can alleviate this issue, they cannot directly solve it. To illustrate this claim, we employ a more recent language encoder [3] (CLIP+LLM2CLIP version), which leverages large language models to enhance the training of linguistic encoders, thereby achieving superior capabilities in language understanding and discrimination. Using this encoder, we retrain our method along with CQL on the BabyAI high-quality dataset. Notably, we introduce an additional baseline, CQL-LLM2CLIP, which directly feeds the LLM2CLIP-embeddings as language embeddings into CQL's policy network without fine-tuning the language encoder. The experimental results are as follows:

Table 2: Success rate of BabyAI tasks with LLM2CLIP.

The experimental results demonstrate that:

• Adopting a stronger pre-trained component shows measurable gains for both CQL-LLM2CLIP and DAIL, confirming that a better pre-trained component can help alleviate the ambiguity problem.

• The CQL vs. CQL-LLM2CLIP comparison demonstrates that despite possessing a powerful pre-trained component, CQL fails to effectively leverage these pre-trained embeddings, with its performance slightly underperforming even the non-finetuned version.

All the above experimental results substantiate our claim that pre-trained components cannot fully resolve the ambiguity problem, emphasizing the critical importance of effectively learning frameworks, precisely the focus of our work.

Thanks again for the valuable comments. We hope our response has cleared your concern. We are looking forward to more discussions.

---

References:

[1] Evaluating the robustness of neural networks: An extreme value theory approach[J]. ICLR, 2018.

[2] Silhouettes: A Graphical Aid to the Interpretation and Validation of Cluster Analysis, Journal of computational and applied mathematics, 1987.

[3] https://github.com/microsoft/LLM2CLIP