SEAL: SEmantic-Augmented Imitation Learning via Language Model

## Abstract

Abstract Hierarchical Imitation Learning (HIL) is a promising approach for tackling long-horizon decision-making tasks. While it is a challenging task due to the lack of detailed supervisory labels for sub-goal learning, and reliance on hundreds to thousands of expert demonstrations. In this work, we introduce SEAL, a novel framework that leverages Large Language Models (LLMs)’s powerful semantic and world knowledge for both specifying sub-goal space and pre-labeling states to semantically meaningful sub-goal representations without prior knowledge of task hierarchies. SEAL employs a dual-encoder structure, combining supervised LLM-guided sub-goal learning with unsupervised Vector Quantization (VQ) for more robust sub-goal representations. Additionally, SEAL incorporates a transition-augmented low-level planner for improved adaptation to sub-goal transitions. Our experiments demonstrate that SEAL outperforms state-of-the-art HIL methods and LLM-based planning approaches, particularly in settings with small expert datasets and complex long-horizon tasks.

## 1 Introduction

The advancement of LLMs brings transformative change to how agents learn to interact and make decisions (Brohan et al., ; Wang et al., ). LLMs like GPT-4 (Achiam et al., ) possess remarkable semantic understanding ability (Liu et al., ), human-like reasoning capability (Wei et al., ), and rich common sense knowledge (Bubeck et al., ), enabling them extracting insights from language instructions to support decision-making agents (Eigner & Händler, ).

A promising path towards LLM-assisted decision-making is to improve Deep Reinforcement Learning (DRL) agents through reward design (Kwon et al., ; Ma et al., ). Yet DRL often suffers from sample inefficiency and still requires extensive interactions with the environment particularly for tasks with long planning horizons and sparse rewards, where agents are rewarded only upon task completion (). Alternatively, Imitation Learning (IL) can be trained to learn generalizable policies from expert demonstrations (Schaal, ). By learning from successful state-action pairs, IL avoids the sample-expensive exploration and exploitation required in DRL. While IL performance can still be limited in long-horizon tasks due to compounding errors with accumulated errors leading to significant deviations from desired trajectories (Nair & Finn, ).
Hierarchical Imitation Learning (HIL) (Le et al., ) leverages the sub-goal decomposition of long-horizon tasks into a multi-level hierarchy and reduces the relevant state-action space for each sub-goal such as goal-states (Ding et al., ) and Task IDs (Kalashnikov et al., ). To address this, language instruction has come as an aid for sub-goal specification, as such instruction can be both informative and flexible (Stepputtis et al., ). Language-conditioned HIL approaches either train the high-level sub-goal encoder and low-level policy agent separately (Prakash et al., ) or jointly (Hejna et al., ). Though both methods have achieved impressive results, learning language-based sub-goals remains challenging, as it requires a large number of expensive sub-goal labels (Chevalier-Boisvert et al., ). To overcome this, various methods have been proposed to infer both sub-goal boundaries and supervision-free representations (Garg et al., ; Jiang et al., ; Kipf et al., ; Simeonov et al., ). Despite these efforts, these language instructions are unstructured, and often fall short of generalization to newer tasks while not seamlessly integrated into IL policy training (Wang et al., ; Mees et al., ).

The capabilities brought by Large Language Models (LLMs) spark new promises for tackling such instruction challenges in IL and especially HIL settings. It has shown leveraging LLM’s strong reasoning and semantic abilities help break down complex, ambiguous language instructions into manageable steps for a high-level plan (Huang et al., ; Ahn et al., ; Huang et al., ). LLM excels at emulating human-like task decomposition, (Huang et al., ; Wei et al., ) This ability has already been harnessed by researchers to produce high-level plans based on textual task instructions (Ahn et al., ; Prakash et al., ; Huang et al., ). However, while these plans reflect some task hierarchy, they are not directly executable due to the reliance on pre-trained low-level policy agents for primitive actions (Prakash et al., ). Moreover, most high-level plans are static and require frequent LLM calls to update as states change, which is highly costly (Song et al., ). These challenges limit applicability of LLM-based approaches in assisting HIL. Motivated by these promises and challenges, we want to address the following question: *“Can pre-trained LLMs serves as a prior to define task’s hierarchical structure, establish sub-goal library autonomously, and use them to closely guide both high-level sub-goal learning and low-level agent?”*

In this paper, we explore the possibility of using LLM-generated high-level plans to assist both sub-goal learning and policy training in Hierarchical Imitation Learning. We introduce SEmantic-Augmented Imitation Learning (SEAL), a novel hierarchical imitation learning framework that utilizes pretrained LLMs to generate semantically meaningful sub-goals. Give the fact that sub-goal are also oftenwise labeled by human and understandable by LLMs, SEAL represents each sub-goal with one-hot latent vectors, and employs these representations to convert expert demonstration states into supervisory labels for high-level latent variable learning. To enhance the learning process, SEAL features a dual-encoder structure for sub-goal representation. One encoder learns sub-goal vectors in a supervised manner using LLM-provided labels, while the other leverages unsupervised Vector Quantization (VQ) (Wang et al., ) to map expert demonstration states to latent sub-goal representations. The effectiveness of the learned latent variables is evaluated within the task environment by comparing their contribution to selecting optimal actions, with success rate as the confidence metric. We design a weighted combination of the two encoders’ outputs as the final sub-goal representation, thereby reducing over-reliance on a weaker encoder, mitigating over-fitting, and improving overall robustness. Additionally, we present a transition-augmented lower-level policy agent that prioritizes intermediate states corresponding to sub-goal transitions by assigning higher weights, reflecting the hierarchical structure of long-horizon tasks.
Extensive experiments on two tasks KeyDoor and Grid-World and show that SEAL can outperform several state-of-the-art supervised and unsupervised HIL approaches. To summarize, our main contributions include:

- •
We propose SEAL, a novel hierarchical imitation learning framework that leverages Language Models to generate high-level plans and provide semantically meaningful sub-goal representations without any prior knowledge of task hierarchical structure.
- •
To enhance SEAL’s effectiveness, we introduce a dual-encoder structure that combines supervised LLM-based sub-goal learning with unsupervised VQ-based representations, ensuring reliability and robustness in sub-goal learning. Additionally, SEAL incorporates a transition-augmented low-level planner for better adaptation to challenging intermediate states where sub-goal transitions occur.
- •
We demonstrate that SEAL outperforms several state-of-the-art HIL approaches, particularly on small expert datasets, and shows superior adaptability to longer-range composition tasks and task variations.

## 2 Related Works

Imitation Learning. Imitation Learning encompasses two primary approaches: Behavioral Cloning (BC) (Bain & Sammut, ) and Inverse Reinforcement Learning (IRL) (Ng et al., ). BC relies on a pre-collected expert dataset of demonstrations, where the agent learns to mimic the actions in an offline manner. While BC is simple to implement, it is prone to compounding errors, particularly when the agent encounters states not present in the expert’s demonstrations (Zhang, ). In contrast, IRL methods (Ho & Ermon, ; Reddy et al., ; Brantley et al., ) involve interacting with the environment to collect additional demonstrations, aiming to infer the underlying reward function that the expert is optimizing. The agent then learns by optimizing this inferred reward. However, IRL approaches are more challenging to implement (Kurach et al., ), typically requiring more computational resources and data. In this work, we primarily adopt the BC architecture in a hierarchical setting, while incorporating insights from IRL by using environment interactions to validate the reliability of learned latent sub-goal variables.

Bi-Level Planning and Execution. Hierarchical Imitation Learning (HIL) enhances the ability of imitation learning agents to tackle complex, long-horizon tasks by breaking them down into smaller sub-goals and conditioning the agent’s behavior on those sub-goals. The high-level agent chooses the sub-goals, while the low-level agent learns to accomplish specific controls under selected sub-goals (Jing et al., ). Many HIL approaches, such as Hierarchical Behavior Cloning (Le et al., ) and Thought Cloning (Hu & Clune, ), rely on supervisory labels for sub-goal learning, but such annotations are often difficult to obtain. To address this limitation, unsupervised methods like Option-GAIL (Jing et al., ), LOVE (Jiang et al., ), SDIL (Zhao et al., ), and CompILE (Kipf et al., ) have been developed to infer sub-goals directly from expert trajectories. However, the lack of labeled guidance in these approaches makes meaningful sub-goal discovery more challenging and hence reduces the reliability of the learned policies.

LLMs for Planning.
Large Language Models (LLMs) have demonstrated significant potential in decision-making processes. Direct generation of action sequences usually do not lead to accurate plans (Silver et al., ; Valmeekam et al., ; Kambhampati et al., ). Recent studies have successfully utilized LLMs to decompose natural language task instructions into executable high-level plans, represented as a sequence of intermediate sub-goals (Ahn et al., ; Prakash et al., ; Huang et al., ). While LLMs can be also applied to translate user-given language instructions to symbolic goals (Mavrogiannis et al., ; Xie et al., ) Additionally, LLMs can function as encoders, identifying current sub-goals based on both observations (sometimes images) and language task descriptions to facilitate high-level plan execution (Fu et al., ; Malato et al., ; Du et al., ). However, these approaches typically still depend on pre-trained low-level planners for generating executable primitive actions. In this work, we leverage LLM-generated high-level plans to assist in learning both sub-goals and low-level actions simultaneously.

## 3 Preliminary

In this paper, we look into the long-horizon, compositional decision-making problem as as a discrete-time, finite-step Markov Decision Process (MDP). MDP can be represented by a tuple $(\mathcal{S},\mathcal{A},\mathcal{T},r,\gamma)$, where $\mathcal{S}$, $\mathcal{A}$ denotes the state and action space, $\mathcal{T}(s_{t+1}|s_{t},a_{t}):\mathcal{S}\times\mathcal{A}\rightarrow
\mathcal{S}$ denotes the transition function, $r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is the reward function and $\gamma\in[0,1]$ is the discount factor.

In standard settings of Hierarchical Imitation Learning (HIL), instead of having access to the reward function $r$, the agent has access to a dataset of expert demonstrations $\mathcal{D}=\{\tau^{e}_{1},\tau^{e}_{2},...,\tau^{e}_{N}\}$, which contains $N$ expert trajectory sequences consisting of state-action pairs $\{(s_{t},a_{t})\}$, where $s_{t}\in\mathcal{S}$, $a_{t}\in\mathcal{A}$, $T$ is the time horizon for planning, $0\leq t\leq T$. In this paper, the expert trajectories are not labeled with any rewards nor subhorizon segments. We assume HIL agents operate in a two-level hierarchy though our method can also be applied to problems with more levels:

- •
High-level Sub-goal Encoder $\pi_{H}(g_{t}|s_{t})$: Selects a sub-goal $g_{t}\in\mathcal{G}$ based on the current states $s_{t}$, where $\mathcal{G}$ is the space of sub-goals.
- •
Low-level Policy Agent $\pi_{L}(a_{t}|g_{t},s_{t})$: Executes actions conditioned on both the current state $s_{t}$ and sub-goal $g_{t}$.

In this work, we focus on settings where agents lack access to the sub-goal space $\mathcal{G}$, relying instead on an oracle full task instruction $\mathcal{M}$ in natural language. While well-defined $\mathcal{G}$ aids efficient HIL agent learning (Hauskrecht et al., ), its acquisition is difficult due to missing task-specific knowledge (Nachum et al., ; Kim et al., ). Natural language task instructions, though easier to obtain as they are common-used commands from human (Stepputtis et al., ), are hard to map to hierarchical structures due to their complex and ambiguous nature (Zhang & Chai, ; Ahuja et al., ; Ju et al., ). In this work, we investigate leveraging LLMs to parameterize $\mathcal{G}$ from language instructions with its powerful semantic and world knowledge, and pre-label states in $\mathcal{D}$ to guide effective learning of $\pi_{H}$ and $\pi_{L}$ in hierarchical imitation learning.

## 4 SEAL for Hierarchical Imitation Learning

Figure: Figure 1: Overview of SEAL Architecture: The LLM aids in discovering sub-goal spaces for the task by semantically decomposing the full-task instruction and labeling each state with a reference latent vector that represents its corresponding sub-goal. These reference labels are then used to train a high-level sub-goal encoder, which works in conjunction with an unsupervised VQ encoder.
Refer to caption: extracted/5897892/Figures/SEAL-New.png

The key idea of SEAL is to learn high-level sub-goal representations using supervisory labels generated by LLMs. In previous works, such labels were typically provided by human experts via instructions (Pan et al., ; Le et al., ), making them expensive to obtain. However, with the assistance of LLMs, we introduce an efficient and reliable method to automatically generate labels that map states to sub-goals. Specifically, LLMs are used to semantically extract a high-level plan from the full-task language instruction $\mathcal{M}$ and map states in expert demonstrations to sub-goals within this plan. Using these learned sub-goal representations, the model then learns the corresponding low-level actions. An overview of our SEAL framework is illustrated in Fig. 1.

### 4.1 Pretrained LLMs for Guiding Sub-Goals Learning

Our key design of leveraging LLMs to guide high-level sub-goals learning can be divided into two stages: (i) Use LLM-generated high-level plan based on full-task instruction as sub-goal space (ii) Use LLMs to encode states in expert demonstrations to sub-goal representations.

Derive Sub-goal Space of Task Prior works have demonstrated that LLMs can establish a meaningful chain of sub-tasks from task instruction as high-level plan (Huang et al. (); Prakash et al. (); Singh et al. ()). Yet few of them incorporate it with Hierarchical Imitation Learning (HIL). In SEAL, we use LLMs to specify the unknown sub-goal space $\mathcal{G}$ in HIL formulations. Feeding LLMs with the full-task language instruction $\mathcal{M}$, we notice that the decomposed sub-goals in high-level plan naturally consist of a language-based sub-goal set: $\{\hat{g}^{1},\hat{g}^{2},...,\hat{g}^{K}\}=f_{llm}(\mathcal{M})$, where $K$ is the total number of generated sub-goals. We treat this estimated sub-goal dataset as the finite sub-goal space: $\mathcal{G}=\{\hat{g}^{1},\hat{g}^{2},...,\hat{g}^{K}\}$.

Labeling Sub-goals for States in Expert Dataset
After devising the sub-goal space $\mathcal{G}$ with LLM-generated sub-goals, we use them to map states $s_{t}\in\mathcal{D}$ to a sub-goal latent space. These LLM-defined labels guide the high-level encoder to learn task-relevant sub-goal representations.
To parameterize the language-based sub-goals $\hat{g}^{i}\in\mathcal{G}$ and facilitate learning, we establish a codebook $\mathcal{C}=\{z^{1},z^{2},...,z^{K}\}$, where each latent variable $z^{i}\in\mathbb{R}^{K}$ is a one-hot vector (i.e. $i$-th element in $z^{i}$ equals to 1, others equal to 0, $i=1,2,...,K$) representing sub-goal $\hat{g}^{i}$ in $\mathcal{G}$. We then prompt the same LLM to perform a encoding function $h_{llm}$, which map $s_{t}$ to latent vector $z_{t}^{(ref)}\in\mathcal{C}$ by checking whether it belongs to sub-goal $\hat{g}^{i}\in\mathcal{G}$:
$z_{t}^{(ref)}=h_{llm}(s_{t},\mathcal{G})$.
We stipulate the output of LLM must be ‘yes’ or ‘no’ and then convert it to integer 1 or 0, as this form of answer has shown to be more reliable than the open-ended answer (Du et al. ()). By repeatedly asking $K$ times we can finally establish the $K$-dim latent variable $z_{t}^{(ref)}$ which represents the sub-goal for all $s_{t}$ in $\mathcal{D}$. We use these LLM-given latent representations $z_{t}^{(ref)}$ as supervisory labels for high-level sub-goal encoder training in HIL. Once we obtain these labels, we have no need to interact to LLMs later.

### 4.2 Dual-Encoder for Sub-goal Identification

Naturally, we consider using these LLM-generated labels for sub-goal representations to train a high-level sub-goal encoder $\pi_{H}(s_{t})$ in a supervised manner. Compared to previous unsupervised approaches, this supervised method helps reduce the randomness of output sub-goals by leveraging the guidance provided by the labels. However, it is prone to over-fitting on the training dataset. To address this challenge, inspired by (Ranzato & Szummer (); Le et al. ()), we propose a Dual-Encoder structure for high-level sub-goal identification. This design integrates both a supervised learning encoder and an unsupervised learning encoder, producing a weighted-average sub-goal representation. The weighted combination allows for flexibility, prioritizing the encoder that performs better for a particular task or dataset, ultimately enhancing robustness and improving generalization.

Supervised LLM-Label-based Encoder Considering that the codebook $\mathcal{C}$, representing the sub-goal space $\mathcal{G}$, is discrete and finite, we formulate the supervised sub-goal learning as a multi-class classification problem. To train this supervised learning encoder $\pi_{H}^{(llm)}$, we define the sub-goal learning objective by maximizing the log-likelihood of the labels generated by the LLMs:

$$ $\mathcal{L}_{H}^{(llm)}=\mathbb{E}_{(s_{t},z_{t}^{(ref)})\sim\mathcal{D}}\;- \text{log}\;\pi_{H}^{(llm)}(z_{t}^{(ref)}|s_{t}).$ (1) $$

Unsupervised VQ Encoder Given the codebook $\mathcal{C}=\{z^{1},z^{2},...,z^{K}\}$, we apply Vector Quantization (VQ) (Van Den Oord et al. ()) to design the unsupervised sub-goal encoder in our SEAL framework. It is a widely used approach that can map the input state $s_{t}$ to a finite, discrete latent space like $\mathcal{C}$. In VQ, the encoder $\pi_{H}^{(vq)}$ first predicts a continuous latent vector: $z^{(con)}_{t}=\pi_{H}^{(vq)}(s_{t})$. This latent vector is then matched to the closest entry in $\mathcal{C}$:

$$ $z^{(vq)}_{t}=\text{argmin}_{z^{i}\in\mathcal{C}}\quad\|z^{(con)}_{t}-z^{i}\|_{ 2}^{2}.$ (2) $$

The learning objective of $\pi_{H}^{(vq)}$, named commitment loss, encourages the predicted continuous latent vector $z^{(con)}_{t}$ to cluster to the final output sub-goal representation $z^{(vq)}_{t}$:

$$ $\mathcal{L}_{H}^{(vq)}=\mathbb{E}_{(s_{t})\sim\mathcal{D}}\quad\|\textbf{sq}(z ^{(vq)}_{t})-z^{(con)}_{t}\|_{2}^{2};$ (3) $$

where $\textbf{sq}(\cdot)$ denotes stop-gradient operation.

### 4.3 Transition-Augmented Low-level Policy

We compute a weighted-average vector $z_{t}$ over $z^{(llm)}_{t}$,$z^{(vq)}_{t}$ obtained by dual-encoders to finalize the predicted sub-goal representation:

$$ $z_{t}=W_{vq}z^{(vq)}_{t}+W_{llm}z^{(llm)}_{t};$ (4) $$

where the weights $W_{vq}$ and $W_{llm}$ quantifies how the predicted sub-goal representations $z^{(vq)}_{t}$ and $z^{(llm)}_{t}$ contribute to the task completion success rate. The weights are updated by validations during the training process. The update details will be demonstrated in Section 4.4.

Given the predicted sub-goal representations $z_{t}$ for each $s_{t}$ in the expert dataset, normally the low-level policy agent follows a goal-conditioned behavioral cloning (GC-BC) architecture. It is trained by maximizing the log-likelihood of the actions in the expert dataset, using the sub-goal representations as auxiliary inputs:

$$ $\mathcal{L}_{GC-BC}=\mathbb{E}_{(s_{t},a_{t},z_{t})\sim\mathcal{D}}\;-\text{ log}\;\pi_{L}(a_{t}|s_{t},z_{t}).$ (5) $$

However, this low-level policy design overlooks the imbalanced distribution and importance of the hierarchical structure captured by high-level sub-goal encoders. Several studies have highlighted that certain states, where transitions between sub-goals occur in long-horizon demonstrations, have a significant impact on the policy agent’s performance (Jain & Unhelkar, ; Zhai et al., ; Wen et al., ). Despite their critical role, these states make up only a small portion of expert demonstrations. Successfully reaching these intermediate states and taking appropriate actions improves sub-goal completion, thereby increasing the overall task success rate. We formally define these states as intermediate states:

Definition 4.3.1. (Intermediate States). Let $s_{t}\in\mathcal{S},0\leq t\leq T$ be a state observed when running the HIL agent, $z_{t}$ is its corresponding latent variable learnt by high-level encoder $\pi_{H}$ that represents sub-goal. $s_{t+1}$ is the following state. The state $s_{t}$ is defined as an intermediate state only when the sub-goal changes: $z_{t+1}\neq z_{t}$.
Due to the scarcity of these intermediate states, it becomes very challenging to imitate the correct behavior in such states. To address this issue, inspired by the practice of assigning extra rewards to sub-goal transition points in hierarchical RL (Ye et al. (); Berner et al. (); Zhai et al. (); Wen et al. ()), we augment the importance of these intermediate states by assigning higher weights to them in the low-level policy training loss:

$$ $\mathcal{L}_{L}=\mathbb{E}_{(s_{t},a_{t},z_{t})\sim\mathcal{D}}\quad-e^{\|z_{t +1}-z_{t}\|_{2}^{2}}\log\;\pi_{L}(a_{t}|s_{t},z_{t});$ (6) $$

where the term $e^{\|z_{t+1}-z_{t}\|_{2}^{2}}$ measures the L2-distance between the current sub-goal representation $z_{t}$ and the next sub-goal $z_{t+1}$. Given that $z_{t}$ is a one-hot vector, we have the term:

$$ $e^{\|z_{t+1}-z_{t}\|^{2}_{2}}=\left\{\begin{aligned} &e,\text{ if $z_{t+1}\neq z _{t}$}\\ &1,\text{ if $z_{t+1}=z_{t}$}\end{aligned}\right.$ (7) $$

Thus, this term can serve as an adaptive weight to enhance the imitation of expert behavior at intermediate states. By incorporating this transition-augmented low-level policy design, we emphasize the importance of sub-goal transitions, and in simulations we also observe this design can greatly help agents make transitions across each sub-goal.

### 4.4 Training

We train our SEAL model end-to-end, jointly updating parameters of $\pi_{H}$ and $\pi_{L}$ by minimizing
the loss function $\mathcal{L}=\beta\mathcal{L}_{H}+\mathcal{L}_{L}$, where $\beta$ is a hyper-parameter that controls the weight of high-level sub-goal learning in relation to the overall training process. Additionally, in order to evaluate the reliability of the latent variables predicted by the VQ encoder and LLM-Label-based encoder and determine the weight combination that can better improve task performance, we keep validating the success rates of those two different latent variables in the environment during training. Based on the validation results, we dynamically update the weights $W_{vq}$ and $W_{llm}$ in Eq. 4.

For validation, we simultaneously execute actions conditioned on both the VQ-encoder and the LLM-label-based encoder: $a_{t}^{(vq)}=\pi_{L}(s_{t},z^{(vq)}_{t})$ and $a_{t}^{(llm)}=\pi_{L}(s_{t},z^{(llm)}_{t})$. We then run episodes to test the different success rates, $SR_{vq}$ and $SR_{llm}$, for completing the full task. The updated weights $W_{vq}$ and $W_{llm}$ are then computed as $W_{vq}={SR_{vq}}/{(SR_{llm}+SR_{vq})};\;W_{llm}={SR_{llm}}/{(SR_{llm}+SR_{vq})}$ respectively. $W_{vq}$, $W_{llm}$ measure the relative task-completion performance of the policy agent under the guidance of $z^{(vq)}_{t}$ and $z^{(llm)}_{t}$, respectively. We refer to these weights as confidences, indicating the preference for trust between $z^{(vq)}_{t}$ and $z^{(llm)}_{t}$.

We also use these weights to finalize the overall training loss of SEAL as a weighted combination of two end-to-end losses under guidance $z^{(llm)}_{t}$ and $z^{(vq)}_{t}$. We finalize the overall training loss of SEAL by using a weighted combination of two end-to-end losses, conditioned on $z^{(llm)}_{t}$ and $z^{(vq)}_{t}$, with the same weights $W_{vq}$, $W_{llm}$ determining the contribution of each loss:

$$ $\displaystyle\mathcal{L}_{vq}=\beta\mathcal{L}_{H}^{(vq)}(s_{t})+\mathcal{L}_{ L}(s_{t},z^{(vq)}_{t});\mathcal{L}_{llm}=\beta\mathcal{L}_{H}^{(llm)}(s_{t})+ \mathcal{L}_{L}(s_{t},z^{(llm)}_{t});\mathcal{L}_{SEAL}=W_{vq}\mathcal{L}_{vq} +W_{llm}\mathcal{L}_{llm}.$ (8) $$

Since the low-level policy agent’s actions are conditioned on the latent sub-goal representations, minimizing this weighted-combination loss $\mathcal{L}_{SEAL}$ allows our SEAL to adapt the trainable parameters of the low-level policy based on task-completion performance. This approach helps the agent make better decisions by adjusting to updated latent predictions $z_{t}=W_{vq}z^{(vq)}_{t}+W_{llm}z^{(llm)}_{t}$ during training process.
As a result, our SEAL framework can continuously adapt both the high-level sub-goal encoders and the low-level policy agent, leading to more reliable and robust sub-goal representations, as well as improved decision-making for action selection. The complete algorithm for SEAL is illustrated in Algorithm 1.

Figure: Algorithm 1 SEmantic-Augmented Imitation Learning (SEAL) via Language Model

## 5 Experiments

In this section, we evaluate the performance of SEAL on two long-horizon compositional tasks: KeyDoor and Grid-World. We compare SEAL’s performance with various baselines, including non-hierarchical, unsupervised, and supervised hierarchical IL methods, in both large and small expert dataset scenarios. Following this, we analyze how SEAL enhances task completion performance.

### 5.1 Simulation Environmental Setup

KeyDoor The MiniGrid Dataset (Chevalier-Boisvert et al., ) is a collection of grid-based environments designed for evaluating reinforcement learning and imitation learning algorithms in tasks requiring navigation, exploration, and planning. Among these environments, we start with KeyDoor, an easy-level compositional task that requires the player to move to the key and pick up it to unlock the door. To add complexity, we enlarge the original 3 $\times$ 3 grid environment to 10 $\times$ 10 size, and randomly initialize the locations of player, key and door for each episode. To facilitate understanding by LLMs, we convert the environment into a vector-based state, with elements including the coordinates of the player, key, and door, as well as the different statuses of the key (picked or not) and door (locked or not). The maximum time-steps $T$ of one episode is set to 100. We evaluate our SEAL on expert datasets with 30, 100, 150, 200 demonstrations generated by an expert bot.

Grid-World The environment is a 10x10 grid world with a single player and multiple objects randomly distributed at various locations. The player’s objective is to visit and pick up these objects in a specific order. This task is more challenging than KeyDoor due to its longer-range compositional nature, involving more sub-goals. In this work, we set the number of objects in the grid world to range from 3 to 5, to test SEAL’s effectiveness in solving longer-range tasks. Similar to KeyDoor, the fully observed environment is converted into a vector-based state, with elements representing the coordinates of the player and objects, as well as their statuses (picked or not). The maximum time-steps per episode is set to 100. We evaluate SEAL on expert datasets with 200, 300, and 400 demonstrations generated by an expert bot.

### 5.2 Baselines

We compare our SEAL with one non-hierarchical learning approach: Behavioral Cloning (BC) (Bain & Sammut, ), two unsupervised learning approaches: LISA (Garg et al., ) and SDIL (Zhao et al., ), and one LLM-enabled supervised learning approaches Thought Cloning (Hu & Clune, ). The detailed settings are listed below:

Behavioral Cloning (BC): The classical imitation learning method with non-hierarchy. We train the policy agent $\pi(a_{t}|s_{t})$ by maximizing the log-likelihood: $\mathcal{L}_{BC}=\mathbb{E}_{(s_{t},a_{t})\in\mathcal{D}}\quad-\text{log}\pi(a
_{t}|s_{t})$.
LISA: A HIL approach with unsupervised VQ-based sub-goal learner. We implement the low-level policy agent with only current state $s_{t}$ as input, rather than a sequence of previous states in original settings, since we assume the task is MDP.
SDIL:A HIL approach with an unsupervised sub-goal learner by implementing only the skill discovery component of SDIL, omitting the skill optimality estimation since our expert bot generating demonstrations is optimal. SDIL selects the final sub-goal representation $z_{t}$ by: $\operatorname*{arg\,max}_{i}\quad\frac{1/D(z^{i},z^{{}^{\prime}}_{t})}{\sum_{i
=1}^{K}1/D(z^{i},z^{{}^{\prime}}_{t})}$, where $D$ denotes the Euclidean distance, $z^{{}^{\prime}}_{t}$ is the continuous output vector of $\pi_{H}$, $z^{i}\in\mathcal{C}$. To make this selection differentiable for gradient back-propagate, we adopt Gumbel-Softmax (Jang et al., ) to replace argmax operation.
Thought Cloning (TC): A HIL approach with supervised sub-goal learner. TC consists of a thought generator $\pi_{u}(th_{t}|s_{t},th_{t-1})$ (where $th_{t}$ equals to our $z_{t}$) and an action generator $\pi_{l}(a_{t}|s_{t},th_{t})$. $\pi_{u}$ requires supervisory labels training. We apply the LLM-generated sub-goal representations $z_{t}^{(ref)}$ as labels.

We also present results from two variants of SEAL: SEAL-L, which relies entirely on the LLM-label-based high-level sub-goal encoder, and SEAL, which uses the dual-encoder design. SEAL-L is compared with TC to highlight the effectiveness of the low-level transition-augmented design in supervised learning, while SEAL demonstrates the superiority of the dual-encoder approach compared with SEAL-L.

### 5.3 Main Results

We first evaluate the effectiveness of our SEAL approach in two environments: KeyDoor and Grid-World with 3 objects. These environments are compositional in nature, and the relatively small number of sub-goals simplifies high-level sub-goal encoder learning in HIL settings. To break down the full task instructions into sub-goals, we utilize the latest OpenAI-developed LLM, GPT-4o (Islam & Moushi ()). The sub-goal count for KeyDoor is $K=4$, and for Grid-World with 3 objects, it is $K=6$. These identified sub-goals are then used to label states in the expert dataset for supervised high-level encoder training. For a fair comparison, we set the number of sub-goals $K$ in the unsupervised baselines LISA and SDIL to match the LLM-determined sub-goal numbers used in SEAL. Additional details are provided in the Appendix.

**Table 1: Simulation Results: Success rates (ranging from 0 to 1) for completing the tasks of KeyDoor and Grid-World (3 Objects), averaged over 5 random seeds. Our SEAL approach outperforms others in most cases. The best-performing method is highlighted in bold.**
| Task | # Traj | BC | LISA | SDIL | TC | SEAL-L | SEAL |
| --- | --- | --- | --- | --- | --- | --- | --- |
| KeyDoor | 30 | 0.09$\pm$0.02 | 0.09$\pm$0.02 | 0.23$\pm$0.05 | 0.26$\pm$0.02 | 0.27$\pm$0.06 | 0.30$\pm$0.04 |
|  | 100 | 0.50$\pm$0.06 | 0.53$\pm$0.05 | 0.45$\pm$0.04 | 0.50$\pm$0.03 | 0.52$\pm$0.02 | 0.56$\pm$0.03 |
|  | 150 | 0.67$\pm$0.05 | 0.66$\pm$0.03 | 0.63$\pm$0.05 | 0.69$\pm$0.05 | 0.68$\pm$0.03 | 0.75$\pm$0.04 |
|  | 200 | 0.74$\pm$0.02 | 0.69$\pm$0.04 | 0.70$\pm$0.04 | 0.70$\pm$0.02 | 0.76$\pm$0.04 | 0.82$\pm$0.04 |
| GridWorld | 200 | 0.26$\pm$0.04 | 0.24$\pm$0.03 | 0.43$\pm$0.04 | 0.44$\pm$0.03 | 0.39$\pm$0.04 | 0.29$\pm$0.03 |
|  | 300 | 0.31$\pm$0.04 | 0.44$\pm$0.05 | 0.48$\pm$0.01 | 0.52$\pm$0.07 | 0.65$\pm$0.04 | 0.61$\pm$0.02 |
|  | 400 | 0.48$\pm$0.04 | 0.53$\pm$0.03 | 0.62$\pm$0.04 | 0.62$\pm$0.04 | 0.83$\pm$0.02 | 0.85$\pm$0.02 |

We train the models using randomly sampled expert demonstrations, with 30, 100, 150, and 200 samples for the KeyDoor environments, and 200, 300, and 400 samples for Grid-World with 3 objects. Since Grid-World is a more complex compositional task, we collect additional expert data for this environment. The task success rates from the simulations are presented in Table 1. The results show in most cases, our approach either outperforms or is competitive with the other baselines.

We observe that with fewer training demonstrations, both the HIL baselines and our SEAL outperform non-hierarchical behavior cloning (BC), as they benefit from the additional information learned from hierarchical sub-goal structures. However, as the amount of training data increases, the gap between BC and HIL methods narrows, especially in simpler environments like KeyDoor, where BC begins to generalize better due to its exposure to more scenarios. Despite this, SEAL maintains a consistent performance advantage do to our Transition-Augmented Low-Level policy agent, which emphasizes imitation the expert actions in challenging sub-goal transition states. As shown in Table 2, SEAL consistently achieves higher sub-goal completion rates, demonstrating its superior ability to select appropriate actions at those critical intermediate states with sub-goal transitions.

**Table 2: Success rates of sub-goals completion in both KeyDoor and Grid-World, averaged over 5 random seeds. For the KeyDoor environment, the sub-goal is to pick up the key, while for Grid-World with 3 objects, the sub-goals are to pick up object 1 and object 2.**
| Sub-goals | # Traj | BC | LISA | SDIL | TC | SEAL-L | SEAL |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Pick up the Key (KeyDoor) | 30 | 0.29$\pm$0.05 | 0.22$\pm$0.03 | 0.42$\pm$0.05 | 0.52$\pm$0.04 | 0.55$\pm$0.10 | 0.56$\pm$0.04 |
|  | 100 | 0.78$\pm$0.04 | 0.77$\pm$0.06 | 0.65$\pm$0.02 | 0.67$\pm$0.02 | 0.82$\pm$0.05 | 0.80$\pm$0.02 |
|  | 150 | 0.81$\pm$0.06 | 0.81$\pm$0.03 | 0.80$\pm$0.02 | 0.82$\pm$0.03 | 0.93$\pm$0.02 | 0.93$\pm$0.04 |
|  | 200 | 0.87$\pm$0.01 | 0.88$\pm$0.03 | 0.83$\pm$0.03 | 0.86$\pm$0.01 | 0.97$\pm$0.02 | 0.98$\pm$0.01 |
| Pick up Object 1 (GridWorld) | 200 | 0.58$\pm$0.09 | 0.58$\pm$0.05 | 0.79$\pm$0.02 | 0.78$\pm$0.04 | 0.83$\pm$0.04 | 0.67$\pm$0.03 |
|  | 300 | 0.64$\pm$0.06 | 0.71$\pm$0.07 | 0.75$\pm$0.02 | 0.85$\pm$0.03 | 0.90$\pm$0.04 | 0.85$\pm$0.04 |
|  | 400 | 0.75$\pm$0.03 | 0.79$\pm$0.03 | 0.85$\pm$0.03 | 0.87$\pm$0.04 | 0.95$\pm$0.02 | 0.98$\pm$0.01 |
| Pick up Object 2 (GridWorld) | 200 | 0.39$\pm$0.09 | 0.36$\pm$0.06 | 0.56$\pm$0.04 | 0.64$\pm$0.05 | 0.61$\pm$0.04 | 0.50$\pm$0.05 |
|  | 300 | 0.44$\pm$0.06 | 0.55$\pm$0.06 | 0.59$\pm$0.03 | 0.52$\pm$0.05 | 0.80$\pm$0.04 | 0.73$\pm$0.03 |
|  | 400 | 0.57$\pm$0.05 | 0.63$\pm$0.04 | 0.70$\pm$0.04 | 0.62$\pm$0.03 | 0.90$\pm$0.02 | 0.89$\pm$0.01 |

We observe that the dual high-level encoder design in our SEAL enhances performance. Compared to relying solely on the LLM-label-based sub-goal encoder, SEAL’s dual-encoder design demonstrates higher success rates across both tasks. As shown in Fig. 2, SEAL achieves greater prediction accuracy in the testing environment than other baselines. Unsupervised approaches like SDIL may struggle to accurately capture the ground-truth hierarchical structure, leading the agent to take irregular actions. Meanwhile, Thought Cloning (TC), although typically better at sub-goal prediction accuracy when aided by LLM-given labels, can also fail to specify the sub-goals of crucial states due to overfitting. This can result in invalid actions and ultimately cause task failure.

Figure: Figure 2: Visualization: Sub-goal selection in an example trajectory instance of Grid-World with 3 Objects. We color-code each sub-goal and black circle marks the final step of each trajectory. The ground-truth is labeled by human in this case, and SEAL achieve the best sub-goal transitions.
Refer to caption: extracted/5897892/Figures/VisualTraj-2.png

Additionally, compared to unsupervised HIL baselines, our SEAL removes the burden for tuning the number of sub-goals $K$, by leveraging concrete suggestions from LLMs. Unsupervised methods like LISA, which rely on a pre-defined codebook of sub-goals, can struggle to match the true task hierarchy when choosing inappropriate hyper-parameter $K$. As shown in Fig. 3, both overestimating and underestimating the number of sub-goals can lead to performance degradation. In contrast, our SEAL avoids this issue and outperforms all unsupervised HIL approaches in most cases.

Figure: (a) #Traj = 200
Refer to caption: extracted/5897892/Figures/Dim-200.png

### 5.4 Performance on longer compositional tasks

**Table 3: Success rates on longer-range compositional tasks (Grid-World) with 3, 4 and 5 objects.**
| Object Num | # Traj | BC | LISA | SDIL | TC | SEAL-L | SEAL |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 3 | 300 | 0.31$\pm$0.04 | 0.44$\pm$0.05 | 0.48$\pm$0.01 | 0.52$\pm$0.07 | 0.65$\pm$0.04 | 0.61$\pm$0.02 |
|  | 400 | 0.48$\pm$0.04 | 0.53$\pm$0.03 | 0.62$\pm$0.04 | 0.62$\pm$0.04 | 0.83$\pm$0.02 | 0.85$\pm$0.02 |
| 4 | 400 | 0.16$\pm$0.03 | 0.13$\pm$0.01 | 0.22$\pm$0.04 | 0.24$\pm$0.05 | 0.26$\pm$0.03 | 0.32$\pm$0.03 |
|  | 500 | 0.39$\pm$0.04 | 0.36$\pm$0.04 | 0.40$\pm$0.01 | 0.39$\pm$0.03 | 0.49$\pm$0.03 | 0.51$\pm$0.03 |
| 5 | 500 | 0.09$\pm$0.02 | 0.11$\pm$0.04 | 0.11$\pm$0.02 | 0.23$\pm$0.06 | 0.25$\pm$0.05 | 0.42$\pm$0.03 |
|  | 600 | 0.30$\pm$0.03 | 0.47$\pm$0.03 | 0.64$\pm$0.05 | 0.35$\pm$0.03 | 0.65$\pm$0.03 | 0.73$\pm$0.03 |

We further investigate whether SEAL can sustain its effectiveness and superiority on longer-range compositional tasks, which involve more sub-goals. To evaluate this, we test our method on Grid-World with 4 and 5 objects, where the LLMs decompose the task instructions into more sub-goals ($K=8$ and $K=10$, respectively). As shown in Table 3, SEAL continues to lead, particularly in cases with smaller expert datasets, demonstrating its adaptability to more complex tasks. In longer-range compositional tasks, managing the increasing complexity of the sub-goal space becomes more challenging, especially for supervised methods like Thought Cloning, as the supervision signal for each sub-goal becomes sparser. SEAL overcomes this by employing a dual-encoder design, which leverages both the flexibility of unsupervised learning to learn sub-goals better. Meanwhile, SEAL focuses more on imitating the few but critical sub-goal transition points in longer-range compositional tasks, avoiding the limitations of signal sparsity faced by other approaches.

### 5.5 Adaptation to Task Variations

We also test the adaptability of SEAL to task variations. To do this, we modify the predefined pick-up order in the Grid-World environment, which includes three objects: A, B, and C. This generates new tasks for evaluation. We create a dataset comprising 400 expert demonstrations for the task with the order ABC, along with few-shot set of 10 expert demonstrations for other orders such as ACB, BCA, and BAC. We then assess the performance of the trained agent on these alternative orders. As shown in Table 4, our method exhibits slightly higher success rates, indicating that SEAL has better adaptability to task variations. However, it is important to note that this conclusion is limited to specific scenarios. In the grid-world, rearranging the order does not introduce new sub-goals, meaning that the sub-goals learned from the training set remain applicable to these new tasks.

**Table 4: Success rates under task variations on Grid-World averaged over 5 random seeds.**
| Test Env | BC | LISA | SDIL | TC | SEAL-L | SEAL |
| --- | --- | --- | --- | --- | --- | --- |
| ABC | 0.48$\pm$0.04 | 0.53$\pm$0.03 | 0.62$\pm$0.04 | 0.62$\pm$0.04 | 0.83$\pm$0.02 | 0.85$\pm$0.02 |
| ACB | 0.01$\pm$0.00 | 0.08$\pm$0.02 | 0.11$\pm$0.04 | 0.13$\pm$0.05 | 0.18$\pm$0.07 | 0.14$\pm$0.03 |
| BAC | 0.01$\pm$0.00 | 0.05$\pm$0.01 | 0.06$\pm$0.02 | 0.09$\pm$0.02 | 0.11$\pm$0.03 | 0.08$\pm$0.02 |
| BCA | 0.00$\pm$0.00 | 0.03$\pm$0.01 | 0.08$\pm$0.03 | 0.08$\pm$0.02 | 0.08$\pm$0.01 | 0.09$\pm$0.03 |

## 6 Conclusion

In this work, we introduce SEAL, a novel HIL framework that leverages LLMs’ semantic and world knowledge to define sub-goal spaces and pre-label states as meaningful sub-goal representations without prior task hierarchy knowledge. SEAL outperforms baselines like BC, LISA, SDIL, and TC across various environments, particularly in low-sample and complex longer-range compositional tasks. Our approach achieves higher success and sub-goal completion rates, with assistance of the dual-encoder proving more robust than the pure LLM encoder and the transition-augmented low-level policy. SEAL also adapts well to varying task complexities and latent dimensions. Our current design still observes training instability, and we are interested in making more efficient SEAL approaches while under partially observed states.

## Appendix A Additional Environment Information

Figure: Figure 4: Examples of compositional-task-related environments used in our experiments. Left: KeyDoor. The player needs to pick up the key and then use it to unlock the door. Right: Grid-World. The player needs to pick up the different objects in a pre-specified order.
Refer to caption: extracted/5897892/Figures/Envs.png

KeyDoor The environment is based on the DoorKey setting from the MiniGrid Dataset, but with modifications to make the state compatible with LLM input for sub-goal mapping. Instead of using image states, we convert the state into an 8-dimensional vector that captures crucial object information: {x-coordinate of key, y-coordinate of key, x-coordinate of door, y-coordinate of door, x-coordinate of player, y-coordinate of player, key status (picked: 1, not picked: 0), and door status (unlocked: 1, locked: 0) }. Wall obstacles are removed to avoid interference. The action space consists of 6 primitive actions: move up, move down, move right, move left, pick up, and unlock. The key can be picked up only when the player reaches the key’s coordinates, and the door can be unlocked only if the player reaches the door’s coordinates with the key already picked up. The language task instruction $\mathcal{M}$ is defined as: *”Pick up the key, and then unlock the door.”* The episode ends when the door is successfully unlocked or the maximum time steps $T=100$ are reached.

Grid-World
The environment is based on the grid world used in (Kipf et al., ; Jiang et al., ). Similar to KeyDoor, the image-based states are converted into a vector format for LLM input, capturing crucial information about objects: x and y coordinates of Object 1, x and y coordinates of Object 2, …, x and y coordinates of the player, status of Object 1 (picked: 1, not picked: 0), status of Object 2, …. For Grid-World with 3, 4, or 5 objects, the state vector has dimensions 11, 14, and 17, respectively. Wall obstacles and irrelevant objects are removed to avoid interference. The action space consists of 5 primitive actions: move up, move down, move right, move left, and pick up. An object can be picked up only when the player reaches its coordinates. The language task instruction $\mathcal{M}$ is defined as: *”Pick up Object 1, then pick up Object 2, then…”* The episode ends when the player picks up all objects in the correct order or after the maximum time step $T=100$. At the start of each episode, the coordinates of all objects and the player are randomly reset.

Sub-goal Spaces Identified by LLMs
We use GPT-4o to decompose the language task instructions for both the KeyDoor and Grid-World environments into their respective sub-goal spaces. In the KeyDoor environment, there are $K=4$ sub-goals: {move to the key, pick up the key, move to the door, unlock the door}. In the Grid-World environment, with 3, 4, and 5 objects, the number of sub-goals is $K=6$, $K=8$, and $K=10$, respectively, including: {move to object 1, pick up object 1, move to object 2, pick up object 2, …}. For both sub-goal spaces, we parameterize each language sub-goal in it by a K-dim one hot vector.

## Appendix B Example Prompts

In SEAL, we prompt LLMs to generate supervisory labels for training the high-level encoder. Fig. 5 illustrates the detailed prompting process. First, we prompt the LLMs to break down the task instruction into a finite set of sub-goals. Then, for each state, the LLM is prompted $K$ times to determine whether it corresponds to each of the decomposed sub-goals, mapping the states to sub-goal representations.
Example prompts for both task decomposition and sub-goal labeling are provided in the following sub-sections.

Figure: Figure 5: A schematic illustrating how LLMs are prompted to define sub-goal spaces from task instructions and map states to sub-goal representations, serving as supervisory labels for training the high-level sub-goal encoder in SEAL.
Refer to caption: extracted/5897892/Figures/Prompt.png

### B.1 Prompts for Task Decomposition

An example prompt used in Grid-World experiment with 3 objects for task decomposition is listed below. We use GPT-4o to produce the answer.

The GPT-4o answer with:

### B.2 Prompts for Mapping states to sub-goal representations

An example prompt used in KeyDoor experiment for mapping states in expert demonstration to sub-goal representations is listed below. We use GPT-4o to produce the answer.

The GPT-4o response with:

## Appendix C Model Implementation Details

We outline the model implementation details for all four baselines and SEAL in the KeyDoor and Grid-World environments. For non-hierarchical BC baselines, we use a two-layer Multi-layer Perceptron (MLP) as the trainable policy agent $\pi(a_{t}|s_{t})$. In HIL approaches like LISA, SDIL, and TC, this same two-layer MLP is used for both the high-level sub-goal encoder $\pi_{H}(s_{t})$ and the low-level policy agent $\pi_{L}(a_{t}|s_{t},z_{t})$. We use Adam (Kingma, ) as the optimizer for all models, with learning rates initialized at 5e-5 for KeyDoor and 5e-6 for Grid-World. To ensure fair comparison, we maintain consistent hyper-parameters across all simulations, including the high-level encoder loss weight $\beta$, the hidden dimensions of the MLPs, and the number of sub-goals $K$ for both HIL baselines and SEAL. Detailed implementations are presented in the following Table 5:

**Table 5: Hyperparameters settings of Model Implementations.**
| Methods | Loss | Hidden dim of MLPs | $\beta$ |
| --- | --- | --- | --- |
| BC | $\mathcal{L}_{BC}=\mathbb{E}_{(s_{t},a_{t})\in\mathcal{D}}-\text{log}\pi(a_{t}|<br>s_{t})$ | [128, 128] | / |
| LISA | $\mathcal{L}_{LISA}=\beta\mathcal{L}_{H}^{(vq)}(s_{t})+\mathcal{L}_{L}(s_{t},z^<br>{(vq)}_{t})$ | [128, 128] | 0.4 |
| SDIL | $\mathcal{L}_{SDIL}=\mathbb{E}_{(s_{t},a_{t})\in\mathcal{D}}\mathbb{E}_{z_{t}<br>\in\pi_{H}(z_{t}|s_{t})}-\text{log}\pi_{L}(a_{t}|z_{t},s_{t})$ | [128, 128] | / |
| TC | $\mathcal{L}_{TC}=\mathbb{E}_{(s_{t},a_{t},z_{t})\in\mathcal{D}}-\text{log}(<br>\beta\pi_{H}(z_{t}|z_{t-1},s_{t})+\pi_{L}(a_{t}|s_{t},z_{t}))$ | [128, 128] | 0.4 |
| SEAL | $\mathcal{L}_{SEAL}=W_{llm}(\mathcal{L}_{H}^{(llm)}+\mathcal{L}_{L}^{(llm)})+W_<br>{vq}(\mathcal{L}_{H}^{(vq)}+\mathcal{L}_{L}^{(vq)})$ | [128, 128] | 0.4 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - How does the learned subgoal representation z_t compare to a representation generated by a LLM? What is the benefit of having a dual encoder compared to simply using open-source LLMs as Llama to represent the subgoals?
- How does this method compare to the rich literature of methods that work on automatic task decomposition and skill library generation (e.g. Voyager)?
- How does the authors propose to do subgoal-state matching when the environment gets more complex?
- How does the authors propose to do task decomposition when the environment is partially observable? Some initially-generated subgoals may not be valid and be changed on the go if the environment can not be fully observed.
- How is the learned subgoal representation actually benefit the low level policy other than assigning higher weights to transition states? Why is there no ablation on this?

Check the weaknesses section for more comments.
soundness: 3
strengths: - Authors study an important problem of representing subgoals from task instructions.
- Automatically does state-sub goal matching from unlabeled trajectory.
- The author's experiments with both supervised subgoal encoder and unsupervised VQ encoder is an interest to the readers, but it is up to debate whether having both is beneficial.
- Experiments are comprehensive.
summary: The authors present a hierarchical imitation learning method that tries to learn a latent subgoal representation.
The subgoal representation is learned from 
1. Generating language subgoals from instructions using LLM prompting.
2. Labeling the state from the expert demonstration with the generated sub-goals using LLM prompting.
3. The labeled states is used to train a supervised subgoal encoder, while the unlabeled state gets used to train a VQ encoder using the subgoal codebook.

The proposed method shows improvement on 2D grid world environment on task success rate.
weaknesses: Author's motivation for the work is quite weak. The authors claim that they want to ask three questions regarding whether pre-trained LLMs can:
1. Serve as a prior to define a task’s hierarchical structure,
2. Establish a sub-goal library autonomously, and
3. Closely guide both high-level sub-goal learning and low-level agent behavior.

There are many existing works (e.g., Voyager, CoT, ToT, LLM-P, etc.) that demonstrate LLMs excel at points 1 and 2, so the main problem is how the authors can make a meaningful contribution to point 3. However, the contributions regarding point 3 are quite limited. The explanation is as follows:

Method:
- The method can be understood as learning an encoder to output a latent representation of a subgoal (which is part of a skill library), a task that existing LLMs already excel at.
- The method appears highly dependent on the quality of LLM-generated state-subgoal matching, which might be manageable in a simple environment and task like those in the paper but is unlikely to transfer well to more complex environments involving multimodal input. The difficulty of such state-subgoal matching has been explored in multiple works (e.g., 1, 2 and more) in robotics.
- The only way the learning of the low-level agent is impacted is by assigning higher weight to states with subgoal transitions. This is a minor contribution.

Experiments:
- The environment and task are quite simple, as the language instructions can be easily decomposed into a limited set of subgoals.
- Although the environment is 2D, it is represented as text, which severely limits the applicability of this method, as the application will be constrained by how the environment can be represented as text.

1. LOTUS: Continual Imitation Learning for Robot Manipulation Through Unsupervised Skill Discovery
2. Skill-based Model-based Reinforcement Learning

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: How did you pick the codebook size for the VQ module? How does that affect the model performance?
soundness: 4
strengths: The strength of this paper lies in its innovative use of Large Language Models (LLMs) for sub-goal generation and labeling within imitation learning, moving beyond the typical application of LLMs as plan generators. By using LLMs to derive semantically meaningful sub-goal commands and annotate demonstration trajectories, the authors introduce a unique way of harnessing LLMs for hierarchical task decomposition. This labeling mechanism directly supports policy training, effectively guiding the learning process in complex, long-horizon tasks without relying on predefined task hierarchies or extensive manual labeling. This approach highlights the potential of LLMs to provide structured supervision in imitation learning.
summary: The SEAL framework (Semantic-Augmented Imitation Learning via Language Model) addresses the challenges of Hierarchical Imitation Learning (HIL) in long-horizon decision-making tasks by leveraging the semantic and reasoning capabilities of Large Language Models (LLMs). SEAL utilizes LLMs to generate meaningful sub-goal representations from language instructions, automatically defining sub-goal spaces without prior task hierarchy knowledge. A dual-encoder structure combines supervised LLM-guided learning with unsupervised vector quantization, improving sub-goal representation reliability. Additionally, SEAL’s transition-augmented low-level planner enhances policy adaptation to sub-goal transitions. SEAL demonstrated improved performance over existing HIL and LLM-based planning methods in tasks requiring complex, long-horizon planning, especially with limited expert datasets.
weaknesses: The paper has several limitations. First, it claims using Vector Quantization (VQ) as part of a key contribution, but there’s no abaltion studies to show if VQ really helps. It would be more convincing if there was an experiment testing VQ against an ablated model that does not have quantization to see if it makes a difference.

Second, the evaluations were done in simple environments with only a few objects (n = 3), which raises concerns. In larger or more complex environments, VQ might not scale well. Also, using LLMs to classify states based on natural language descriptions could cause problems if there are many similar-looking objects (e.g., think about a Sokoban environment where there are many boxes of the same shape/color on the ground) or if the environment has continuous states and actions, like in robotics. Using language as the intermediate representation might not handle these situations "accurately".

Based on the Appendix, the prompts used for SEAL are very specific to the environments tested, which makes it unclear if they would work well in other, different environments without adjustments. Or, how should a user decide on what to include in the prompts?

Finally, if data efficiency is a main advantage of SEAL, it should be compared to other methods for data-efficient RL/IL, like Glide (https://arxiv.org/abs/2403.17124 which also uses LLMs to decompose tasks) and EfficientZero (https://arxiv.org/abs/2111.00210 which focuses on data-efficient model-based RL). Comparing SEAL to these would make it clearer how it performs against other approaches focused on efficient learning.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Why is LLM necessary? What's the benefit of using LLM over abstract task descriptions provided by humans [1]? Do you have quantitative or qualitative results supporting your claim?
2. Assuming LLM is necessary, why is each component of SEAL necessary to learn semantically meaningful subgoals? Can you not use simpler learning objectives [2] to achieve the same goal?
3. What are concrete examples to highlight the drawbacks of using supervised loss or unsupervised loss alone? How to ensure that during training supervised learned goals and unsupervised learned goals are aligned so the weighted average make sense?
4. How do you use LLM to map states to subgoals? While I understand LLM can propose a discrete set of subgoals, but how do you prompt LLM to map states to one of these subgoals? Do you have concrete examples to clarify the mechanism?
5. How do you know where the intermediate states are so you can apply transition adaptive weights? 
6. What are the motivations for using each of the current baselines? What claim does comparison with each of the baselines support?
7. Could you include more complex experiments or domains to show the benefits of your approach more convincingly?
soundness: 2
strengths: The paper works on an important problem of learning semantically meaningful subgoals for demonstrations without dense expert labels. The attempt to leverage LLM knowledge about discrete task structure is well-positioned, and authors took the effort to compare SEAL to many baselines.
summary: This paper tries to leverage LLM to improve sub-goal learning from unsegmented demonstrations for hierarchical imitation learning. The proposed method SEAL combines supervised subgoal learning using LLM labels and unsupervised subgoal learning using vector quantization. The method is tested on two discrete environments KeyDoor and Grid-World with improved results over baselines.
weaknesses: The biggest concern I have for the paper is that the algorithm 1 looks very complex yet the experiments are only two simple discrete grid-world domains. I am not fully convinced why such simple domains require such sophistication in loss term designs while there exist simpler algorithms for more complex domains [1, 2]. Perhaps you should consider citing them, explaining the differences in related works or using them as potential baselines, which might be more relevant than your current baselines. While I understand you have motivated why you need each loss term through high-level description, you should show concrete examples or ablation results to support the claims more clearly. 

[1] Learning Rational Subgoals from Demonstrations and Instructions
[2] Grounding Language Plans in Demonstrations Through Counterfactual Perturbations

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 4
strengths: SEAL effectively leverages LLMs to autonomously define and label sub-goals, enabling the framework to operate without detailed task hierarchy knowledge. I believe this is a crucial step to eliminate the needs to extensively labeling.

The dual-encoder system, combining LLM-based supervision with Vector Quantization, strengthens sub-goal representation and improves model robustness.

SEAL desires several easy but non-trivial tasks to prove it effectiveness. The  transition-augmented low-level planner adapts well to intermediate task transitions, enhancing performance on complex and long-horizon tasks.
summary: SEAL  introduces a novel framework leveraging LLMs to enhance Hierarchical Imitation Learning (HIL) by addressing long-horizon decision-making tasks. 

SEAL utilizes LLMs to generate semantically meaningful sub-goal representations, thus eliminating the need for extensive expert demonstrations and detailed supervisory labels. It employs a dual-encoder structure combining supervised LLM-guided sub-goal learning with unsupervised Vector Quantization for robust sub-goal representations. 

Additionally, SEAL features a transition-augmented low-level planner to improve adaptation to sub-goal transitions. Experimental results indicate that SEAL outperforms existing HIL methods, especially in scenarios with limited expert data and complex tasks.
weaknesses: SEAL heavily relies on the semantic abilities of LLM's quality to generate meaningful sub-goals and guide the hierarchy.

Need more complicated envs. An env with continuous space and action could prove more.

Meta-review: The authors present a hierarchical imitation learning method that tries to learn a latent subgoal representation. The subgoal representation is learned from by (1) generating language subgoals from instructions using LLM prompting. (2) labeling the state from the expert demonstration with the generated sub-goals using LLM prompting. (3) the labeled states is used to train a supervised subgoal encoder, while the unlabeled state gets used to train a VQ encoder using the subgoal codebook. (4) The overall proposed method shows improvement on 2D grid world environment on task success rate.

While the paper studies an important problem the overall motivation of the work as well as the need of using LLM is not well explained.

Decision: Reject