LLM-based Skill Diffusion for Zero-shot Policy Adaptation

Authors: kyung  woo
  kim
youngseok  lee
jooyoung  kim
honguk  woo
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:50:39

Abstract: Recent advances in data-driven imitation learning and offline reinforcement learning have highlighted the use of expert data for skill acquisition and the development of hierarchical policies based on these skills. However, these approaches have not significantly advanced in adapting these skills to unseen contexts, which may involve changing environmental conditions or different user requirements. In this paper, we present a novel LLM-based policy adaptation framework LDuS which leverages an LLM to guide the generation process of a skill diffusion model upon contexts specified in language, facilitating zero-shot skill-based policy adaptation to different contexts. To implement the skill diffusion model, we adapt the loss-guided diffusion with a sequential in-painting technique, where target trajectories are conditioned by masking them with past state-action sequences, thereby enabling the robust and controlled generation of skill trajectories in test-time. To have a loss function for a given context, we employ the LLM-based code generation with iterative refinement, by which the code and controlled trajectory are validated to align with the context in a closed-loop manner. Through experiments, we demonstrate the zero-shot adaptability of LDuS to various context types including different specification levels, multi-modality, and varied temporal conditions for several robotic manipulation tasks, outperforming other language-conditioned imitation and planning methods.

Main paper: 
 introduction: 
 Skill-based learning has demonstrated its potentials in generalizing to novel downstream tasks by leveraging pre-trained skills learned from the offline dataset. Furthermore, the integration of skillbased learning and natural language realizes the remarkable ability to perform practical tasks via the provision of a human-oriented interface where agents are controlled by instructions describing the goals of the task. Building upon this notion, several previous studies have investigated bridging the gap between the human instructions and physical world manipulation by learning semantically meaningful skills given the language-annotated dataset [1, 2, 3] . However, due to the inherently open-ended nature of language, it is impractical to obtain a dataset annotated with a sufficiently wide range of contexts, encompassing various environmental conditions and user requirements, to develop a versatile language-conditioned policy capable of accommodating such diverse contexts. Consequently, as shown in the left side of Figure 1 , these prior works are limited to processing the narrow scope of instructions that primarily convey only the goal of the task without any contextual information (presented as case 1).To address the challenges in language-conditioned skill learning, we explore large language model (LLM)-based policy adaptation approaches that enable zero-shot adaptation to contexts specified in Figure 1 : Zero-shot policy adaptation to contexts: In case 1, the instruction includes only the task goal. In cases 2 and 3, the instruction is supplemented by the task goal with the context. Conventional language-conditioned skill approaches struggle to generate trajectories well aligned with the contexts, and typically succeed only for instructions as in case 1. Conversely, our LLM-based policy adaptation approach effectively adapts to the contexts in a zero-shot manner across all cases. language at test-time. Leveraging the effectiveness of diffusion models in controlling their generation process via loss functions [4] , and the code generation capability of LLMs [5] , we adapt diffusion with guidance by the loss function generated through LLMs. As illustrated in the right side of Figure 1 , our LLM-based policy adaptation is capable of adapting to diverse contexts in a zero-shot manner, such as "target speed of an agent should be 8m/s" or "minimize the power usage of an agent".To this end, we present a novel LLM-based Skill Diffusion (LDuS) framework, designed to facilitate zero-shot adaptation to unseen contexts by generating skill trajectories that are controllable through loss-guided diffusion. Specifically, we devise a hierarchical skill learning structure in which a diffusion model is employed as a skill planner with sequential in-painting. This in-painting method sequentially substitutes consecutive state-action pairs with their originals and learns on the remaining parts, thus allowing for robust trajectory generation by conditioning the trajectory on past experiences. For evaluation, LDuS provides an interface where a context specified in language is translated into a loss function to guide the generation process of the skill diffusion planner. The generation process is then continuously refined through an iterative process using an LLM as a self-critic, ensuring that the generating skill trajectory aligns with the given context. As such, our framework stands apart from existing language-conditioned skill imitation approaches, as it enables the zero-shot adaptation of skill-based policies to various contexts that extend beyond the training dataset.The contributions of our work are summarized as follows.• We present the LDuS framework to address a novel challenge of zero-shot policy adaptation to unseen contexts specified in language.• We develop a hierarchical skill learning structure that adapts the skill diffusion planner with sequential in-painting, enabling robust skill trajectory generation.• We devise an interface that utilizes LLMs to translate a context into a loss function, which is then used to control the generation process of the skill diffusion planner. This is further validated via iterative refinement to adequately align with the given context.• We experimentally show that LDuS achieves superior performance in zero-shot adaptation to a wide range of contexts including different specification levels, multi-modality, and varied temporal conditions for robotic manipulation tasks.2 Related work 
language-conditioned skill learning: 
 In the domain of sequential decision-making, several researches have explored techniques for learning language-conditioned skills [1, 2, 3, 6, 7, 8] . LISA [2] employs hierarchical skill learning to achieve a language-conditioned policy through discretized skill codes. Recently, LCD [3] , PlayFusion [6] ,and SkillDU [8] commonly adopt diffusion models as a language-conditioned policy to address highdimensional vision inputs or to leverage play datasets collected by human. While these approaches primarily concentrate on learning by direct supervision from language-annotated datasets, our LDuS aims at adapting to unseen contexts that convey varying environment conditions or different user requirements beyond the datasets.Meanwhile, several studies have harnessed the code generation capabilities of LLMs to ground language instructions to actionable skills [9, 10, 11, 12, 13, 14] . For example, in Kinematic-LLM [14] , an LLM is prompted with in-context samples to generate the waypoints for robotic manipulation with pre-defined code primitives. The performance of this in-context learning often depends on the quality and relevance of selected samples, leading to limited generalization to unseen contexts. Unlike this waypoint-based high-level planning with pre-defined code primitives, our LDuS adapts the code generation capabilities of LLMs with iterative refinement to enable the fine-grained control of trajectories, particularly suited for adjusting the generation process of diffusion models.Diffusion models have shown promising results in various areas including computer vision [15, 16, 17, 18] , offline reinforcement learning (RL) [19, 20, 21] , and long-horizon planning [22, 23, 24, 25] . The strong generation ability of diffusion models leads to robust adaptability at deployment. In classifier guidance [26] , diffusion models are controlled at test-time, thereby supporting the generation of images belonging to specific classes. This controlled adaptation concept has been further investigated for text-driven image generation [27, 28] and noisy inverse problems [29, 30] . Recently, LGD [4] presents a loss-guided diffusion mechanism, by which diffusion models can be controlled via differentiable loss functions without additional training. To facilitate the querycompliant scene generation, CTG [5] leverages LLMs as a loss function generator for user queries.In the RL domain, such guidance schemes for controlling diffusion-based policies have been investigated with pre-trained dynamics models [31] and value functions [22, 23] . Yet, these schemes rarely accommodate language-specified contexts. Our LDuS is the first to integrate the reasoning capabilities of LLMs and the controlled generation capabilities of diffusion models, thus enabling zero-shot policy adaptation to language-specified contexts in the domain of sequential decision-making. 
problem formulation: 
 Contextual Markov Decision Process (MDP). We consider a task as a contextual MDP [32, 3] (C, S, A, G, P c , r c , γ, ρ 0 ) where c ∈ C is a context space, s ∈ S is a state space, a ∈ A is an action space, g ∈ G is a goal space, P c : C × S × A × S → [0, 1] is a transition probability conditioned on the context and the goal,r c : C × S × A × G → R is a reward function conditioned on the context, γ ∈ [0, 1] is a discount factor, and ρ 0 : S → [0, 1] is an initial state distribution. Here, we consider a goal to be specified in language, denoted as g l , such as "open the drawer" or "close the door". Moreover, we assume a context is provided in language [3] , denoted as c l , which can affect either the reward function or the transition probability.Policy adaptation to contexts. We assume access to a dataset D = {τ i } i≤N , where each trajectory τ i is represented as a sequence of state and action pairs with a goal {(s t , a t , g l )} t≤T for a T -length episode without any contextual information.We consider task evaluation scenarios with a context c l which conveys environmental conditions or user requirements, along with a goal g l . Then, our objective is to develop a policy adaptation framework ϕ(g l , c l ) that maps both goal g l and context c l to a policy π c maximizing the return of context-conditioned rewards.EQUATIONFigure 2 : Concept of LDuS with skill learning and adaptation phasesDiffusion models have been explored for task planning [22, 23, 24, 33] , offline RL [19, 20, 21] , and language-conditioned skill learning [3, 6, 8] . In [22] , a diffusion model for planning was introduced, which generates h-length state and action sequences, denoted as x in a two-dimensional array.EQUATIONIn particular, the diffusion model ϵ(x k , k) based on the U-net architecture [34] predicts a noise η ∼ N (0, I) given a noise-corrupted trajectory x k as an input [15, 16] , and it is optimized through the following loss.EQUATIONHere, k ∈ [1, K] is a denoising step, andx k = √ α k x + √ 1 -α k ηis generated by adding a Gaussian noise η to the original trajectory x with a variance schedule parameter α k . At sampling, the diffusion model generates trajectory x from a random noisy input x K ∼ N (0, I) by sequentially denoising it,EQUATIONwhere σ k is a parameter for a variance schedule.4 Our approach 
overall framework: 
 To address the zero-shot adaptation to various contexts, we develop the LDuS framework comprising two phases: (i) skill learning via the skill diffusion planner, and (ii) policy adaptation to unseen contexts via LLM-guided diffusion, as illustrated in Figure 2 . (i) In the learning phase, we establish a hierarchical structure in which skills are learned, conditioned on a goal. This structure includes a skill encoder, a skill prior, and a skill diffusion planner that generates skill trajectories. These components are learned on the dataset that contains only goals, without any contextual information. Additionally, we employ a sequential in-painting technique when training the skill diffusion planner to enhance robustness in skill trajectory generation. (ii) In the adaptation phase, we guide the generation process of the skill diffusion planner upon a language-specified context, by harnessing the code generation capabilities of LLMs. The generated code serves as a loss function that guides the skill diffusion planner to generate skill trajectories at every denoising step. This facilitates alignment between the skill trajectories and the given context. Furthermore, LDuS employs an iterative refinement process, in which generated skill trajectories are repeatedly validated in a closed-loop manner to achieve robust alignment with the context.To facilitate skill learning using a diffusion model, we adopt a variational autoencoder (VAE) [35] architecture with three components: a skill encoder q(z|s t:t+h , a t:t+h ), a skill prior p l (z|s t , g l ), and a skill diffusion planner ϵ(x k , k, z). Given an h-length skill trajectory x = {s t , a t } t≤h , the skill encoder q predicts a skill embedding z, then the skill diffusion planner reconstructs the h-length skill trajectory x based on z. To optimize both the skill encoder and skill diffusion planner, we employ a conditional VAE objective that combines a diffusion reconstruction term in (3) and a prior regularization term such as min q,ϵ E x∼D,z∼q,k∼ [1,K] ||ϵ(x k , k, z) -η|| 2 2 + βD KL (q(z|s t:t+h , a t:t+h ), p(z))Figure 3 : LDuS framework: In (i), given the dataset annotated with the goals, skills are learned through the hierarchical structure employing a skill diffusion planner with sequential in-painting techniques. In (ii), the context specified in language is translated into a loss function, which is then used to guide the generation process of the skill diffusion planner. This process is further validated with closed-loop iterative refinement to better align skill trajectories with the given context.where p(z) is a unit Gaussian N (0, I), D KL is the Kullback-Leibler (KL) divergence, and β is a weight for regularization [36] . To establish a versatile skill embedding space encapsulating common skills across multiple tasks, we use the skill encoder without goal conditions. In addition, to learn skills conditioned on the goal, the skill prior p l , which is conditioned on state s t and goal g l , is tailored to align with the output of the skill encoder. To handle the goal provided in language, we also use a pre-trained language encoder Φ l such as CLIP [37] that produces language embeddings of g l for the skill prior. Then, the skill prior is jointly trained with the skill encoder and skill diffusion planner by minimizing the distance with the skill encoder, i.e., min p l D KL (q(z|s t:t+h , a t:t+h ), p l (z|s t , Φ l (g l ))) .Using the skill prior, we predict an appropriate skill embedding at deployment based on the current state and specified goal.Sequential in-painting. The in-painting technique is adopted in [22] to address goal-conditioned problems, in which a diffusion model is conditioned by replacing the last state of a generating trajectory x k with the goal state. We adapt this in-painting to train the skill diffusion planner with a sequential replacement mechanism, where a sequence of m ∼ [1, h] states and actions of x k is substituted with the corresponding original state-action pairs to learn the remaining portion of x k , as illustrated in the bottom left side of Figure 3 . Then, for evaluation, unlike conventional methods that constrain the diffusion model only with the current state [22, 23] , our approach constrains the skill diffusion planner using the previously encountered m states. This in-painting method enables the diffusion planner to generate more robust and contextually aligned skill trajectories, as described in Section 5.3.In general, it is challenging for a model to directly acquire the zero-shot adaptation ability for various contexts, particularly when the training dataset has limited coverage on those contexts. To tackle this challenge, we harness the controlled generation capabilities of diffusion models [26] along with the code-generation capabilities of LLMs [5] .In LDuS, for zero-shot adaptation to various contexts, we employ three procedures: translation of a given context to its corresponding loss function, loss-guided trajectory generation, and closed-loop iterative refinement, as illustrated in the right side of Figure 3 . Initially, the LLM is tasked with translating the context into a loss function. This loss function then guides the skill diffusion planner in its trajectory generation process. To ensure the accuracy and relevance of generated skill trajectories, the process is validated by the LLM in a closed-loop manner, querying the LLM to verify whether the generated trajectories meet the specifications of the given context.1: Inputs: skill prior p l (z|s t , g l ), skill diffusion planner ϵ(x k , k, z), goal g l , context c l , LLM Φ LLM , guidance weight δ 2: Obtain loss function ℓ(x 0 ) using Φ LLM through (7) 3: for every environment step t do 4:z ∼ p l (z|s t , g l ) 5:while not validate do 6:Sample trajectory x 0 using ϵ(x k , k, z) without guidance through (4) 7:Sample guided trajectory x0 using ℓ(x 0 ) and ϵ(x k , k, z) through (8) 8:Validate whether satisfy the context via LLM as Φ LLM (x 0 , x0 , g l , c l )9:Execute a m in x 0 to the environment Translation of context to loss function. For converting the context into a loss function in code, we employ chain-of-thought prompting [38] with a pre-defined list of queries u = [u 1 , .., u n ]. The queries are designed to capture the specifications of the agent and the desired format of loss functions. Then, the queries are sequentially prompted to the LLM Φ LLM in conjunction with a goal g l and a context c l , i.e., {y j |y j = Φ LLM (g l , c l , {u i , y i } i<j )} (7) where {u i , y i } i<j represents a set of prompted queries u i and their respective responses y i from the LLM. The final response y n is then used as the loss function to guide the skill diffusion planner.Loss-guided skill trajectory generation. Similar to prior work [29, 4] , we implement the loss-guided skill trajectory generation, where the guidance is computed as the gradient of the loss function with respect to x k at each denoising step k.EQUATIONHere, x0 = 1 √ ᾱk (x k + (1 -ᾱk )ϵ(x k , k, z)) is an approximation of x 0 given x k [29], and δ is a hyperparameter to modulate the strength of the guidance.Closed-loop iterative refinement. In the closed-loop iterative refinement, we employ the LLM as a self-critic to evaluate both the loss functions and generated trajectories. Specifically, we prompt the LLM with unguided trajectory x 0 , guided trajectory x0 , goal g l , and context c l , i.e., Φ LLM (x 0 , x0 , g l , c l ). The LLM then checks for errors. If errors are detected in the loss function, it is regenerated. If there is a mismatch between the trajectory and the context, the frequency of the guidance application is increased. This ensures continuous improvement in the accuracy and relevance of trajectories generated by the skill diffusion planner. The process of zero-shot policy adaptation is summarized in Algorithm 1. 
experiment settings: 
 Datasets. We use the MetaWorld benchmark [39] , specifically with 10 different robot manipulation goals. We also utilize long-horizon goals from the multi-stage MetaWorld, where each goal comprises a sequence of short-horizon manipulation sub-goals. For data collection, we emulate rule-based expert policies. For each goal, we collect 60 trajectories, varying the speed of the agent as well as the position and weight of the objects being manipulated.Contexts. We use two context groups: (i) language context where the context is solely specified in language, (ii) multi-modal context where additional information is provided through image input to assist in resolving the given context. These contexts are used to direct the agent with instructions such as moving below or above a specific speed, adjusting its speed faster or slower along a specified axis, or exerting more or less force on heavy objects. The context details are in Appendix A.4.Evaluation metrics. We use two metrics to assess the zero-shot performance of LDuS and the baselines. Success Rate (SR) quantifies the percentage of goals or sub-goals that are successfully completed. Context Reward (CR) evaluates the average reward achieved based on how effectively the models satisfy the given context. Baselines. For comparison, we use several language-conditioned imitation and planning methods. 1) LangDT [40] is an imitation learning method that utilizes a language-conditioned decision transformer, 2) LISA [2] is a hierarchical skill imitation framework that learns discredited skill codes conditioned on language instructions, 3) LCD [3] is a hierarchical planning framework that reconstructs state sequences using a diffusion model conditioned on a language input, 4) Diffuser [22] is a task planning framework based on diffusion models.For LLMs, we use GPT-3.5 [41] which is capable of generating loss functions in the form of executable code. Moreover, for multi-modal contexts, we use GPT-4 [42] . Since these baselines rarely account for zero-shot adaptation to contexts, we adopt the same diffusion guidance used in LDuS for those (i.e., Diffuser and LCD) employing diffusion models. In the cases where the diffusion model is guided by a hand-designed loss function, which is considered optimal, we specify such baselines with the additional label of Guidance.Zero-shot performance. In these experiments, LangDT and LISA exhibit the lowest performance, even for the cases without contexts. This is attributed to the multi-modality in the dataset, which tends to hinder the learning of policies built with multi-layer perceptrons or transformers. In contrast, the baselines employing diffusion models, such as Diffuser and LCD, show improved performance. However, none of these Various context types. In Table 2 , we evaluate the performance across several context types such as precise, abstract, and temporal contexts, while all given contexts are specified in language. Specifically, precise contexts include detailed user requirements, such as a specific target speed, e.g., "the agent speed should move between 5m/s and 6m/s." In contrast, abstract contexts lack specific details. For instance, if the user wants the agent to increase its speed, the abstract context could be phrased as "I am very busy; the agent needs to hurry." Temporal contexts are dynamic and vary over time, which are particularly relevant in long-horizon goals. As shown, the results indicate that LDuS significantly enhances CR, with an increase of 43.7% at average compared to Diffuser+Guidance.Comparing with waypoint generation. In addition to the learning-based baselines compared previously, we compare our approach with Kinematic-LLM [14] by which waypoints for pre-defined skill primitives are generated through an LLM with in-context samples. To implement Kinematic-LLM, we define basic skill primitives such as move, push, and pull, and use the same samples for prompting, which are used for LDuS. As shown in Table 3 , Kinematic-LLM shows comparable performance in SR for MetaWorld, but lags in multi-stage MetaWorld. This is attributed to the increased complexity of planning with LLMs for long-horizon goals. Regarding CR, Kinematic-LLM consistently demonstrates lower performance compared to LDuS. This result stems from a lack of versatile samples and skill primitives that are necessary to effectively adapt to various contexts. Inference Time. In Table 4 , we present the average inference time (in milliseconds) required per timestep for LDuS and the baselines. The measurements are conducted on a system equipped with an Intel(R) Core(TM) i9-10980XE CPU and an NVIDIA RTX A6000 GPU, and we use GPT-3.5 for the LLM. As LDuS requires both diffusion sampling time and LLM inference, we measure these component separately, denoted as "Diffusion" and "LLM" in the parenthesis. Diffuser and LCD exhibit the shortest inference time, as these baselines do not require loss guided sampling or LLM inference. When considering only the diffusion sampling time excluding LLM inference, LDuS demonstrates an inference time comparable to the baselines that use the loss guidance. However, the full inference time of LDuS is longer due to its LLM-based code generation and iterative refinement process. This overhead can be mitigated by using a smaller language model, which can be obtained by distilling the essential knowledge required for LDuS. We collect successful skill trajectories from LDuS, using different contexts and varying the guidance weight δ ranging from 0.05 to 0.4. This weight setting regulates the strength of gradient application, as described in (8) . As observed, the embeddings are expanded via LDuS, specifying that LDuS is capable of rendering novel skill trajectories, which are not presented in the dataset but necessary to adapt to different contexts. This demonstrates the versatility of LDuS that arises from the generation capabilities of the diffusion model.LLM-guided skill diffusion planner. Table 5 shows the impact of our LLM-guided skill diffusion planner. In this ablation study with MetaWorld, we evaluate two variants of LDuS, one without loss guidance and the other without iterative refinement. LDuS achieves improved CR of 37.8% over the variant without the loss guidance for language contexts. It also achieves improved CR of 46.8% on average over the variant without iterative refinement. These results underscore the effectiveness of LDuS with loss guidance and iterative refinement for zero-shot adaptation. Sequential in-painting. Figure 5 shows the effect of our sequential in-painting technique. With multi-stage MetaWorld, the hatched bars denote the performance of LDuS and Diffuser without sequential in-painting, while the solid-colored bars indicate that of these models with sequential inpainting. SR and CR are significantly improved for both LDuS and Diffuser, when the in-painting technique is applied. This is attributed to conditioning on the past m experiences, which function as long-term memory. This memory feature enhances performance, particularly for long-horizon goals like multi-stage MetaWorld. Furthermore, the performance enhancement is more pronounced for LDuS, as it benefits from learning common skills from the dataset containing trajectories of multiple goals. 
conclusion and limitations: 
 In this work, we presented the LDuS framework for zero-shot skill-based policy adaptation to contexts specified in language. The framework employs a hierarchical structure for skill learning, in which the skill encoder learns task-agnostic skill abstractions and the skill diffusion planner generates various skill trajectories. The skill diffusion planner is enhanced with sequential in-painting, thus enabling context-aligned trajectory generation for the skills. At test-time, given a specific context describing environmental conditions or user requirements, LDuS directly influences the generation process of the skill diffusion planner, allowing for skill-based policies to adapt to the context. This zeroshot adaptation is achieved by a combination of LLM-based loss function generation and iterative refinement, along with the controllable structure of the skill diffusion planner. LDuS stands apart from other language-conditioned approaches, which are limited to certain variations of the instructions present in the dataset and generalize insufficiently to a range of unseen contexts.Limitations. LDuS has several limitations which direct us to future work. One limitation is related to the inference time, as LDuS relies on iterative LLM inferences for refinement. This issue could be mitigated by distilling only essential knowledge, such as code generation and verification capabilities of an LLM, into a smaller language model. Furthermore, since LDuS relies on the LLM for multiple components, including code generation and iterative refinement, another limitation is its robustness, which can be affected by the variability in the LLM's performance and the design of the prompts. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: In this section, we provide detailed information about our environment settings, dataset collection strategy, and context configurations used for zero-shot evaluation.For evaluation, we use MetaWorld [39] , where an agent is tasked with manipulating an object to achieve a given goal. We For zero-shot evaluation, we utilize all manipulating goals for language contexts and select two manipulating goals for multi-modal contexts.To evaluate long-horizon secenarios, we modify MetaWorld to configure the long-horizon goals. Each goal in the multi-stage Metaworld consists of four existing MetaWorld sub-goals including slide puck, close drawer, push button and insert peg. The agent is then tasked to complete these sub-goals in a specified order. We use 3 different long-horizon goals in this multi-stage MetaWorld, each with a unique sub-task completion sequence. The goal description for multi-stage MetaWorld is formed by concatenating the descriptions of the sub-goals in the specified order, such as "close drawer and insert peg and push button and slide puck". The right side of Figure 6 illustrates an example of long-horizon goal used in the multi-stage MetaWorld. For zero-shot evaluation, we utilize all three long-horizon goals for language contexts and one long-horizon goal for multi-modal context. We configure various contexts specified in natural language to evaluate the zero-shot adaptability of LDuS and the baselines. These contexts are defined based on different specification levels (abstract and precise), multi-modality, and varied temporal conditions. To assess whether LDuS and the baselines align with the given context, we manually design context-specific reward functions for each scenario.For precise contexts, a concrete user requirement is given with a numerical value. We define four different precise contexts as follows:• The agent should move at a speed slower than x m/s• The agent should move at a speed greater than x m/s• The agent should move at a speed faster than x m/s but slower than x m/s• The agent should move at speed along the y-axis a minimum of x m/s.Here, x is filled with different numerical values, depending on the specific goal and the context. To configure the zero-shot adaptation settings, we test the model with speeds different from those used in the training dataset.For abstract contexts, a user requirement is conveyed with some degree of ambiguity. We define two different abstract contexts as follows:• I want to be relaxed, but the agent is too distracting.• I'm very busy, the agent should hurry up.For multimodal context, the agent is required to exert greater force only when in context with the manipulating object, aiming to minimize energy consumption of the end executor. In this scenario, the agent utilizes image input to determine whether it is in context with the manipulating object.• The agent should exert a greater speed of x m/s only when in contact with the object.For temporally varying context, the context changes over time to one of the precise contexts.In this section, we provide the implementation details of LDuS and other baselines, along with the hyperparameter settings used for training. All experiments are conduced on a system equipped with an Intel(R) Core(TM) i9-10980XE CPU and an NVIDIA RTX A6000 GPU.We implement LangDT based on Decision Transformer [40] . LangDT consists of a single decision transformer, built on GPT-2 [43] , conditioned by the goal specified in language. Since LangDT is trained solely by conditioning on goals, which are the only language annotations in the dataset, it struggles with unseen contexts. Therefore, for zero-shot evaluation, we condition the model solely on the given goal. The hyperparameter settings for LangDT are summarized in Table 6 Note that short refers to MetaWorld, while long refers to multi-stage MetaWorld. We implement LISA [2] using the open source project2 . LISA consists of a skill predictor and a policy, both implemented with causal transformers. The skill predictor generate discrete skill codes given a goal and a current state. Then, the policy generates a h-length skill based on the predicted skill code. Similar to LangDT, LCD is evaluated by conditioning the model solely on the goal. The hyperparameter settings for LISA are summarized in Table 7 . We implement LCD [3] using the open source projects Jax3 and Haiku4 . LCD consists of a diffusion model, which generates a state trajectory given a current state and a goal, and an inverse dynamics mode, which reconstructs actions given a pair of states. For zero-shot adaptation, we implement LCD controlled with a loss function, denoted as LCD+Guidance in the main manuscript. The loss functions are hand-designed for each context, and loss guidance is applied at each denoising timestep of the diffusion model, following the same approach as in the loss-guided trajectory generation process of LDuS. The hyperparameter settings for LCD are summarized in Table 8 . We implement Diffuser [22] using the open source projects Jax and Haiku. Diffuser is butile with a diffusion model that generates a trajectory given a current state and a goal. Similar to LCD, we implement Diffuser controlled with a manually designed loss function, which is denoted as Diffuser+Guidance in the main manuscript. The hyperparameter settings for Diffuser are summarized in Table 9 .We implement Kinematic-LLM [14] using GPT-3.5 [41] . Kinematic-LLM consists of a kinematic knowledge parser and a kinematic-aware planner. The kinematic knowledge parser generates an XML file that describes the current state of the environment, including details about the kinematic properties of both the manipulating object and the agent. Following this, the kinematic-aware planner generates an abstract manipulation sequence along with 3D manipulation waypoints. The abstract manipulation sequence provides the description of the procedure, while the 3D manipulation waypoints give the numerical values of the skill primitives in a structured format. The LLM is prompted with the generated XML file containing current kinematic information, a goal description, the available skill primitive, and sufficient in-context samples. For skill primitives, we use the 6 different skills defined for rule-based expert policy emulation, as described in Section A.3. To evaluate zero-shot performance on unseen contexts, the language context is concatenated to the goal description. Below, we provide the prompt used for Kinematic-LLM. Listing 1: The prompt used for Kinematic-LLM# # In -context Sample 2 ... # # In -context Sample N ... --------------------------------------------------------------------- # ItWe implement LDuS using the open source projects Jax and Haiku. LDuS consists of a skill encoder based on an LSTM, a skill prior based on MLPs, and a skill diffusion planner. At test-time, the skill prior is used to predict the skill embedding given a goal and a current state. Then, the skill diffusion planner generates skill trajectory based on the skill embedding. To embed goals specified in language, we utilize CLIP [37] . The hyperparameter settings for LDuS are summarized in Table 10 . Translation of context to loss function. We manually design queries that describe the specifications of the agent and the desired format of the loss function. These queries are used in chain-of-thought (CoT) prompting to guide the LLM in generating a loss function form a given context. We utilize GPT-3.5 [41] . Below, we present the prompt used to translate contexts to loss functions for LDuS.Query 1: " The shape of the action sequence of the agent is (B , H , 4) where B is the batch size , H is the number of sequences , and 4 represents the (x , y , z , grab on / off ) in corresponding agent coordinate of the agent . Explain the configuration the agent 's action " Query 2: " The agent should satisfy the given user requirement . The user requirement is given as follows : [ context ] . Note that the speed of the agent is determined by the L2 norm of the actions . In order to satisfy the user requirement , what should be considered ? " Query 3: " Now , generate a loss function that guides the generating trajectory to satisfy the given user requirement . Generated Python based loss function should follow the following format : ' def _loss_fn (x , obs_dim ) : act = [ x : ,: , obs_dim ] return loss ', where act is a numpy array representing the action sequences . "Listing 2: The prompt used for generating loss function Then, the generated loss function is used to guide the generation process of our skill diffusion planner. Below, we present an example of code generated by the LLM. Loss-guided skill trajectory generation. We implement the loss-guided skill trajectory generation based on previous studies [4, 29] . In practice, instead of predicting noise from a noise-corrupted trajectory x k , our skill diffusion planner reconstructs the original trajectory x 0 to ensure robust generation. Thus, we do not need to explicitly calculate x0 from x k , and the loss-guidance is directly applied to the output of the skill diffusion planner. The frequency of guidance applications, denoted as n, is determined by an iterative refinement procedure, where n is gradually increasedfrom 1 to a pre-defined maximum value until the generated trajectory aligns with the given context. The procedure of loss-guided skill trajectory generation is summarized in Algorithm 2.Closed-loop iterative refinement. We use an LLM to detect errors in the loss function and identify any mismatch between the generated trajectory and the given context. Below, we present the prompt used to regenerate the loss function when LDuS detects such errorsn.Query : " The previously generated code has an error . Regenerate the Python code . Note that the user requirement was [ context ] . "Listing 4: The prompt used to detect error in the loss function Below, we present the prompt use to detect a mismatch. If a mismatch is found between the generated trajectory with guidance and the context, the frequency of the guidance application is increased. This validation process is conducted every h steps using the LLM.Algorithm 2 Loss-guided skill trajectory generation 1: Inputs: skill diffusion planner ϵ(x k , k, z), loss function ℓ, total denoising timestep K, guidance weight δ, frequency for guidance application n 2: x K ∼ N (0, I) 3: for k = K -1, ..., 0 do 4:x k s0:sm ← s 0:m 5:EQUATIONfor i = 1, ...n, do 7:ε ← ε -δ∇ x k ℓ(x 0 = ε) 8:η ∼ N (0, I)9: Listing 5: The prompt used to detect a mismatch between generated trajectory and the context Multi-modal context. For a multi-modal context, LDuS needs to determine whether the agent is in contact with the manipulating object from an image input. We utilize GPT-4 [42] , as it can concurrently process both image and text inputs. If the LLM detects that the agent is in contact with the object, the loss-guidance is applied. Initially, we generate a caption describing the image using the LLM. Then, we prompt the LLM to assess whether the agent is in contact with the object based on the caption. To further enhance the reasoning capability of the LLM, we prompt the LLM with the previous image and the corresponding answer to better understand the context of the current image. Below, we present the prompt used for a multi-modal context.x k-1 ← √ α k (1-ᾱk-1 ) 1-ᾱk x k + √ α k-1 β k 1-ᾱk ε + σ k η 10: return x 0Caption : " In the image , the red robot arm is attempting to close the green drawer with white drawer handle . The robot 's arm has a cylindrical blue and white tool attached to its end , which is called as an end effector . The background is an indoor setting with a wooden floor and a gray wall . " Table 11 and 12 shows the detailed experiment results in MetaWorld and multi-stage MetaWorld, respectively, under the conditions where no contexts are provided. As shown, LDuS consistently achieves the best performance in SR for all goals. The low performance of LangDT and LISA is attributed to the multi-modal nature of the dataset, which hinders the learning of policies built with MLPs or transformers. LCD and Diffuser mitigate this issue to some extent by employing the diffusion models, but their performance is still lower compared to LDuS. This performance gap is more pronounced in multi-stage MetaWorld. This is because LDuS leverages as sequential in-painting method, which contributes to more robust trajectory generation. as they are the most comparable baselines. As shown, LDuS consistently delivers robust zero-shot performance in both CR and SR across all configurations. In multi-stage MetaWorld, the performance fluctuates in the baselines, due to their limited ability to generate plans for long-horizon goals. In Table 15 , we present additional experiments involving two different context types: energy constraints and spatial limitations. We conduct these experiment on a single task in MetaWorld. For the energy context, the agent aims to minimize its energy consumption by reducing acceleration or deacceleration [44] . For the spatial context, the agent is tasked with stying within a specified spatial boundary without crossing it. As shown, LDuS outperforms the baselines in CR SR, demonstrating its scalability cross diverse context types.However, for complex tasks such as dexterous control, interactions with the environment [45] or AP functions [46] are necessary. This is because LLMs are not inherently grounded in these complex environments, making it challenging for them to directly generate directly loss functions. We believe that incorporating such approaches will enable LDuS to accommodate more complex tasks, which we plan to explore as part of our future research directions. Justification: We include the details of our computer resources in Appendix B.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There is no societal impact of the work performed in our workGuidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).https://github.com/Div99/LISAhttps://github.com/google/jaxhttps://github.com/google-deepmind/dm-haiku 


Decision: This paper presents an LLM-based policy adaptation framework for language-specified context in robotic tasks. 
It combines a skill-based diffusion policy learned with in-painting, All reviewers agree that this is an innovative integration of LLM reasoning and loss-guided diffusion, and the paper demonstrates robustness across context types, with a clear presentation. The approach shows good performance  in skill-based learning task with context-specific instructions, particularly for speed variations in robotic manipulation tasks. Although, limitations exist in prompt dependency. 

However the rebuttal did a good job of providing clarifications with more than one reviewers updating the scores leaning more positively.  All reviewers lean accept after rebuttal phase, and the Metra-reviewer concurs on the overall agreement. . 
The authors are advised to review the final comments, and update the manuscript accordingly.. In addition to including clarifications in the main paper, the authors are also advised to describe the possible limitations of this work in more detail.

Reviews: review: paper_summary: paper_summary: This paper presents an LLM-based policy adaptation framework for a language specified context. Here, context is a slight variation in how the task is performed. It has two stages, in skill-learning phase, a skill-based diffusion policy is learned with in-painting technique. Here, skill is a VAE encoding of a sequence of states-actions. In adaptation phase, the generation process of diffusion policy is guided with loss-guided diffusion (LGD) mechanism, where loss is generated by LLM conditioned on the context. Experimental results generally show that their skill-based policy along with LLM-generated-LGD is better at context following that other baselines.
strength_weakness: strength_weakness: 1. The scope of "context" is only limited to the agent's speed in the experiments. Meaning, the only type of adaptation/user-requirement studied in the paper is speed variation. Generating loss function for this is an easy task for an LLM given we have already seen objective generation on much harder tasks in works like [Eureka] and [Lang-to-Rew]. Perhaps experimenting with other types of context would have offered better insight into scalability of such approaches. 
2. The methods has only been shown to work with low-dimensional state and not with images which raises question on it's real-world applicability. Although for given contexts, proposed framework might work with image-action data (as actions are sufficient to define trajectory for loss generation by LLM) but proof of that is missing. Perhaps an ablation where no low-dimensional state is used elsewhere except in skill-encoder and skill-prior, with only actions in trajectory, could have been performed to study the reliance of performance on state-info in both stages.

[Eureka]: https://arxiv.org/abs/2310.12931 
[Lang-to-Rew]: https://arxiv.org/abs/2306.08647
questions: questions: 1. Line 197-198: How's error in the loss function identified if true loss function is not known? Line 638 mentions "if LDuS detects errors in the loss function", how is this detected? If this step is not robust then how will it be sure if mismatch is due to loss function or due to guidance weight?
2. Line 207: If speed variation is already captured in the training data then how is the evaluation ensured too be zero-shot? Details about the range of speed values during data collection and evaluation should be mentioned. t-SNE figure does show OOD generation but does not imply if all evaluation instances were OOD.
3. LDuS's no-iterative-refinement version outperforms diffuser suggesting skill-based diffusion is contributing a lot to the performance gain. In this regard, [Skill-diffuser] is a more relevant baseline to compare as it also conditions the action diffuser on latent skill computed using lang and image obs. It would have been a more fairer comparison in skill-based category as compared to LCD, as LCD is diffusion in latent skill space with RNN action head. Do you agree with this? If no, why? If yes, why wasn't this baseline compared?

[Skill-diffuser]: https://arxiv.org/abs/2312.11598
limitations: limitations: Image-based policy experiments are important to ensure real-world applicability of robotic methods and their abscence should have been addressed.
paper_summary: paper_summary: LDuS is a diffusion-based approach for offline skill learning that adopts several advances to improve performance in goal-driven settings. The main contribution is the adoption of LLM-based guidance, allowing to comply with contextual information/conditions, while achieving the given goal.
strength_weakness: strength_weakness: The work presents some weaknesses, other than the limitations presented by the authors:
* **Problem**: the main problem addressed by this work (how to "stylize behaviours" given context information) is interesting but it's somehow narrow. Nonetheless, the method seems to perform better also in settings where context is not provided 
* **Prompting**: as many works employing LLMs, the way the LLM is prompted seems to be crucial for correctly guiding the diffusion process. This means that changing the LLM, a new specific prompt structure would be required
questions: questions: * what is the current inference time and how does it compare with the other approaches?
* line 461 multimodal, should be multi-modal to align with the rest of the text
* typo "Planing", in Appendix
limitations: limitations: Limitations are described in the paper, though the authors may consider adding some of the Weaknesses described above (e.g. prompting).
paper_summary: paper_summary: This paper presents LDuS, a framework that adapts skill diffusion models to unseen contexts in a zero-shot manner. The proposed method first couples VAE with a diffusion planner for hierarchical skill learning. To perform zero-shot policy adaptation to the language-specified context, LDuS translates contexts to loss functions using LLM and leverages loss-guided diffusion for controllable trajectory generation. Furthermore, LDuS adopts LLM to iteratively evaluate and improve the alignment between the generated trajectory and the given context. The method is extensively evaluated in two zero-shot settings of MetaWorld, and achieves superior performance in both success rate and context reward compared to the baselines.
strength_weakness: strength_weakness: - Tested contexts seem to be mainly limited in speed specifications.
- More seeds should be tested in experiments to draw more convincing and robust conclusions.
- The application of loss guidance may harm the model's capability of completing the goal/task (i.e. success rate).
questions: questions: - Would sequential in-painting increase the training cost (e.g. training time) of the learning phase?
- How is the context reward defined specifically for each context?
- In line 228, how is the loss function manually designed for the "Guidance" setting?
- How many times of refinement are usually needed before a generated trajectory is considered contextually aligned? Would that be too time-consuming for long-horizon tasks?
- How sensitive is LDuS to loss guidance strength?
limitations: limitations: The authors adequately discussed the limitations of the proposed method.
paper_summary: paper_summary: This paper presents a novel framework called LLM-based Skill Diffusion (LDuS), designed to enable zero-shot skill-based policy adaptation to various contexts specified in natural language. It leverages LLMs to guide a skill diffusion model, allowing for the generation of adaptable skill trajectories. The framework uses a hierarchical skill learning structure and a loss-guided diffusion process with sequential in-painting techniques. The authors demonstrate the effectiveness of LDuS in adapting to diverse contexts through experiments on robotic manipulation tasks, showing improved performance compared to other methods.
strength_weakness: strength_weakness: ### Major comments: 
* The combination of a diffusion model and a VAE to learn the skills from a dataset of trajectories is not well justified. In particular, Equation (5), which is the loss function for learning the skills, informally replaces the reconstruction loss in ELBO of VAE with the noise prediction loss of a diffusion model. It is unclear if this is mathematically and conceptually correct. If this equation is adapted from a paper, it should be cited. Otherwise, more justification for this combination is needed. The authors are encouraged to look into works that combine these two models, such as:

Kingma, Diederik, Tim Salimans, Ben Poole, and Jonathan Ho. "Variational diffusion models." Advances in neural information processing systems 34 (2021): 21696-21707.


* The seeding procedure is unclear and possibly incorrect. In the caption of Table 1, the authors state that the methods are "evaluated" on 3 seeds. This suggests that they trained on a single seed and evaluated on 3 seeds, which is not the correct use of seeding. The proper approach is to train and evaluate for each seed. Additionally, more seeds should be used to ensure robustness. While computation costs are a concern, using at least 5 seeds, and preferably 10, is recommended.
* The contribution of the paper does not appear to be significant. It combines multiple components into a framework. In some cases these combinations are not well-justified either (e.g., VAE+diffusion mentioned above). While I understand the results are promising, I have concerns about the significance of the contributions.


### Minor comments:
* Notation fix: In Equation 1, the action sampling needs to be added to the subscript of the expectation: $a_t \sim \pi_c(\cdot | s_t, g_t)$. 

* Given that VAE is one of the main components of the method, its paper should be cited: 

Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." arXiv preprint arXiv:1312.6114 (2013).

* The cross-reference to Table 3 in the related works section is out of place. Table 3 presents empirical results, and referring to it during the literature review breaks the coherence of the paper and the order of cross-references.

* A bit nit picky, but the abbreviation "LDuS" for "LLM-based skill diffusion" is unclear and should be clarified.
questions: questions: 1. What is the justification for Equation (5)? 
2. Please explain the seeding procedure.
limitations: limitations: The authors briefly discuss one limitation of their work in line 320. However, more discussion is needed given the multiple components and heavy reliance on LLMs for tasks such as encoding the goal with CLIP or encoding the context with GPT. Additionally, potentially negative societal impacts are not discussed.

Comments: title: Rebuttal Acknowledgement
comment: We are very encouraged by your satisfaction with our rebuttal. As addressed in response to question 1, we will include the inference time comparison in our paper. We deeply appreciate for your thoughtful feedback, as well as your decision to the recommend acceptance of our work.
I am satisfied with the author's rebuttal.

I would recommend including the inference time comparison in the paper, as this shows one of the current limitations of the approach.

I will keep my score and recommend acceptance of the work.
We are very satisfied that your concerns have been addressed and appreciate your decision to increase the score. We are also greatly encouraged by your recognition of our work as promising.

For the final version of our paper, we will include experiments with the 'spatial' and 'energy' contexts, as well as the image-based version of LDuS. Additionally, we will include details of our zero-shot evaluation settings.

Once again, we truly appreciate your thoughtful and insightful feedback.
I thank authors for running additional experiments, especially the spatial context ones, it definitely builds more confidence into the method. I believe the proposed work is a promising direction towards incorporating various user-contexts and safety constraints. All my concerns have been addressed and I have updated my score.
We deeply appreciate your positive rating of our paper and are very satisfied that your concerns have been addressed. Due to the text limit on the rebuttal, we did not include the standard deviation; however, it will be included in our final paper.

Regarding to the scope of the context, we will include experiments both 'spatial' and 'energy' contexts in our final paper. Additionally, we will address the several questions raised by the reviewer, such as the definition of context reward, the number of iterative refinement required, and sensitivity to guidance strength in our final paper.

We thanks for the reviewer once again for your thoughtful response.
I want to thank the authors for their efforts in adding experiments and detailed responses to my questions. I suggest that the authors also include the standard deviation in the final paper for better comparison. Most of my concerns have been addressed, I will maintain my positive rating towards the paper.
We are very satisfied that your concerns have been addressed. We truly appreciate your valuable feedback, as well as your decision to increase the score.

We fully agree with your perspective that a mathematical justification for the VAE+Diffusion model is required. As mentioned in our rebuttal, the reconstruction loss in the VAE can be seamlessly replaced by the diffusion loss. We will include this justification in the final version of our paper.

We once again deeply thank you for your thoughtful feedback on our paper.
I thank the authors for their reply. Most of my concerns were addressed and I have increased my score.