DAIL: Beyond Task Ambiguity for Language-Conditioned Reinforcement Learning

TL;DR

DAIL addresses the ambiguity arising from flexible language instructions in language-conditioned RL by combining a distributional policy that preserves full return distributions with a trajectory-wise semantic alignment that maximizes mutual information between trajectories and instructions. The authors provide theoretical insights on sample efficiency for distributional vs value-based methods and validate that DAIL yields clearer, non-ambiguous task representations while improving performance on offline benchmarks like BabyAI SynthLoc and ALFRED. Empirically, DAIL outperforms strong baselines in both in-distribution and out-of-distribution settings, with visualization and clustering analyses illustrating sharper instruction-task alignment. The approach is modular and plug-in friendly, offering a practical route to tackle language-driven task ambiguity in offline and potentially online RL, albeit with acknowledged limitations around real-world deployment and online analysis in future work.

Abstract

Comprehending natural language and following human instructions are critical capabilities for intelligent agents. However, the flexibility of linguistic instructions induces substantial ambiguity across language-conditioned tasks, severely degrading algorithmic performance. To address these limitations, we present a novel method named DAIL (Distributional Aligned Learning), featuring two key components: distributional policy and semantic alignment. Specifically, we provide theoretical results that the value distribution estimation mechanism enhances task differentiability. Meanwhile, the semantic alignment module captures the correspondence between trajectories and linguistic instructions. Extensive experimental results on both structured and visual observation benchmarks demonstrate that DAIL effectively resolves instruction ambiguities, achieving superior performance to baseline methods. Our implementation is available at https://github.com/RunpengXie/Distributional-Aligned-Learning.

Figures (22)

• Figure 1: The framework of our method. The key ideas are: (1) use the distributional language-guided policy to aid task discrimination, (2) use the trajectory-wise semantic alignment module to extract task representation.
• Figure 2: Left: The green flags in the map denote the accessible goals. Middle: An illustration of the mapping between goal positions and instructions. Right: Average success rates over 100 evaluations for each number of instructions and 3 seeds.
• Figure 3: An illustrative case where $Q(s,a)$ functions in two tasks share the same expectations but have different distributions. In this case, traditional RL cannot discriminate between the tasks while distributional RL can. Please refer to Appendix \ref{['appendix:theory']} for more details.
• Figure 4: Left: An example of the SynthLoc task in the BabyAI environment: "put the yellow key next to a green ball". Right: Two example scenes with instructions from ALFRED. Agents are asked to finish specific tasks in the 3D household environment according to received instructions.
• Figure 5: Training curves on in-distribution BabyAI tasks. Success rates are evaluated over 3 seeds.

Theorems & Definitions (7)

• Definition 1: Semantics Instructions Distinction
• Theorem 1: Sample Complexity for Task Instruction Disambiguation
• Definition 1: Restatement of Definition \ref{['def: distintion']}
• Theorem 1: Sample Complexity for Task Instruction Disambiguation
• proof
• Corollary 2
• proof