Few-Shot Task Learning through Inverse Generative Modeling

Authors: aviv  netanyahu
yilun  du
antonia  bronars
jyothish  pari
joshua  tenenbaum
tianmin  shu
pulkit  agrawal
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 19:47:55

Abstract: Learning the intents of an agent, defined by its goals or motion style, is often extremely challenging from just a few examples. We refer to this problem as task concept learning and present our approach, Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM), which learns new task concepts by leveraging invertible neural generative models. The core idea is to pretrain a generative model on a set of basic concepts and their demonstrations. Then, given a few demonstrations of a new concept (such as a new goal or a new action), our method learns the underlying concepts through backpropagation without updating the model weights, thanks to the invertibility of the generative model. We evaluate our method in five domains -- object rearrangement, goal-oriented navigation, motion caption of human actions, autonomous driving, and real-world table-top manipulation. Our experimental results demonstrate that via the pretrained generative model, we successfully learn novel concepts and generate agent plans or motion corresponding to these concepts in (1) unseen environments and (2) in composition with training concepts.

Main paper: 
 introduction: 
 The ability to learn concepts about a novel task, such as the goal and motion plans, from a few demonstrations is a crucial building block for intelligent agents -it allows an agent to learn to perform new tasks from other agents (including humans) from little data. Humans, even from a young age, can learn various new tasks from little data and generalize what they learned to perform these tasks in new situations [1] .In machine learning and robotics, this class of problems is referred to as Few-Shot Learning [2] . Despite being a widely studied problem, it remains unclear how we can enable machine learning models to learn concepts of a novel task from only a few demonstrations and generalize the concepts to new situations, just like humans do. Common approaches learn policies either directly, which often suffer from covariate shift [3] , or via rewards [4] [5] [6] , which are largely limited to previously seen behavior [7] . In a different vein, other work has relied on pretraining on task families and assumes that task learning corresponds to learning similar tasks to ones already seen in the task family [8, 9] .Inspired by the success of generative modeling in few-shot visual concept learning [10] [11] [12] , where concepts are latent representations, in this work, we investigate whether and how few-shot task concept learning can benefit from generative modeling as well. Learning concepts from sequential demonstrations rather than images is by nature more challenging due to sequential data often not Figure 1 : Few-shot concept learning. Given paired task demonstration τ (e.g., 'walk') and concept c (a latent representation of the task), we train a generative model G θ to generate behavior from a concept. Then, given demonstrations of a new behavior τ (e.g., 'jumping jacks') without its concept label, we aim to learn its concept representation by optimizing concept c as input to frozen G θ .Figure 2 : Experiment Domains. We extensively evaluate our approach for various domains.satisfying the i.i.d. assumption in machine learning [13] . In particular, we assume access to a large pretraining dataset of paired behaviors and task representations to learn a conditional generative model that synthesizes trajectories conditioned on task descriptions. We hypothesize that by learning a generative model conditioned on explicit representations of behavior, we can acquire strong priors about the nature of behaviors in these domains, enabling us to more effectively learn new behavior that is not within the pretraining distribution, given a limited number of demonstrations, and further generate the learned behavior in new settings.To this end, we propose Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM).In our approach, we first pretrain a large conditional generative model which synthesizes different trajectories conditioned on different task descriptions. To learn new tasks from a limited number of demonstrations, we then formulate few-shot task learning as an inverse generative modeling problem, where we find the latent task description, which we refer to as a concept, which maximizes the likelihood of generating the demonstrations. This approach allows us to leverage the powerful task priors learned by the generative model to learn the shared concepts between demonstrations without finetuning the model (Figure 1 ). We demonstrate this approach in various domains: object rearrangement, where concepts are relations between objects, goal-oriented navigation, where concepts are target attributes, motion capture, where concepts are human actions, autonomous driving, where concepts are driving scenarios, and real-world table-top manipulation where concepts are manipulation tasks (Figure 2 ).New concepts are either (1) compositions of training concepts (e.g., multiple desired relations between objects that define a new object rearrangement concept) or (2) new concepts that are not explicit compositions in the natural language symbolic space of training concepts (e.g., a new human motion 'jumping jacks' is not an explicit composition of training concepts 'walk', 'golf' etc.) Thanks to generative models' compositional properties that enable compositional concept learning [14] , in addition to being able to learn a single concept from demonstrations directly, FTL-IGM learns compositions of concepts from demonstrations that, when combined, describe the new concept.We show that our approach generates diverse trajectories encapsulating the learned concept. We achieve this due to two properties of generative models. First, these models have shown strong interpolation abilities [15, 16] , which allow generating the new concept on new initial states they were not demonstrated from. Second, these models have compositional properties that enable compositional trajectory generation [17] , which allow composing learned concepts with training concepts to synthesize novel behavior that was not demonstrated (e.g., 'jumping jacks' and 'walk'), see Figure 3 . We further demonstrate that our approach addresses a unique challenge introduced in learning task concepts: we utilize plans generated by learned concepts in a closed-loop fashion.Our main contributions are (1) formulating the problem of task learning from few demonstrations as Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM), (2) adapting a method for efficient concept learning to this problem based on the new formulation, and (3) a systematic evaluation revealing the ability of our method to learn new concepts across a diverse set of domains. 
related work: 
 Learning from few demonstrations. Our problem setting is closely related to learning from few demonstrations. There has been much work on learning to generate agent behavior given few demonstrations. There are several common approaches to this problem. First, behavior cloning (BC) is a supervised learning method to learn a policy from demonstrations that predicts actions from states. Similar to our framework, goal-conditioned BC can predict states from task representations and states [18] . Finetuning these models to learn new behaviors requires labeled demonstrations of the new task. We assume unlabeled demonstrations. BC often suffers from covariate shift [3] and fails to generate the demonstrated behavior in novel scenarios. This can be mitigated by assuming access to a human in the loop [19] . Second, the inverse reinforcement learning (IRL) framework learns a policy that maximizes the return of an explicitly [5, 6, 20] or implicitly [21] learned reward function. These works learn a reward for a single task or for a set of tasks (e.g., goal-conditioned IRL [22, 23] and multi-task IRL [24] ). While IRL is more data efficient than BC, it is computationally costly due to learning a policy every iteration via an inner reinforcement learning (RL) loop. Additionally, it requires access to taking actions in the environment during training and when faced with a new task, we have to retrain the reward again. A third approach is inverse planning [25] [26] [27] , which can robustly infer concepts such as goals and beliefs even in unseen scenarios. However, it assumes access to a planner, knowledge about environment dynamics, and the task/goal space. Finally, in-context learning approaches [8, 28, 29] learn actions in a supervised manner by representing the task with demonstrations. This allows few-shot generalization without further training.In contrast, we do not learn an action-generating policy directly or via a reward function. For concept learning, we do not assume having access to any given planner, world model, actions, rewards, or prior over the task space. Instead, we learn concepts (task representations) from demonstrations via a pretrained generative model that takes a concept as input and directly produces state sequences. We then input the learned concept into the generative model to produce behavior similar yet diverse to the demonstrated one. We further demonstrate how to use these state sequences with a planner to take actions and achieve the desired behavior. The idea of concept learning via generative models has been explored for computer vision applications [11, 12] . We build on this work and show how to extend it to learn agent task concepts. Our work also differs from prior works on learning trajectory representations [30] [31] [32] [33] [34] . These works focus on learning plans over trajectory embeddings, whereas we learn a task representation from demonstrations on which we condition to generate behavior.Making. There has been work on generative modeling for decisionmaking, including generative models for single-agent behaviors, such as implicit BC [35] , Diffuser [36, 17] , Diffusion Policy [37] , Decision Transformer [38] [39] [40] , and for multi-agent motion prediction such as Jiang et al. [41] . The success of diffusion policy in predicting sequences of future actions has led to 3D extensions [42] , and combined with ongoing robotic data collection efforts [43] and advanced vision and language models, has led to vision-language-action generative models [44] [45] [46] . In this work, we utilize a conditional generative model for the inverse problem, i.e., learning concepts from demonstrations.Composable representations. There has been work on obtaining composable data representations. β-VAE [47] learns unsupervised disentangled representations for images. MONet [48] and IODINE [49] decompose visual scenes via segmentation masks and COMET [50] and [12] via energy functions.There is also work on composing representations to generate data with composed concepts. Generative models can be composed together to generate visual concepts [14, [51] [52] [53] [54] [55] and robotic skills [17] . The generative process can also be altered to generate compositions of visual [56] [57] [58] [59] and molecular [60] concepts. We aim to obtain task concepts and generate them in composition with other task concepts. 
formulation: 
 Inspired by recent success in large generative models, we propose a generative formulation for learning specific behavior given a small set of demonstrations, which we term Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM). In our formulation, we assume access to a large pretraining dataset D pretrain = {(τ i , c i )} N i=1 of state-based sequences τ i = {s 0 , s 1 , ...} ⊆ T of states from state space S annotated with meta-data "concepts" c i ∈ C ⊆ R n describing trajectories. This assumption is often not prohibitive in practice. There is typically a vast amount of existing data collected from the internet or prior exploration in an environment, which may only need to be weakly annotated to characterize the trajectory, e.g., the goal state. Given D pretrain , we learn a conditional generative model G θ : C ×S → T conditioned on concepts and initial states, which learns to generate future trajectories. We train the parameters of G θ to maximize likelihoodarg max θ E τ,c∼Dpretrain [log G θ (τ |c, s 0 )].Then, given an unlabeled demonstration dataset D new ∼ D, we formulate learning a new concept c that is used to sample trajectories from D as inverting the generative model. In particular, we learn new concept c so that our frozen conditional generative model G θ maximizes the likelihood of trajectories in D new , corresponding to arg max c E τ ∼Dnew [log G θ (τ |c, s 0 )]. We find that this design choice enables us to leverage the priors learned by G θ from D pretrain to effectively learn concepts from D new given very few demonstrations, even if the demonstrated D new deviates from the concept labels c seen in D pretrain . For evaluation in closed loop, we further assume access to a planner that given two states plans which action to take in the environment, sometimes via access to simulation in the environment. We use this planner sequentially to make decisions in the environment.The key difference between our approach to few-shot adaptation from demonstrations and prior approaches is the assumption and usage of a large pretraining dataset of paired behaviors and concepts D pretrain combined with an invertible generative model. We learn new concepts solely from demonstrations without finetuning model weights or taking actions in the environment by relying on the pretrained concept space. 
few-shot concept learning based on ftl-igm: 
 We adapt a few-shot concept learning method to task concepts based on the FTL-IGM framework.During training we learn a generative model G θ from training {(τ i , c i )} i pairs. We then freeze G θ , and given demonstrations of a new task {τ } i , optimize a concept c to produce the new behavior via G θ . We then generate a diverse set of behaviors via G θ , either for the learned concept c conditioned on new initial states or for compositions of c with other concepts.A diffusion model is a generative model that given a forward noise adding process q(x t |x t-1 ) := N (x t ; √ 1 -β t x t-1 , β t I) starting from data x 0 according to a variance schedule β 1 , ..., β T , learns the reverse process p θ (x t-1 |x t ) := N (x t-1 ; µ θ (x t , t), Σ θ (x t , t)). Ho et al. [61] simplify the training objective to estimate noise E t∼U {1,T },x0,ϵ∼N (0,I) ||ϵ -ϵ θ (x t , t)|| 2 where x t is produced by adding noise ϵ to data x 0 by the forward noising process at diffusion step t, q(x t |x 0 ) := N (x t ; √ ᾱt x 0 , (1ᾱt )I) where ᾱt := Π t s=1 (1 -β s ). Dhariwal and Nichol [62] enable conditioned generation by guiding the reverse noising process with classifier gradients. The noise prediction becomes ε = ϵ θ (x t , t) -ω √ 1 -ᾱt ∇ xt log p ϕ (y|x t , t) where classifier p ϕ (y|x t , t) is trained on noisy images, and ω is the guidance scale. Ho and Salimans [63] introduce classifier-free guidance that achieves the same objective without the need for training a separate classifier. This is done by learning a conditional and unconditional model by removing the conditioning information with dropout during training. The noise prediction is then ε = ϵ θ (x t , t) + ω(ϵ θ (x t , y, t) -ϵ θ (x t , t)). Ramesh et al. [64] and Nichol et al. [65] demonstrate how this idea can be used to generate images conditioned on a class. Diffusion models have recently shown success as generative models for decision making [36, 17] . Specifically, Ajay et al. [17] used a conditional classifier-free guidance diffusion model [63] to generate trajectories of future states to reach given an input observation. We adopt this objective and learn a denoising model ϵ θ conditioned on latent concepts and initial observed states to estimate noise of a future state trajectory:EQUATION)where p is the probability of removing conditioning information which is then replaced by dummy condition c ∅ , and s 0 is the initial state corresponding to trajectory τ . x t (τ ) is obtained from x 0 = τ by the forward noising process. We then extend this approach for the inverse problem, namely, learning a concept from demonstrations.Gal et al. [11] use a frozen generative model to learn visual concept representations from few images depicting the concept by optimizing the model's input v * = arg min v E x0,v,ϵ∼N (0,1),t ||ϵ - ϵ θ (x t , c θ (v), t)|| 2 2, where c θ and ϵ θ are fixed. Liu et al. [12] extend this and learn visual concept compositions with a pretrained diffusion model in an unsupervised manner. Namely, from a set of images that depict various concepts, for each image x i they learn a set of weights ω i k and a shared set of visual concepts for all images c k , ε = ϵ(x i t , t)+ K k=1 ω i k (ϵ(x i t , c k , t) -ϵ(x i t , t)).We extend these formulations to inferring multiple concepts, whose composition describes a single task concept, from few demonstrations of a task.Given a trained diffusion model ϵ θ and demonstrations of a new concept {τ } i from D new , we learn concepts {c 1 , ..., cK } for K ≥ 1 and their weights {ω 1 , ..., ω K } that best describe the demonstrations. Starting from uniformly sampled concept embeddings ck ∼ U([0, 1] n ), we freeze ϵ θ , and optimize ck and ω k :E ϵ∼N (0,I) [||ϵ -(ϵ θ (x t (τ ), c ∅ , s 0 , t) + K k=1 ω k (ϵ θ (x t (τ ), ck , s 0 , t) -ϵ θ (x t (τ ), c ∅ , s 0 , t)))|| 2 ]. (2)We find that this compositional approach enables us to effectively represent and learn new demonstrations, even when demonstrations are substantially different than those seen in training tasks.After learning concepts ck , whose composition describes the new task c, we evaluate the behavior it generates by initializing x T (τ ) ∼ N (0, αI), and compute x t ∼ N (µ t-1 , αΣ t-1 ) iteratively as a function of the estimated denoising function ε(ϵ θ ), where µ and Σ are the mean and variance that define the reverse process, and α ∈ [0, 1) is a scaling factor that leads to lower temperature samples, until generating x 0 = τ representing the trajectory of the agent. The denoising function is constructed by fixed or learned weights as defined in Eq. 2 and by any number of concepts ≥ 1. The applications of the generation procedure can be summarized as:Learned concept and demonstrated initial states. We apply our learned concept to a set of demonstrated initial states. In domains where the initial state and concept jointly determine optimal behavior, the generated trajectory corresponds to optimal actions to execute (e.g. goal-oriented navigation). In contrast to other domains where the initial state is irrelevant for a task due to the randomness in sampling x T (τ ) (e.g. motion capture), generated trajectories correspond to diverse plausible behaviors exhibiting the learned concept.Learned concept and novel initial states. We further apply the generation procedure from our learned concept on novel initial states, to generate trajectories of new behaviors exhibiting our conditioned concept. Prior methods may suffer from covariate shift in this setting [21] . We empirically show that our method is less prone to this problem.Learned concept composed with other concepts. Finally, we modify the generation procedure of our newly learned concept to generate trajectories that simultaneously exhibit other concepts. To generate a trajectory with an added another concept, we add another term to the sum in Eq. 2 where the learned concept is composed with a training concept c k and its weight ω k : ω k (ϵ θ (x t (τ ), c k , s 0 , t)ϵ θ (x t (τ ), c ∅ , s 0 , t)). This modified generation procedure constructs trajectories which exhibits behavior that has a composition of the learned concept and the other specified concepts [14] . Similarly to Ajay et al. [17] , in environments where an inverse dynamics model is provided, we generate trajectories in a closed loop. We execute actions calculated by the inverse dynamics given the predicted plan by the model and then repeatedly replan given new observations. 
experiments: 
 We demonstrate results in four domains where concept representations are T5 [66] embeddings of task descriptions in natural language for training, and empty string embeddings for the dummy condition. During few-shot concept learning, we are provided with three to five demonstrations of a composition of training concepts or of a novel concept that is not an explicit composition of training tasks in natural language symbolic space. We ask a model to learn the concept from these demonstrations.Learning concepts describing goals that are spatial relations between objects. Object rearrangement is a common task in robotics [67] [68] [69] and embodied artificial intelligence (AI) [70, 71] , serving as a foundation for a broader range of tasks such as housekeeping and manufacturing. Here, we use a 2D object rearrangement domain to evaluate the ability of our method to learn task specification concepts. Given a concept representing a relation between objects, we generate a single state describing that relation. The concept in a training example describes the relation (either 'right of' or 'above') between only one pair of objects (out of three objects) in the environment. Then, a model must learn compositions of these pairwise relations and new concepts such as 'diagonal' and 'circle' (see Figure 4 ). The results in Figures 5 and 6 demonstrate that our method learns unseen compositions of training concepts and new concepts. They further demonstrate how our method composes new concepts with learned concepts. For additional qualitative results, please refer to Appendix A.While successful in most cases, there are also a few failure examples. The accuracy for the new 'circle' concept is low (0.44) compared to the mean over task types in Figure 6 Object Rearrangement New Concept (0.82 ± 0.09). This is most likely due to this concept lying far out of the training distribution. The task 'square right of circle ∧ triangle above circle' has low accuracy for 2 concepts (0.32) compared to the mean in Table 2 Object Rearrangement Training Composition (0.75 ± 0.11). This may arise from the combined concept-weight optimization process -as there is no explicit regularization on weights, they may converge to 0 or diverge. In Figure 12 , we show that concept components may or may not capture new concept relations.Learning concepts describing goals based on attributes of target objects. We test our method in a goal-oriented navigation domain adapted from the AGENT dataset [72] , where an agent navigates to one of two potential targets. Conditioned on a concept representing the attributes of the desired target object and initial state, we generate a state-based trajectory describing an agent navigating to the target. Each object has a color and a shape out of four possible colors and four shapes. During training, we provide 16 target-distractor combinations that include all colors and shapes (but not all combinations), and a concept is conditioned on one of the target's attributes (e.g., color). We introduce new concepts defined by both target attributes, including (1) unseen color-shape target combinations and (2) new target-distractor combinations. Figure 7 shows an example. In training, we see bowl and red object targets. A new concept includes a novel composition as the targetred bowl. The new concept distractor objects (green bowl and red sphere) were introduced during training, but they were not paired with a red bowl as the target. As Figure 13 shows, our method successfully learns concepts where targets are new compositions of target attributes in settings with new target-distractor pairs and generalizes to new initial object states. We further evaluate our model and baselines in closed loop (Figure 6 ) by making an additional assumption that a planner is provided. The planner produces an action given a current state and a future desired state predicted by a model.Learning concepts describing human motion. Unlike prior work on learning to compose human poses from motion capture (MoCap) data [e.g., 73 , 74], here we focus on the inverse problemlearning new actions from MoCap data. In particular, we use the CMU Graphics Lab Motion Capture Database (http://mocap.cs.cmu.edu/). We train on various human actions in the database and few-shot learn three novel concepts (see Appendix B.3 for details). Learning tasks from few demonstrations is especially beneficial in this domain since describing motion concepts in words could be hard. In Table 5 .1, we ask five human volunteers to select generated behaviors that depict training and new concepts. We demonstrate quantitatively that our method generates human motion that captures the desired behavior across training and new behaviors. We qualitatively demonstrate how our method generates learned new concepts ('jumping jacks' and 'breaststroke') from new initial states and composes 'jumping jacks' with training concepts 'walk', 'jump', and 'march'.Results are best viewed on our website website. We compare motion generated by our method to various baselines on new initial states for 'jumping jacks' and 'breaststroke'. Our method is noisy yet captures the widest range of motion. While other methods often produce smoother trajectories, they mostly capture local (VAE) or degenerate (BC, In-Context, Language) motion.Learning concepts describing driving scenarios. In an Autonomous Driving domain [75] , an agent acts in a challenging multi-agent environment to complete a driving task. We train on several driving scenarios ('highway', 'exit', 'merge', and 'intersection') and learn a new driving scenario ('roundabout') from several demonstrations (see Figure 8 and further details in Appendix B.4). We evaluate this scenario in closed loop on new initial states, assuming access to a planner that can simulate taking actions in the environment. Over two evaluation metrics (crash and task completion rate), our method achieves overall best results (Figure 9 ).Learning concepts describing real-world table-top manipulation tasks. We evaluate our method's capability to learn a novel concept for real-world table-top manipulation with a Franka Research 3 robot. Training concepts include 'pick green circle and place on book', 'pick green circle and place on elevated white surface', 'push green circle to orange triangle' and 'push green circle to orange triangle around purple bowl'. The new scenario includes pushing the green circle to the orange triangle on a book (Figure 10 ). We evaluate training pushing in closed loop and and achieve success (90% accuracy). We evaluate the new concept, elevated pushing, against a baseline that conditions on the training 'push green circle to orange triangle' representation in the new scenario setup where the objects are placed on a book. Learning the new concept succeeds (55% accuracy) while using the training representation fails (15% accuracy) and the robot often pushes the book instead of the object. Details are in Appendix B.5 and qualitative results are on our website.Goal conditioned behavior cloning. We compare our method with goal-conditioned behavior cloning (BC), which, given a condition and a state, outputs the next state in a sequence. It is trained on our paired pretraining dataset and learns concepts by optimizing the input condition to maximize the likelihood of new concept demonstration transitions. We test its ability to compose new learned concepts with training concepts naively by adding conditions that are then inputted into the model. We observe that even though goal-conditioned BC has access to the pretrained dataset and conditions, and while it may learn new concepts, it suffers from covariate shift on new initial states and lacks the ability to generalize to novel compositions of the learned concept with training concepts (Figures 6, 9, 13). To achieve these generalization capacities, we need a model that can process interpolated (initial states) and composed (concepts) conditions, such as our generative model.Learning the concept space with a VAE. We compare our method with VAE [76] that does not utilize the concepts in the paired pretraining data but rather learns the concept space by reconstructing pretraining data trajectories through their encoded representation z. Trajectories are generated via a decoder conditioned on an initial state and z with added noise. z is obtained by encoding a trajectory for training evaluation and by fixing the decoder and optimizing z to generate a given demonstration when learning a new concept. We find that the VAE model learns a latent space that captures training and new concepts but does not enable generalization to new initial states (Figures 6, 9 completion rate yet a high crash rate. In-context has a low crash rate yet 0% success rate -typically, the controlled vehicle reaches the roundabout's center but does not complete the crossing. Overall, our method learns to complete the roundabout crossing with competitive crash and success rates.Training concepts New concept closed loop evaluation pick&place on book push on surface push around pick&place on table push on book In-Context learning. We compare our approach with training a method to in-context learn from demonstrations. Specifically, we compare our approach to Xu et al. [28] using the pretrained dataset for few-shot behavior generation without further training. We sequentially predict states from demonstrations of a concept and window of current states and show that our method adapts better to new concepts (Figures 6, 9, 13) . This emphasizes the need to learn explicit concept representations.Conditioning on Language Descriptions of New Concepts. There has been work on generating actions from language instructions [77] [78] [79] . We demonstrate that in our setup, merely providing new concept language instructions embedded with T5 (as in our pretraining dataset) is insufficient, and generalization is better when learning concepts from few demonstrations. For AGENT, training compositions on new initial states has an average accuracy and standard error of 0.63 ± 0.07, lower than ours (0.73 ± 0.07). For Object Rearrangement, training compositions (0.2 ± 0.07), new concept (0.2 ± 0.05), and new and training concept compositions (0.13 ± 0.04) accuracies are significantly lower than ours (0.9 ± 0.04, 0.82 ± 0.09 and 0.8 ± 0.02). For MoCap, we demonstrate qualitatively that instead of capturing new human actions, the agent transitions into walking. Results are best viewed on our website.When learning weights together with concepts, we check the effect of the number of learned concepts and weights. We report results in Table 2 for Object Rearrangement, AGENT, and Driving, and find that, on average, learning two concepts improves concept learning. We demonstrate qualitatively for MoCap that learning two conditions is preferable. In 'jumping jacks', we observe that the motion lacks raising and lowering both arms and in 'breaststroke', it lacks complete arm and upper torso movement. Results are best viewed on our website. 2 : Ablation on the number of learned concepts. We test the effect of the number of learned concepts and weights in FTL-IGM on the generation accuracy of new learned concepts. On average, learning two concept components and their weights is preferable to learning one concept component and its weight. We report average accuracy and standard error over task types for Object Rearrangement and AGENT. Driving includes a single new concept and we report accuracy only.learned component captures a single composed concept. In other cases, a single learned component captures both concepts (Figure 12 ).New concepts that are not explicit compositions of training concepts in natural language symbolic space. In Figure 11  
discussion and limitations: 
 In this work, we formulate the problem of new task concept learning as Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM). We adapt a method for concept learning based on this new formulation and evaluate task concept learning against baselines in four domains. Our extensive experimental results show that, unlike the baselines, FTL-IGM successfully learns novel concepts from a few examples and generalizes the learned concepts to unfamiliar scenarios. It also composes learned concepts to form unseen behavior thanks to the compositionality of the generative model. These results demonstrate the efficacy, sample efficiency, and generalizability of FTL-IGM.However, our work has several limitations. First, while our framework is general for any parameterized generative model, our implementation with a diffusion model incurs high inference time. We note that there is still space for improvement in the MoCap generation quality and in the compatibility rate of demonstrations generated by composing learned and training concepts. In addition, we assume that learned concepts lie within the landscape of training concepts to learn them from a few demonstrations without retraining the model. We have approached the question of what new concepts can be represented by compositions of concepts in this landscape empirically, leaving a theoretical analysis as future work. We are hopeful that with the continued progress in the field of generative AI, more powerful pretrained models will become available. Combined with our framework, this will unlock a stronger ability to learn and generalize various task concepts in complex domains.Object Rearrangement. We display additional states generated by our model for various new concepts. We generate learned concepts that are compositions of training concepts (Figure 14 ), new concepts that are not explicit compositions of training concepts (Figure 15 ), and a new concept composed with a training concept (Figure 16 ). We analyze the learned concepts qualitatively (Figure 12 ) and quantitatively (Table 3 ).Goal-Oriented Navigation. We provide accuracy for open loop evaluation in Figure 13 . We further analyze the learned concepts quantitatively (Table 3 ). • Five new concepts (new concept in Figure 6 ). 'circle': circle, triangle, and square lie in equal intervals on the circumference of a circle with radius 1. We verify that objects do not overlap and that no training relations unrelated to the specified new concepts exist in the demonstrations between objects.State space. The state space is a 21-tuple describing three shapes (circle, triangle, and square), each represented by a 7-tuple: their center 2D position, size, angle, and one-hot shape type. A -y B | ≤ 2•max{r A , r B }, A is considered 'above' B if y A > y B and |x A -x B | ≤ 2•max{r A , r B }.A, B and C lie on a 'circle' if their centers form a circle of radius r where |r -1.67| < 0.3 and A is 'diagonal' to B if x A > x B , y A > y B and they lie within a 2 • max{r A , r B } margin of f (x) = x.Training. To collect demonstrations, we follow the data generation process in the AGENT benchmark environment [72] , which provides a planner for navigation given the desired target. In the provided environment, the agent's initial position a p t=0 = (0, 0.102, -3.806), color (yellow) and shape (cone) are fixed. There are two objects: the target and a distractor. = {1, 2} \ i), 'go to yellow bowl' (o c i , o c j = yellow, o s i = bowl, o s j ∈ {cone, sphere} or o s i , o s j = bowl, o c i = yellow, o c j ∈ {purple, green}), 'go to red cube' (o c i , o c j = red, o s i = cube, o s j ∈ {cone, sphere} or o s i , o s j = cube, o c i = red, o c j ∈ {purple, green}), 'go to yel- low cube' (o c i , o c j = yellow, o s i = cube, o s j ∈ {cone, sphere} or o s i , o s j = cube, o c i = yellow, o c j ∈ {purple, green}), 'go to purple cone' (o c i , o c j = purple, o s i = cone, o s j ∈ {bowl, cube} or o s i , o s j = cone, o c i = purple, o c j ∈ {red, yellow}).Note that in each scenario, the distractor object either has the same color or shape as the target, and combined with its other attribute (shape or color), it is within the training distribution (i.e., red and yellow cones and spheres, and purple and green bowls and cubes). There are five demonstrations of each novel composition.• Conditioning each novel composition concept on novel initial states sampled from the novel concept distribution (training composition new initial state in Figure 6 ).State space. The initial state space we condition on is a 52-tuple based on the state space in Shu et al. [72] describing the agent and two objects. The agent is represented by its 3D position, quaternion ∈ R 4 , velocity ∈ R 3 , angular velocity ∈ R 3 , one-hot type (representing the agent and two objects), rgba color, and one-hot shape (representing the four possible shapes). Similarly, each object is represented by its 3D position, type, color, and shape. The demonstration state space is a 13-tuple representing the first 13 dimensions of the initial state space (agent position, quaternion, velocity, and angular velocity). Demonstrations i ∈ [N ] have horizons H i ≤ 150 and are padded to length 150 using the final state. During training, we sample the demonstrations to generate subtrajectories of length 128.Evaluation data and metrics. For training, we generate 50 trajectories conditioned on uniformly sampled training concepts and initial states. For each new concept, we generate five trajectories conditioned on the learned concept and five initial states from the new concept demonstrations. To evaluate each concept composition on new initial states, we generate trajectories conditioned on the learned concept and 50 initial states sampled from the new concept distribution. Results in Figure 6 are reported for best learned ω and fixed ω ∈ {1.2, 1.4, 1.6, 1.8} -we learn two concepts with classifier-free guidance weight ω = 1.6. We report accuracy based on whether the agent in the generated trajectory has made progress towards the desired target (|a p t=128 -o p target | < |a p t=0 -o p target |). Accuracy for closed loop evaluation (Figure 6 ) is reported based on optimally reaching target o target before a distractoro distractor (|a p t=128 -o p target | < 0.365 ∧ ∀t < 128 : |a p t -o p distractor | ≥ 0.365).Closed loop evaluation. The action space A ⊆ R 2 represents forces applied to the agent. We take actions a t = τ t+5 -s t where s t is an observation in the environment, and τ t+5 is a future step in the open loop plan generated by G conditioned on observation s t and learned concept c. We execute actions in the environment until reaching the target or a maximum number of steps and report the evaluation metric described above. Closed loop evaluation is done in pybullet simulation [81] for efficiency.Training. All human actions in the CMU Graphics Lab Motion Capture Database (http://mocap.cs.cmu.edu/) are included in the training set except three new scenarios. We further discard videos with less than 128 frames. The training set includes 2210 demonstrations.• Three new concepts. 'jumping jacks' (3 demonstrations), 'chop wood' (2), and 'breaststroke'• Conditioning each novel concept on novel initial states -the last states in the new concept demonstrations.• Composing new learned concept 'jumping jacks' weighted by learned concept weights with three training concepts ('walk', 'jump' and 'march') weighted by ω = 1, and conditioned on a training initial state from the training concepts.State space. The initial state space is a 42-tuple representing the 3D position of 14 joints. The original dataset contains 31 joints. Our version is adapted with Tanke et al. [82] to reduce the number of joints to 14 and to remove rotation and translation. The demonstration state space is the same with horizonsH i , i ∈ [N ].We subsample the original trajectories every 4 steps and further sample trajectories to generate subtrajectories of unified length 32 on which we train.Evaluation data and metrics. For each new concept, we generate trajectories conditioned on the learned concept and an initial state from the new concept demonstrations. For evaluating each new concept on new initial states, we generate trajectories conditioned on the learned concept and novel initial states from the new concept demonstrations that were not used as initial states during concept learning, specifically the last state in each demonstration. For new concepts composed with training concepts, we generate trajectories, each conditioned on the learned concept, a uniformly sampled training concept, and training initial state. Results are reported for best learned ω and fixed ω ∈ {1.2, 1.4, 1.6, 1.8} -we learn two concepts with classifier-free guidance weight ω = 1.8.Human experiment. We show five humans (aged 23-28, four male) videos of training and new concepts with their natural language labels and ask whether each video depicts the concept and which depicts it best. The participants gave their consent, and the study was approved by an Institutional Review Board. We show three training concepts: 'march', 'run', and 'walk', and three new concepts: 'jumping jacks', 'chop wood', and 'breaststroke' on new initial states.Training. We use the following scenarios from the HighwayEnv driving simulation environment [75] . The training scenarios include four driving scenarios: 'highway', 'exit', 'merge', and 'intersection'. In 'highway', the objective is driving on a highway at high speed on the rightmost lanes while avoiding collisions. The highway has four lanes and 50 vehicles. The initial lane and position of all vehicles are sampled, as well as non-controlled vehicle speeds. An episode ends if the controlled vehicle crashes or a time limit is reached. In 'exit' the objective is to take a highway exit while driving on a four-lane highway with an exit lane and 20 vehicles. The controlled vehicle is rewarded for exiting at high speed and driving on the rightmost lanes while avoiding collisions. The initial position of all vehicles is sampled. An episode ends if the controlled vehicle crashes or a time limit is reached.In 'merge', the objective is driving on a highway with three lanes and a merging lane and three vehicles, one of them merging. The controlled vehicle is rewarded for driving at high speed while avoiding collisions and allowing another vehicle to merge into the highway. The lane, position, and speed are sampled for non-controlled vehicles. An episode ends if the controlled vehicle passes the merging lane or crashes. In 'intersection', the objective is making a left turn at an intersection with four two-way roads and 10 vehicles. The controlled vehicle is rewarded for crossing the intersection at high speed while staying on the road and avoiding collisions. The controlled vehicle's position is sampled, as well as the lane, position, and speed of other vehicles. An episode ends if the controlled vehicle completes crossing the intersection or crashes. Expert demonstrations were collected using a deterministic tree search planner provided in [83] .New Task. The new scenario includes:• A novel driving scenario: 'roundabout'. In this scenario, the objective is to take the second exit at a roundabout with four exits and four vehicles. The position and speed of noncontrolled vehicles are sampled. An episode ends if the controlled vehicle crashes or a time limit is reached. The novel concept is conditioned on novel initial states sampled from the new concept distribution. There are five demonstrations of this new concept.State space. The initial state space we condition on is in R 5×7 , the controlled vehicle and the four closest vehicles. Each vehicle is represented by a 7-tuple, including whether it is present on the road, its x and y positions, x and y velocities, and cosine and sine heading directions. The demonstration state space observations of the controlled vehicle in R 7 for horizonsH i , i ∈ [N ].We sample trajectories to generate subtrajectories of length 8, which we train on.Evaluation data and metrics. For evaluating the new concept on new initial states, we evaluate in closed loop on 50 initial states sampled from the new concept distribution. Task return, crash, and success rates are calculated based on the rewards described above. Note that the 'highway' scenario doesn't have a success score as there is no final state to reach in this scenario. Results are reported for best learned ω and fixed ω ∈ {1.2, 1.4, 1.6, 1.8} -for training, we report ω = 1.8 and during new concept learning, we learn two concepts and their corresponding classifier-free guidance weights.Closed Loop Evaluation. The action space is discrete: move to the left lane, stay idle, move to the right lane, drive faster, drive slower. We assume access to a planner P that given two states plans which action to take in the environment via access to simulation in the environment, a t = P(τ t+1 , s t ).The planner simulates the possible actions from observed state s t and randomly selects an action out of the ones closest to the model predicted next state τ t+1 that does not result in the controlled vehicle crashing. We execute actions until reaching a maximum number of steps or until the controlled vehicle crashes.Training. We collect 214 expert demonstrations with a Franka Research 3 robot via teleop with a Spacemouse for four table-top manipulation tasks: 'pick green circle and place on book' (29 demonstrations), 'pick green circle and place on elevated white surface' (30), 'push green circle to orange triangle' (124), and 'push green circle to orange triangle around purple bowl' (31) . Examples of these tasks are best viewed on our website.New task. The new scenario includes pushing the green circle to the orange triangle on a book. We provide ten demonstrations of this task.State space. The initial state space we condition on includes an overhead RGB image of the scene (Figure 10 ) and the robot's end effector pose, a 7-tuple of its 3D position and quaternion, and gripper state in R. We verify that most states are fully observable. The demonstration state space includes observations of the end effector pose and gripper state in R 8 for horizons H i , i ∈ [N ]. We sample trajectories to generate subtrajectories of length 32, which we train on.Evaluation data and metrics. We evaluate 20 episodes for training 'push green circle to orange triangle' and new concept 'push green circle to orange triangle on book' on new initial states sampled from the concept distributions. An episode is successful if the green circle touches the orange triangle before a maximum horizon is reached. Results for training are reported for ω = 1.8 and during new concept learning, we learn two concepts and their corresponding classifier-free guidance weights.Closed loop evaluation. The action space is an 8-tuple of the end effector pose and gripper state, equivalent to the predictions of the model. Given predicted action τ t+1 and end effector pose s t , we assume access to a planner that linearly interpolates between the current and next predicted pose. We make predictions and roll out the 32 predicted states in closed loop until it succeeds or executes a maximum number of steps H = 160.Diffusion model. We represent the noise model ϵ θ with an MLP for the object rearrangement domain and with a temporal U-Net for the AGENT, MoCap and Driving domains as implemented in Ajay et al. [17] . For Manipulation, we use a temporal U-Net where images are processed by a pretrained resnet18 [84] that is finetuned with the model during training. We use the same hyperparameters as in Ajay et al. [17] except for the probability of removing conditioning information, p, which we set to 0.1. In Table 4 we demonstrate the effect of choosing different classifier-free guidance weights ω. We are overall better or comparable to the baselines. BC. We implement behavior cloning (BC) [85] . The BC model is deterministic. In the Object Rearrangement domain, we only condition on concept c. To add stochasticity, when evaluating BC, we average results over 50 different seeds. We use an MLP with one hidden layer of size 512, ReLU activations, and AdamW [86] with learning rate 6 • 10 -4 . In AGENT, at each step, the input to the model is the fully observable initial state (52-dim) and the condition, and the model learns to predict the next partially observable agent state (13-dim). Composing concepts is implemented by adding conditions.VAE. We implement a conditional variational autoencoder model (CVAE) [76] 2 . We use an MLP for the architecture with AdamW [86] and learning rate 6 • 10 -4 . We sample 50 trajectories per concept and average the results. Composing concepts is implemented by adding conditions.In-Context. 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). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We discuss limitations in Section 6 and obtain an IRB for a small-scale human experiment.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). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: In this paper, we demonstrate that effectively learning behavior (or concept) representations from a few demonstrations is possible given a pretraining dataset, dependent crucially on the quality of training and new concept demonstrations. While such an effective few-shot learning system may have many societal implications, none are specific to this work that we feel a need to highlight.Guidelines:• 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). 
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: There are no high risks we feel the need to highlight. 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? Answer: [Yes] Justification: We add citations and links throughout the paper to assets used in this work. 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.based on the implementation in https://github.com/pytorch/examples/blob/main/vae/main.py 
 
 Appendix:  


Decision: The submission introduces a novel few-shot learning framework for understanding agent intents such as goals or motion styles from few examples, for multiple application domains (e.g., motion capture and autonomous driving). The proposed method leverages a pre-trained invertible neural generative model, initially trained on a variety of basic tasks and their demonstrations. By inverting the diffusion model, the method can obtain new concept embedding, which then can be used to generate the new trajectories (given any new state). The authors validated the method in diverse tasks including object rearrangement, goal-oriented navigation, and human action simulations. 

All reviewers agreed that the paper showed good results and the method is timely interesting to the community. So, I recommend accept. There are a few suggestions, such as improving the writing of the method section and including more baselines. Please add them to your final copy.

Reviews: review: paper_summary: paper_summary: The paper focuses on learning concepts that describe behaviors seen in state-based trajectories, such as mocap data or simplified autonomous driving simulation. The proposed approach uses a generative model to predict state trajectories based on concepts annotated with natural language. Next, new trajectories without annotated concepts are provided to the model. The method uses gradient descent to invert the network and optimize new concept embeddings for the novel trajectories. Then, the learned concepts embeddings are shown to generate correct trajectories in novel starting states. Moreover, the authors demonstrate composable concept embeddings.
strength_weakness: strength_weakness: 1. Section 4.2 could use more explanation. In particular, is there something specific about the diffusion model that makes it “invertible”? As far as I understand, optimizing with respect to the input using gradient descent is possible with all neural networks. Usually, the community uses the term “invertible networks” (this term is used in the paper) to mean a neural network F for which computing F^{-1} is easy [1, 2]. Moreover, Equation 2 could be explained better so that the paper is more self-contained.

2. Only state-based domains are used. It is unclear if this approach learns meaningful concepts when used with an image or a video generative model.

3. The exposition could be improved. The methods section is very short; a background section could prepare readers to understand it better. Moreover, the related work section would benefit from further discussion of generative models and composable representations.

References:

[1] Invertible Residual Networks. Behrmann et al. 2018.

[2] Analyzing inverse problems with invertible neural networks. Ardizzone et al. 2018.
questions: questions: 1. Is unconstrained gradient descent with respect to the input concept embedding the right way to learn concepts? For example, in the “deep dream” literature [3], gradient descent is used to create an input image that maximizes a particular neuron in the network. In this case, normal gradient descent does not work very well because it introduces only high frequency patterns in the image. There are various approaches that add noise to the gradients for smoothing, etc.

[3] https://research.google/blog/inceptionism-going-deeper-into-neural-networks/
limitations: limitations: The limitations of the method are addressed.
paper_summary: paper_summary: The paper addresses the challenge of learning the intents of an agent, such as its goals or motion style, from a few examples. The proposed approach, Few-Shot Task Learning Through Inverse Generative Modeling (FTL-IGM), leverages invertible neural generative models to learn new task concepts. The method involves pretraining a generative model on a set of basic concepts and their demonstrations, and then learning new concepts from a few demonstrations without updating the model weights. The approach is evaluated in four domains: object rearrangement, goal-oriented navigation, motion capture of human actions, and autonomous driving.
strength_weakness: strength_weakness: 1. **Abstract Concept:** The concept of "task concept learning" is not rigorously defined and remains too abstract. The distinction between "concept" and "context" or "contextual" is unclear, which could lead to confusion about the precise nature and scope of the proposed methodology.

2. **Scalability Concerns:** Reinforcement learning environments vary widely in complexity and characteristics. The scalability of the pretrained model to diverse and complex real-world environments is not well addressed, raising concerns about its practical applicability.
questions: questions: 1. **Baseline Comparisons:** The paper compares the proposed method only against BC and VAE. Why were these baselines chosen, and how does the method compare to other relevant techniques in the field? Including a broader set of baselines would provide a more comprehensive evaluation of the approach's effectiveness.
limitations: limitations: The paper tackles an important problem in reinforcement learning with an innovative approach that leverages invertible neural generative models. While the visualizations and generalizability are strong points, the abstract nature of the core concept and scalability concerns in diverse environments need to be addressed. Additionally, expanding the set of baseline comparisons would strengthen the empirical validation of the method.
paper_summary: paper_summary: The paper presents an approach to a new approach to few-shot learning. Specifically, the proposed method first pre-trains a conditional classifier-free guidance diffusion model on task concepts and their corresponding demonstrations. Few-shot learning is possible by inverting the diffusion model by optimizing for the task concepts.
strength_weakness: strength_weakness: 1. For each new task, we need to find the concept vector that corresponds to the new task. This might introduce further complications at inference time compared to few-shot learning in LLM works, the language-conditioned and the in-context learning baseline.
questions: questions: 1. Are concept representations updated during training or are they fixed (as the T5 embeddings)? 
2. How important is the classifier-free guidance weight? Does it require task-specific tuning?
limitations: limitations: Limitations are present in section 6.

Comments: title: Raising the score
comment: The responses address my concerns. I thus raise the score.