Discrete Latent Plans via Semantic Skill Abstractions

Authors: Haobin Jiang, Jiangxing Wang, Zongqing Lu
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Hierarchical Learning, Skill Learning, Imitation Learning

Abstract:
Skill learning from language instructions is a critical challenge in developing intelligent agents that can generalize across diverse tasks and follow complex human instructions. Hierarchical methods address this by decomposing the learning problem into multiple levels, where the high-level and low-level policies are mediated through a latent plan space. Effective modeling and learning of this latent plan space are key to enabling robust and interpretable skill learning. In this paper, we introduce LADS, a hierarchical approach that learns language-conditioned discrete latent plans through semantic skill abstractions. Our method decouples the learning of the latent plan space from the language-conditioned high-level policy to improve training stability. First, we incorporate a trajectory encoder to learn a discrete latent space with the low-level policy, regularized by language instructions. Next, we model the high-level policy as a categorical distribution over these discrete latent plans to capture the multi-modality of the dataset. Through experiments in simulated control environments, we demonstrate that LADS outperforms state-of-the-art methods in both skill learning and compositional generalization.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730170782022

thread_type:
official_review

<thread_object>
creator: reviewer_n4ju

<content>
summary:
Learning skills from language instructions poses a significant challenge for current algorithms, particularly when dealing with complex tasks that require sequential execution of multiple skills. Previous research (Garg et al. 2022, Liang et al. 2024, Ju et al. 2024) has addressed this issue by unsupervised learning of constituent skills from language-labeled trajectories. Notably, LISA (Garg et al. 2024) has developed a hierarchical approach to learn interpretable skill codes for each task; however, it encounters stability issues. This work improves upon LISA by decoupling the learning of high-level and low-level policies and introducing additional loss terms to enhance stability. Additionally, it addresses dataset multimodality by modeling the high-level policy as a categorical policy. Their approach, **LADS**, demonstrates more stable training, leading to improved generalization and compositionality performance, with ablation studies analyzing the impact of each modification made.

strengths:
* The authors have pinpointed a specific and relevant issue in the prior algorithm, LISA (Garg et al. 2022), and have methodically sought to address it. They hypothesize that the end-to-end training of high-level and low-level policies in LISA leads to instability, and they have developed a new algorithm to overcome this challenge.

* The paper is well-written, with each loss term clearly explained and supported by thorough proofs and mathematical details.

* The authors have conducted extensive experiments, comparing their method to various baselines from previous research, which bolsters the credibility of their approach and hypotheses. The metrics and evaluations align with prior work, facilitating straightforward comparisons.

* They have achieved impressive results in both compositionality and generalization tasks during evaluation.

* The comprehensive ablation studies provide insights into the sources of performance improvements, allowing readers to gain a deeper understanding of the proposed algorithm.

weaknesses:
* The paper argues that instability issues in LISA stem from the end-to-end training of high-level and low-level policies. While LADS decouples these policies, it still trains the trajectory encoder and low-level policy together in an end-to-end manner. How does this approach enhance stability compared to LISA's method? Additionally, the trajectory encoder in LADS has the benefit of looking ahead by $H$ steps before predicting the next skill code $z$. Could this be LADS's primary advantage over LISA?

* The interpretability results presented in Section 5.5 are somewhat unclear and could be better articulated. I have specific questions in the Questions section below, but I am curious about how the skill codes in LADS compare to those in LISA. Are they interpretable? LISA demonstrated a mapping from skill codes to word clouds and behaviors. Do we observe similar behaviors in LADS? Are the skill codes collapsing like in LISA, or do they exhibit more stability?

* The argument for multi-modality is currently presented in a somewhat unconvincing manner. For example, in Section 4.3, the authors state that the dataset may perform the same skill, such as "open drawer," either quickly or slowly, leading to multimodality. How does this impact the categorical representation of skills, considering both variations would still be captured under the same skill? Furthermore, how does the distribution of skills matter if the low-level policy acts on only one skill at a time? Although the ablation results indicate that the categorical distribution performs slightly better, could this be due to the inherent difficulty of predicting the latent code $z$ directly for the high-level policy?

* Semantic regularization appears to play a crucial role in this algorithm, but a compelling rationale for its necessity is lacking. Why do we require both $J_p$ and $J_q$? Why not use just one? What is the significance of the stop gradients? The explanation in lines 474-476 is confusing; it seems that even $J_p$ alone could suffice for language alignment. Why is $J_q$ also needed?

questions:
In addition to the questions mentioned in the Weaknesses section, I have a few specific inquiries about the paper:

* L252: What does it mean when you say the high-level policy regularizes the latent space?

* I find Figure 4 quite confusing. Could you provide a clearer explanation? Specifically, how is the composite distribution calculated? What does the x-axis represent? Why does the probability of the composite distribution approach 0?

* Figure 5 is also somewhat unclear. What are the labels for the axes, and what do the curves represent? Why are we only considering the first 10 steps?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730501193230

thread_type:
official_review

<thread_object>
creator: reviewer_ayjy

<content>
summary:
The authors introduce a two stage method which starts by learning a discrete latent space for compressing trajectories and then learns a high-level categorical policy in this latent space. They also add a semantic alignment loss to align latent plans with language instructions. The authors show good results in some language-conditioned BC tasks.

strengths:
Originality:
 * The semantic alignment objective seems to be novel

Quality:
 * The authors ablate the main parts of their method and see that they are helpful
 * The authors achieve strong performance compared to the baselines they include
 * The authors provide a nice experiment showing the multimodal capabilities of their work.

Clarity
 * Overall the paper and its conclusions are easy to follow

Significance
 * Language conditioned imitation learning is an important problem and this paper does a good job trying to address it. I find the insight about the semantic alignment interesting.

weaknesses:
Important issues:
* A key weakness of the work is that it is very similar to several recent works that it does not compare to. In particular, [VQ-BeT](https://sjlee.cc/vq-bet/), [QueST](https://arxiv.org/abs/2407.15840) have very similar setups with a discrete skill learning step followed by a high-level policy learning step. QueST is relatively newer but VQ-Bet is a pretty big oversight, and there should be a comparison to one of or both of these works. 
 * VQ-BeT also evaluates in the kitchen environment and it would seem that their results, as well as the results from several baselines, are far better than yours. In particular, you report an average of 2.25 goals achieved while they report 3.78, and their worst-performing non-naive baseline reports 3.09. It's important to add this comparison and explain why that is the case to justify the publication of this paper.

Minor issues:
 * As I mention in the following section, it's difficult to compare the LOREL results to other works because your evaluation schema is different.

questions:
* I find it very surprising that the version of this method without semantic regularization does so poorly. Do you have any insights as to why that might be? Why do other baselines, like VAE, not have a similar problem? Why do algorithms like VQ-BeT, which has a more expressive high-level policy head, not suffer from the same overfitting.
 * Can you please confirm the difference between your work and similar works like VQ-BeT? As far as I can tell, aside from some architectural differences the main difference is the semantic alignment.
 * Line 348 you mention that your evaluation scheme follows previous work, but none of the papers you subsequently cite have the same 'atomic seen / atomic rephrased / composite' division. Can you please explain the relationship between your scheme and the scheme used by LOREL and SkillDiffuser (seen / unseen noun / unseen verb / unseen noun / human provided)? This will help me more directly compare with the results in those works.


Small comments / questions:
 * Equation 10: I don't believe you define $\mathcal{L}_\textrm{CL}$ anywhere
</content>
</thread_object>

</thread>
<thread>
creation date:
1730518157951

thread_type:
official_review

<thread_object>
creator: reviewer_3qsh

<content>
summary:
The paper introduces a hierarchical skill learning framework called LADS (LAnguage-conditioned Discrete latent plans via semantic Skill abstractions) aimed at learning skill abstractions from language instructions to enhance skill generalization across diverse tasks. This method is designed to help AI agents interpret and follow human instructions by modeling a discrete latent plan space, which decouples high-level language-based policy from low-level control.
LADS uses a VQ-VAE model to encode sequence of actions into a discrete latent skill space, which allows the high-level policy to predict categorical distributions over potential skills. The hierarchical structure includes three components: a high-level language-conditioned policy, a low-level skill policy, and a trajectory encoder, all of which contribute to stable training. The method demonstrates its effectiveness in robotic control environments like LOReL and Kitchen, outperforming baselines in skill learning and generalization to new task compositions. The key contributions include a decoupled hierarchical policy, improved discrete skill space interpretability, and enhanced multi-modal skill selection.

strengths:
- **Clear Presentation:**
    - The motivation is clearly presented.
    - The related work section is very well-organized, outlining the relationship between this work and others across three different aspects, which helps readers—especially those unfamiliar with the field—understand the context more clearly.
    - The methods are well-decomposed.

- **Significant Performance Improvement on the Benchmark**

weaknesses:
- **Concern about Scalability**

    According to Appendix B.2, the codebook size for the trajectory encoder is 20, which is relatively small. With a larger dataset that includes more diverse atomic instructions, the method would likely require a larger codebook size, increasing the output dimension of the high-level policy. In such cases, I wonder if:
    
    - *LADS without CP* might outperform, as the high-level policy dimension would not necessarily need to increase.
    - *LADS with VAE* might be advantageous due to its continuous latent space, which doesn’t require an expanded codebook size.
    - A combination of the two, *LADS without CP and with VAE*, might also offer a performance benefit.

- **Concern about Discreteness**

    The set of atomic instructions used in each task appears discrete, meaning each instruction is mutually exclusive. However, if we were to encounter instructions with overlapping meanings but different magnitudes, such as:
    
    - *open the drawer 5 cm, open the drawer 10 cm, open the drawer 15 cm, ...*
    - *move white mug right 5 cm, move white mug right 10 cm, move white mug right 15 cm, ...*

    In such cases, I’m uncertain that having discrete latent plans would be beneficial for task-solving. Perhaps in these scenarios, *LADS with VAE* (as discussed in Section 5.4), with its continuous latent space, could be more effective. This raises the possibility that the proposed method may be overly tailored to tasks with discrete instructions.

- **Need for Pseudocode**

    Given that the proposed framework involves training multiple components, with each component requiring several neural networks, including pseudocode would provide readers with a clearer, holistic understanding of the process.

- **Minor Correction**

    - *Line 192:* The objective function does not include the parameter θ. It may be clearer to adjust the notation to explicitly include θ in the objective function.

questions:
Could you respond to the following two concerns I raised in the weaknesses section?
- *Concern about Scalability*
- *Concern about Discreteness*
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652979051

thread_type:
official_review

<thread_object>
creator: reviewer_rvu8

<content>
summary:
This paper introduces a method for for language-guided. Hierarchical RL. The paper learns a high-level and low-level policy connected through a latent plan space. Unlike previous approaches, these policies are not learned end-to-end but are decoupled by learning a discrete latent space for the high-level policy to sample from. The paper shows results on two language-instructed control environments.

strengths:
The paper is extremely well presented and written. The methodology is quite complex, but the good use of figures (Fig. 1) and a clearly written methodology section (that explains the motivation behind each step of deriving the losses rather than just throwing math at you) makes the method easy to follow.

The paper is clear on what it sets out to do: learn a better way of jointly learning HRL policies in a latent planning space while avoiding some of the mode collapse and other issues of previous works and this goal is validated through experiments.

The experimental evidence of the claims of the paper appear to be well backed up in two environments. There are a few possibly concerns in the experiments, or ways to maybe make this stronger (see weaknesses), but holistically the experiments do show the methodology is effective, which is good.

weaknesses:
The related work on language-guided HRL is a little bit sparse and skews very recent. Some early important works here such as [1], [2], [3] should be included and also possibly serve as inspiration for more simple baselines (see below) for language-based HRL.

I do wonder if there are some simpler, language based HRL baselines that could have been compared to. Language DT is a good baseline, so appreciate that, but I wonder if there is a simple hierarchical baselines that could be done here as well. From the above works, people had much simpler methods of avoiding the collapse between the hierarchical policies. For instance, in [2], the high-level policy was learned through IL and then held fixed while the low-level policies were fine-tuned, with the intermediate space being a hidden state, rather than in language space. 

There are quite a few loss functions that go into this method, L_VQ, the VQ-VQA objective, L_BC the behavioral cloning objective, L_CE, the categorical objective and L_align, the alignment objective. I was really hoping to see a more thorough ablation of all of these losses, but did not see them. Some of these, (w/o VAE or CP) sort of get at these, but not entirely. I would have in particular been curious on how important the BC loss was to these results (it's kind of a footnote in the description of the method but I suspect it's rather important to getting it to work), and how important some of these other losses are, the alignment loss in particular. Especially because this work is quite complicated and incorporates so many losses, it seems important to know how vital each of them actually is (even if for some it will quite obviously fail without, that's good to know).

[1] Yiding Jiang, Shixiang Shane Gu, Kevin P Murphy, and Chelsea Finn. Language as an abstraction for hierarchical deep reinforcement learning. In Advances in Neural Information Processing
Systems, pp. 9414–9426, 2019.
[2] Chen, Valerie, Abhinav Gupta, and Kenneth Marino. "Ask Your Humans: Using Human Instructions to Improve Generalization in Reinforcement Learning." International Conference on Learning Representations 2021.
[3] Hengyuan Hu, Denis Yarats, Qucheng Gong, Yuandong Tian, and Mike Lewis. Hierarchical decision making by generating and following natural language instructions. In Advances in neural
information processing systems, pp. 10025–10034, 2019.

questions:
Line 362ish
What does "procedurally label[ing]" the demonstrations mean? Like, use the simulator to manually label what each of the trajectories are? This isn't really explained.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a method for for language-guided. Hierarchical RL. The paper learns a high-level and low-level policy connected through a latent plan space. Unlike previous approaches, these policies are not learned end-to-end but are decoupled by learning a discrete latent space for the high-level policy to sample from. The paper shows results on two language-instructed control environments.

Overall, the paper seemed to provide a clear and novel approach that significantly outperformed baselines. In terms of weaknesses, the overall method design is the paper is a bit complex, and upon close reading of the paper, some choices seemed a bit arbitrary. For instance, the authors use CLIP for semantic alignment due to the claim that language is very high dimensional, but many of the language instructions in the considered environments are fairly simple. In addition, the experimental evaluation is on existing environments where performance gains are marginal. To further strengthen the paper, the authors could show some novel capabilities the approach enables.