Few-Shot Task Learning through Inverse Generative Modeling
Authors: Aviv Netanyahu
Yilun Du
Antonia Bronars
Jyothish Pari
Joshua B. Tenenbaum
Tianmin Shu
Pulkit Agrawal
Venue: NeurIPS 2024 poster (2024)
Keywords: few shot learning
inverse generative modeling
Primary area: reinforcement_learning

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.

Few-Shot Task Learning through Inverse Generative Modeling
Aviv Netanyahu', Yilun Du Antonia Bronars', Jyothish Pari', Joshua Tenenbaum', Tianmin Shu', and Pulkit Agrawal
Massachusetts Institute of Technology
Harvard University
Johns Hopkins University
Abstract
Learning the intents of an agent, defined by its goals or motion style, is often extremely challenging from just few examples. We refer to this problem task concept learning and present our approach, Few-Shot Task Learning through Inverse Generative Modeling (FTL-IGM), which learns new task concepts by lever- aging invertible neural generative models. The core idea is to pretrain a generative model on set of basic concepts and their demonstrations. Then, given few demonstrations 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 manip- ulation. 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.
Introduction
The ability to learn concepts about a novel task, such as the goal and motion plans, from a few demonstrations is 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 Few-Shot Learning [2]. Despite being 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-6], which are largely limited to previously seen behavior [7]. different vein, other work has relied on pretraining on task families and assumes that task learning corresponds to learning similar tasks ones already seen in the task family [8, 9].
Inspired by the success of generative modeling in few-shot visual concept learning [10-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
"Correspondence to Aviv Netanyahu <avivn @ mit.edu:. Project website https:/aaiivnegiihubiiim hwt.:/n.iiituuuu.ooiiiii.....
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Pretraining
Few-Shot Concept Learning
given trajectories T and concepts C
given few demonstrations of a new concept
learn generative model Go
+ frozen model Goa - learn new concept č
C
Go
ic
Goa
T
forward pass
back propagation
Figure 1: Few-shot concept learning. Given paired task demonstration T (e.g., 'walk') and concept (a latent representation of the task), we train generative model Ge to generate behavior from concept. Then, given demonstrations of new behavior F (e.g., 'jumping jacks') without its concept label, we aim to learn its concept representation by optimizing concept è as input to frozen Go.
Object Rearrangement
Navigation to Goal
MoCap
Driving
Table-Top Manipulation
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 large pretraining dataset of paired behaviors and task representations to learn conditional generative model that synthesizes trajectories conditioned on task descriptions. We hypothesize that by learning generative model conditioned on explicit representations of behavior, we can acquire strong priors about the nature behaviors these domains, enabling us to more effectively learn new behavior that is not within the pretraining distribution, given 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 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 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 relations between objects, goal-oriented navigation, where con- cepts 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 manipulation tasks (Figure 2).
New concepts are either (1) compositions of training concepts (e.g., multiple desired relations between objects that define new object rearrangement concept) or (2) new concepts that are not explicit compositions in the natural language symbolic space of training concepts (e,g., 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 Iearn 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: utilize plans generated by Iearned concepts in 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.
2 Related Work
Learning from few demonstrations. Our problem setting is closely related to learning from few demonstrations. There has been much work learning to generate agent behavior given few demonstrations. There are several common approaches to this problem. First, behavior cloning (BC) is supervised learning method to learn 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 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 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 new task, we have to retrain the reward again. A third approach is inverse planning [25-27], which can robustly infer concepts such as goals and belief's even in unseen scenarios. However, assumes access to planner, knowledge about environment dynamics, and the task/goal space. Finally, in-context learning approaches [8, 28, 29] learn actions in 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 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 produce behavior similar yet diverse to the demonstrated one. We further demonstrate how to use these state sequences with planner to take actions and achieve the desired behavior. The idea of concept learning via generative models has been explored for computer vision applications [1 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-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.
Generative Models in Decision Making. There has been work on generative modeling for decision- making, including generative models for single-agent behaviors, such as implicit BC [35], Diffuser [36, 17], Diffusson Policy [37], Decision Transformer [38-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 aasion-language- In this work, we utilize a conditional generative model for the inverse problem, i.e., learning concepts from demonstrations.
generative models [44-46].
Composable representations. There has been work on obtaining composable data representations. 3-VAE [47] learns unsupervised disentangled representations for images. MONet [48] and TODINE [49] decompose visual scenes via segmentation masks and COMET [50] and [12] via energy functions. There also work on composing representations to generate data with composed concepts. Generative models can be composed together to generate visual concepts [14, 5 51-55] and robotic skills [17]. The generative process can also be altered to generate compositions of visual [56-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 small set of demonstrations, which we term Few-Shot Task Learning through
3

---

Few-Shot Concept Learning learn a new concept
Learned Concept Generation generate trajectories from the learned mociel concept from new initial states
in composition with training concepts
80 Cwalk
Go
Go
Go @
Figure 3: Diverse learned concept generation. We generate versions of the new behavior conditioned on the learned concept and (1) new initial states and (2) composed with other concepts.
Inverse Generative Modeling (FTL-IGM). In our formulation, we assume access to large pretraining } N of state-based sequences т; {So. s1,... C of states from state space S annotated with meta-data "concepts" G E C C R describing trajectories. This assumption often not prohibitive in practice. There is typically a vast amount of existing data collected from the internet 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 conditional generative model Go C + TT conditioned on concepts and initial states, which learns to generate future trajectories. We train the parameters of Go to maximize likelihood arg max, Epeennn [log Ge(Tle so)]-
dataset pretriin = {(T; c3)
Then, given an unlabeled demonstration dataset Dnew ~ D. we formulate learning new concept è that is used to sample trajectories from D as inverting the generative model. In particular, we learn new concept è so that our frozen conditional generative model Go maximizes the likelihood of trajectories in Deew. corresponding to arg max; E.wD.. [log s.)) We find that this design choice enables us to leverage the priors learned by Go from D pretmin to effectively learn concepts from Dnew given very few demonstrations, even if the demonstrated Drew deviates from the concept labels e seen in D pretrain- For evaluation in closed loop, we further assume access to a planner that given two states plans which action take in the environment, sometimes via access to simulation 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.
4 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 Iearn generative model Go from training {(T,, c))} pairs. We then freeze Go. and given demonstrations of new task {f). optimize a concept ē to produce the new behavior via Go. We then generate diverse set of behaviors via Go. either for the learned concept č conditioned on new initial states or for compositions of č with other concepts.
4.1 Training a diff'usion model to generate behavior
diffuusin model is a generative model that given a forward noise adding process ((1,[1;,1) == ((z v1 - 1. B.II) starting from data zo according to variance schedule $1...., 3r. learns the reverse process po (#t-1 == (Z He (2zz, t), E. (#,, f)). Ho et [61] simplify the training objective to estimate noise (4 "o..))).....e((..... [||e eg (x,, 1)] where x, is produced by adding noise to data zo by the forward noising process at diffusion step t, (=x))=0) :=N(xx, vār.o. (1 I) where āe := 11 (1 {.). Dhariwal and Nichol [62] enable conditioned generation by guiding the reverse noising process with classifier gradients. The noise prediction becomes ?= co(2,,t) t) w/1 a, V x, log (o(y22,, t) where classifier (sllaa t) is trained on noisy images, and w 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 é ep ( l) + w(en (2t, y, L) ca(s,, !)). Ramesh er al. [64] and Nichol et al. [65] demonstrate how this idea can used to generate images conditioned on class. Diffusion models have recently shown success 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 future state trajectory:
(1) where p the probability of removing conditioning information which is then replaced by dummy condition co. and 8p is the initial state corresponding to trajectory т. (x( is obtained from vo =T by the forward noising process. We then extend this approach for the inverse problem, namely, learning concept from demonstrations.
(prre)-DDommmee-d (0 (.)) [|| (o(2 -)) yoB, So. t)]l²)
4.2 Few-shot concept learning
Gal al. [11] use frozen generative model to learn visual concept representations from few images depicting the concept by optimizing the model's 's input v, = arg Big,, (T co(v), t)](2) where Cj 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 set of images that depict various concepts, for each image x' they learn a set of weights wf. and shared set of visual concepts for all images cj. é c(x), t) + Σk= w (e(x), Ck. t) e(z], t)). We extend these formulations to inferring multiple concepts, whose composition describes single task concept, from few demonstrations of task.
[[lc
Given a trained diffusion model co and demonstrations of a new concept (+); from Dnew: we learn concepts (a1 .... ...,ax) } for K 2 1 and their weights (wi, Wx] that best describe the demonstrations. Starting from uniformly sampled concept embeddings č₁ ~ (100 1]m),, we freeze co. and optimize ō% and wt:
E.. - (=a(<x(F), co, 80, 1) + wh (x)(())))) Cx s0, () - (F( co. 80, ()) (2)
1=1
We find that this compositional approach enables to effectively represent and learn new demonstra- tions, even when demonstrations are substantially different than those seen in training tasks.
4.3 Generating the learned concept
After learning concepts CR. whose composition describes the new task ē, we evaluate the behavior it generates by initializing (T(T) ~ (100, aI), and compute x, ~ ~N (H4-1: aE,, iteratively as a function of the estimated denoising function é(ea), where μ and Σ are the mean and variance that define the reverse process, and a E [0, 1) is a scaling factor that leads to lower temperature samples, until generating xo 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. coie-toee navigation). In contrast to other domains where the initial state is irrelevant for task due to the randomness in sampling (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 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 trajectory with an added another concept, we add another term to the sum in Eq. where the learned concept is composed with training concept e₁ and its weight wi: (x(e Ck. 80, () (x(xx( Cg. 80, 1)). This modified generation procedure constructs trajectories which exhibits behavior that has composition of the learned concept and the other specified concepts [14].
5

---

Training
New Concepts
single reloton between two objects
training relation compostions
novel concepts
frkangin right square
circle above square
friingl right sqjuene and circle above sepuare
liner circle right angle and triangle right of square
frisng diagonal to sqqueee
circle triingle and aquare on circle of rodiuss
Figure 4: Object rearrangement. Training concepts are single pairwise relations ('A right of/above B'), and new concepts are either compositions of training concepts ('A right of/above 'B right of/above C*) or new relations ('A diagonal to B", 'A, B, C on circle circumference of radius r').
squeeee diagonal to triangle
squaare diagonal to triangle
circle right of square
Figure 5: Object rearrangement new concept qualitative evaluation. Learning the new concept "square diagonal to triangle' and composing with the training concept 'circle right of square'.
Similarly to Ajay et al. [17]. in environments where an inverse dynamics model is provided, we generate trajectories 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 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 learn the concept from these demonstrations.
5.1 Task-Concepts
Learning concepts describing goals that are spatial relations between objects. Object rearrange- ment is common task in robotics [67-69] and embodied artificial intelligence (AI) [70, 71], serving as foundation a broader range of tasks such as housekeeping and manufacturing. Here, we use 2D object rearrangement domain to evaluate the ability of our method to learn task specification con- cepts. Given a concept representing relation between objects, we generate single state describing that relation. The concept in a training example describes the relation (cither 'right of" 'above") between only one pair objects (out of three objects) in the environment. Then, a model must learn compositions of these pairwise relations and new concepts such 'diagonal' and 'circle' (see Figure 4). The results in Figures and demonstrate that our method learns unseen compositions of training concepts and concepts. They further demonstrate how method composes 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 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 A triangle above circle' has low accuracy for 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 weights, they may converge 0 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 goal-oriented navigation domain adapted from the AGENT dataset [72], where an agent navigates to one of two potential targets. Conditioned on concept representing the attributes of the desired

---

1,00
0.8
Austary
0.25
0.6
0.50
0.25
0.2
0.00
Treining
Componiton
Traking
New Concepi
New and Triining Concrge Compostion
0.0
Training
"Thining Compostion
Training Compoition
Ours
VAE
Oure
BC
VAR
In-Coniextt
Figure 6: Object rearrangement (Ieft) and AGENT closed loop (right) quantitative evaluation on training and few-shot novel concept learning. Accuracy of FTL-IGM (ours), BC, VAE, and In-Context over concept generation of training concepts, novel compositions, novel concepts, and new initial states. We plot the average and standard error over new task types. Full details of the evaluation metries appear Appendix B and for baselines implementation Appendix C.
target object and initial state, we generate state-based trajectory describing an agent navigating to the target. Each object has 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 concepts defined by both target attributes, including (1) unseen color-shape target combinations and (2) new target-distractor combinations. Figure shows an example. training, see bowl and red object targets. A new concept includes novel composition the target red bowl. The new concept distractor objects (green bowl and red sphere) were introduced during training, but they were not paired with red bowl the target. As Figure 13 shows, method successfully learns concepts where targets are new compositions of target attributes in settings with new target-distractor pairs and generalizes new initial object states. We further evaluate our model and baselines in closed loop (Figure 6) by making an additional assumption that planner is provided. The planner produces action given a current state and future desired state predicted by model.
Learning concepts describing human motion. Unlike prior work on learning to compose human poses from motion capture (MoCap) data [c.g.. 73, 74]. here focus the inverse problem learning new actions from MoCap data. In particular, we use the CMU Graphics Lab Motion Capture Database (http://mocap.co www.mm..euuu We train 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 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 'breaststrokee) from new initial states and composes 'jumping jacks' with training concepts 'walk', 'jump', and 'march'.
Results are best viewed on website website. We compare motion generated by 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 a challenging multi-agent environment to complete a driving task. We train on several driving scenarios ('highway', 'exit', 'merge', and "intersection') and learn new driving scenario ('roundabout)) from several demonstrations (see Figure 8 and further details in Appendix B.4). We evaluate this scenario in closed loop new initial states, assuming access 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 novel concept for real-world table-top manipulation with Franka Research 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 book (Figure 10). We evaluate training pushing in closed loop and and achieve success
7

---

Begget defieed by single colonthacee aesibuuee
Sarae
Dutracts ves ooject
Durge Seftecciyy conpollcooottiiooooooo ofarritueee
Terget
@.- pnllonctoral cace
yullin
recisce 68481
Cnarpl cone
100 purpleccorn GPILALL.
Figure 7: Goal-oriented navigation. In training. targets are defined by single attribute (color or shape), and in new concepts, targets are defined by novel combination of color and shape attributes. To make the setting more challenging, distractor objects in new concept demonstrations have combination of attributes that are within the training distribution.
Setting
BC
VAE
Ours
Training
20 + 16.3
0 =
80 + 16.3 100 +
66.7 + 0
46.7 16.3
New Concept
26.7 = 13.3
0= 0
73.3 + 13.3
New Initial State
40 + 24.9
26.7
13.3
80 + 16.3
Table 1: MoCap human experiment. For training concepts and new concepts on new initial states, we report (top) percentage of time each method is the most successful at depicting a concept and (bottom) percentage of time each method depicts a concept. Mean and standard deviation are calculated over the number of scenarios in each setting and human subjects.
(90% accuracy). We evaluate the new concept, elevated pushing, against baseline that conditions on the training 'push green circle to orange triangle' representation in the new scenario setup where the objects are placed 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 Appendix B.5 and qualitative results are on our website.
5.2 Baselines
Goal conditioned behavior cloning. We compare method with goal-conditioned behavior cloning (BC), which, given a condition and state, outputs the next state in a sequence. It is trained 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, 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 9, 13). To achieve these generalization capacities, we need a model that can process interpolated (initial states) and composed (concepts) conditions, such our generative model.
Learning the concept space with VAE. We compare method with VAE [76] that does 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 decoder conditioned on an initial state and 2 with added noise. is obtained by encoding a trajectory generate given demonstration
for training evaluation and by fixing the decoder and optimizing when learning a new concept. We find that the VAE model learns latent space that captures training and new concepts but does not enable generalization to new initial states (Figures 9, 13). This highlights the importance of concept representations in the pretrained data.
Traning
Haghwwyy
merga
intersection
Figure 8: Autonomous driving. A controlled agent (green) completes various driving objectives as fast as possible while sharing the road with other vehicles (blue) and avoiding collisions. Training concepts: 'highway', 'exit', 'merge', and 'intersection' new concept: 'roundabout'.

---

60
20
0 Toining
Concept Crush
Twiining Succes
New Concrpe Saoore
VAE
In-Conicct
Figure 9: Driving crash (left) and success (right) rates, Crash rate (lower better) and task completion rate (higher is better) averaged over training tasks. We report standard errors training tasks and accuracy over 50 trajectories generated from the learned new concept. VAE has a high completion rate yet high crash rate. In-context has low crash rate yet 0% success the controlled vehicle reaches the roundabout's s center but does not complete the crossing. Overall, our method learns complete the roundabout crossing with competitive crash and success rates.
typically,
Training concepts
New concept closed loop evaluation
picispace book
push
push
Figure 10: Table-top manipulation. Training concepts: bock-and-place onto elevated surfaces and table-top pushing. New concept: pushing on an elevated surface.
In-Contes. learning. We compare our approach with training method to in-context learn from demonstrations. Specifically, we compare our approach Xu 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, 13). This emphasizes the need to learn explicit concept representations,
Conditioning Language Descriptions of New Concepts. There has been work on generating actions from language instructions [77-79]. We demonstrate that in our setup. merely providing new concept language instructions embedded with T5 (as in our pretraining dataset) 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 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.
5.3 Learning two concepts yields higher accuracy than concept
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
website.
5.4 How are learned new concepts related to training concepts?
New concepts that are compositions of training concepts. We analyze what the learned two concepts in Object Rearrangement and AGENT learn for novel concept compositions (e.g., 'red bowl'). For each concept (e.g., 'red' and 'bowl'), we generate two sets of 50 samples from the learned components. Table 3 shows accuracy for these sets over the concepts. some cases (most notably the 'line' concept in Object Rearrangement, 'circle right of triangle and triangle right of square"), each

---

Environment
Setting
1 Concept
2 Concepts
Training Composition
0.28
0.75:1..11
Object
New Concept
0.49:10.11
0.77:10.09
Rearrangement
New+Training
0.38++0.04
0.73+10.04
Training Composition
0.52:#0.08
0.68+0.13
AGENT
New Initial State
0.54-10.07
0.67:10.05
Driving
New Initial State
0.14
0.24
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 weight. We report average accuracy and standard error over task types for Object Rearrangement and AGENT. Driving includes single new concept and we report accuracy only.
Table 2:
learned component captures a single composed concept. In other cases, single learned component captures both concepts (Figure 12).
New concepts that are not explicit compositions of training concepts in natural language sym- bolic space. In Figure we visualize t-SNE [80] embeddings for T5 training concept representations and learned concept component representations. We note that learned components relatively close to training concepts, maintaining the model's 's input distribution, yet capture concepts that are not explicit compositions of training concepts.
Object Rearrangement
MoCap
Driving
training concapt
loorned new concapt
Figure 11: t-SNE embeddings of new concepts that not explicit compositions of training concepts. See interactive version for detailed labels on our website.
6 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 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 few examples and generalizes the learned concepts to unfamiliar scenarios. It also composes learned concepts to form unseen behavior thanks to the compositionality 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 general for any parameterized generative model, our implementation with 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 few demonstrations without retraining the model. We have approached the question of what new concepts can be represented by compositions of concepts 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 stronger ability to learn and generalize various task concepts in complex domains.
10

---

7 Acknowledgments
This research was also partly sponsored by the United States Air Force Research Laboratory and the United States Air Force Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750-19 2-1000. The views and conclusions contained this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes, notwithstanding any copyright notation herein. acknowledge support from ONR MURI under N00014-22 1-0000----------- and ARO MURI under Fellowship. We would like to thank Abhishek Bhandwaldar for help with the AGENT environment, and Anurag Ajay, Lucy Chai, Andi Peng, Felix Yanwei Wang, Anthony Simeonov and Zhang- Wei Hong for helpful discussions.
1I-NF-23-11 1--2-111-77 Yilun is supported in part by NSF Graduate Research
References
[1] Brenden M Lake, Tal Linzen, and Marco Baroni. Human few-shot learning of compositional instructions. arXiv preprint arXiv:1901.04587, 2019. I
[2] Archit Parnami and Minwoo Lee. Learning from few examples: A summary of approaches to few-shot learning. arXiv preprint arXiv:2203.0429) 2022.
[3] Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likeline log-likelihood function. Journal of Statistical Planning and Inference, 90(2):227-244, 2000. 1,
3
[4] Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, volume 8, pages 1438, 2008.
[5] Andrew Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Interna- tional Conference on Machine Learning, volume I, page 2, 2000. 3
[6] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. In International Conference on Learning Representations, 2018. 1,
[7] Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence mini- mization perspective on imitation learning methods. In Conference on Robot Learning, pages 1259-1277. PMLR, 2020. 1
[8] Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. Advances in Neural Information Processing Systems, 30, 2017. 1, 3, 25
[9] Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning In Conference on Robot Learning, pages 357-368. PMLR, 2017.
[10] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 35066266) 2015.
[11] Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. In International Conference on Learning Representations, 2023. 3,
[12] Nan Liu, Yilun Du, Shuang Li, Joshua Tenenbaum, and Antonio Torralba. Unsupervised compositional concepts discovery with text-to-image generative models. In International Conference on Computer Vision, 2023. 3,
[13] Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv: 00009, 2016. 2

---

[14] Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. arXiv preprint arXiv:2206.07714, 2022.
2, 5,5
[15] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8823 PMLR, 2021. 2
[16] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, volume 35, pages 36479-3400 -34944, 2022. 2
[17] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In International Conference on Learning Representations, 2023. 2, 5, 6, 24, 25
[18] Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. Advances in Neural Information Processing Systems, 32, 2019. 3
[19] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction no-regret online learning. In Proceedings of the International Conference on Artificial Intelligence and Statistics, pages 627-635, 2011.
[20] Minghuan Liu, Tairan Hc, Minkai Xu, and Weinan Zhang. Energy- imitation learning, In International Conference on Autonomous Agents and Multiagent Systems, 2021.
[21] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems, 29, 2016. 3,
[22] Dzmitry Bahdanau, Felix Hill, Jan Leike, Edward Hughes, Arian Hosseini, Pushmeet Kohli, and Edward Grefenstette. Learning to understand goal specifications by modelling reward. In International Conference on Learning Representutions, 2019.
[23] Justin Fu, Anoop Korattikara, Sergey Levine, and Sergio Guadarrama. From language to goals: Inverse reinforcement learning for vision-based instruction following. In International Conference on Learning Representations, 2019. 3
[24] Adam Gleave and Oliver Habryka. Multi-task maximum entropy inverse reinforcement learning. In Proceedings of the /st Workshop on Goal Specifications for Reinforcement Learning, 2018. 3
[25] Chris L Baker, Rebecca Saxe, and Joshua B Tenenbaum. Action understanding as inverse planning. Cognition, 113(3):329-3444 2009.
[26] Chris L Baker, Julian Jara- Ettinger,, Rebecca Saxe, and Joshua B Tenenbaum. Rational quantita- tive attribution of beliefs, desires and percepts in human mentalizing. Nature Human Behaviour, 1(4):0064, 2017.
[27] Tan Zhi-Xuan, Jordyn Mann, Tom Silver, Josh Tenenbaum, and Vikash Mansinghka. Online bayesian goal inference for boundedly rational planning agents. Advances in Neural Information Processing Systems, 33:19238-19252 2020.
[28] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In International Conference Machine Learning, pages 24631-24645, PMLR, 2022. 9,
[29] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? arXiv preprint arXiv:2209. 14935, 2022. 3
[30] Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In International Conference Machine Learning, pages 8302-8313, PMLR, 2021.
[31] Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with goal- conditioned policies. Advances in Neural Information Processing Systems, 32, 2019.
12

---

[32] John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder. Hierarchical reinforcement learning with trajectory embeddings. In International Conference on Machine Learning, pages 1009-1018. PMIR, 2018.
[33] Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Learning Representations, 2018
[34] Divyansh Garg, Skanda Vaidyanath, Kuno Kim, Jiaming Song, and Stefano Ermon. Lisa: Learning interpretable skill abstractions from language. Advances in Neural Information Processing Systems, 35:21711-2111 2022. 3
[35] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022. 3
[36] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine Planning with diffusion for flexible behavior synthesis. In International Conference Machine Learning, 2022. 3,
[37] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems, 2023. 3
[38] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning vía sequence modeling. Advances in Neural Information Processing Systems, 2021. 3
[39] Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefen- stette, and Yuandong Tian. Efficient planning in compact latent action space. International Conference on Learning Representations, 2023.
[40] Shyam Sudhakaran and Sebastian Risi. arXiv:2301.13573 2023.
Skill decision transformer. arXiv preprint
[41] Chiyu Jiang, Andre Cormman, Cheolho Park, Benjamin Sapp. Yin Zhou, Dragomir Anguelov, al. Motiondiffuser: Controllable multi-agent motion prediction using diffusion. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9644 9653, 2023.
[42] Yanjic Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu. 3d diffusion policy. arXiv preprint 24t.:::2000033954., 2024. 3
[43] Abhishek Padalkar, Acom Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310.08864 2023.
[44] Haoyu Zhen, Xiaowen Qiu, Peihao Chen, Jincheng Yang, and Xin Yan. Yilun du, yining hong, and chuang gan. 3d-vla: A 3d vision-language--ction generative world model. arXiv preprint 0rXiv:2403.0963/ 7(8), 2024.
[45] Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, al. Openvla: An open-source 2024.
vision-language-action model. arXiv preprint arXiv:2406.
[46] Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees. Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprint arXiv:2405.1 /2213, 2024.
[47] Irina Higgins, Loie Matthey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. International Conference on Learning Representa- tions. 3, 2017.
13

---

[48] Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representa- tion. arXiv preprint arXiv: /901. 1390, 2019. 3
[49] Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object repre- sentation learning with iterative variational inference. In International Conference on Machine Learning, pages 2424-2433. PMLR, 2019.
[50] Yilun Du, Shuang Li, Yash Sharma, Josh Tenenbaum, and Igor Mordatch. Unsupervised learning of compositional energy concepts. Advances in Neural Information Processing Systems, 34: 15608- 15608- 2021.
[51] Nan Liu, Shuang Li, Yilun Du, Josh Tenenbaum, and Antonio Torralba. Learning to compose visual relations. Advances in Neural Information Processing Systems, 34:23166-23178 2021.
[52] Weili Nie, Arash Vahdat, and Anima Anandkumar. Controllable and compositional generation with latent-space energy-based models. Advances in Neural Information Processing Systems, 34:13497-13510, 2021.
[53] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with eneryy-based diffusion models and meme. In International Conference on Machine Learning, pages 8489-8510 PMLR, 2023.
[54] Zihao Wang, Lin Gui, Jeffrey Negrea, and Victor Veitch. Concept algebra for (score-based) text-controlled generative models. Advances in Neural Information Processing Systems, 36, 2024.
[55] Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation and inference with energy based models. arXiv preprint arXiv:2004.060300 2020. 3
[56] Weixi Feng, Xuehai He, Tsu-Jui Fu, Varun Jampani, Arjun Akula, Pradyumna Narayana, Sugato Basu, Xin Eric Wang, and William Yang Wang. Training- free structured diffusion guidance for compositional tex--to-mmage synthesis. arXiv preprint arXiv:22 22t:20022..0 2022. 3
[57] Changhao Shi, Haomiao Ni, Kai Li, Shaobo Han, Mingfu Liang, and Martin Renqiang Min. Exploring compositional visual generation with latent classifier guidance. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 853 033-862,45 2023.
(58) Yuren Cong, Martin Renqiang Min, Li Erran Li, Bodo Rosenhahn, and Michael Ying Yang. Attribute-centric compositional text-to-image generation. arXiv preprint arXiv:2301.01413 2023.
[59] Lianghua Huang, Di Chen, Yu Liu, Yujun Shen, Deli Zhao, and Jingren Zhou. Com- poser: Creative and controllable image synthesis with composable conditions. arXiv preprint arXiv:2302.09778, 2023. 3
[60] Timur Garipov, Sebastiaan De Peuter, Ge Yang, Vikas Garg, Samuel Kaski, and Tommi Jaakkola. Compositional sculpting of iterative generative processes. Advances in Neural Information Processing Systems, 36: 2265- 2023. 3
[61] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020. 4
[62] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[63] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint 22X.:22077./2598 2022. 4,5
[64] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint MoXiv:2204.061222 1(2):, 2022. 4
14

---

[65] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Iya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprini arXiv:21 10741, 2021. 4
[66] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, (1):555555555, 2020.
[67] Ankit Shah, Pritish Kamath, Julie Shah, and Shen Li. Bayesian inference of temporal task specifications from demonstrations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems
31, pages 3804-3813. Curran Associates, Inc., 2018. 6
[68] Fujian Yan, Dali Wang, and Hongsheng He. Robotic understanding of spatial relationships using neural-logic learning. In 2020 JEEE/RSJ International Conference on Intelligent Robots and Systems, pages 8358-8365. IEEE, 2020.
[69] Ryan Rowe, Shivam Singhal, Daqing Yi, Tapomayukh Bhattacharjee, and Siddhartha S Srinivasa. Desk organization: Effect of multimodal inputs on spatial relational learning. In 2019 28th IEEE International Conference on Robot and Human Interactive Communication, pages 1 1-8. IEEE, 2019. 6
[70] Xavier Puig, Tianmin Shu, Shuang Li, Zilin Wang, Yuan-Hong Liao, Joshua B Tenenbaum, Sanja Fidler, and Antonio Torralba. Watch-and-helpp A challenge for social perception and human-ai collaboration. arXiv preprint arXiv:2010.09890. 2020.
[71] Sanjana Srivastava, Chengshu Li, Michael Lingelbach, Roberto Martin-Martín, Fei Xia, Kent El- liott Vainio, Zheng Lian, Cem Gokmen, Shyamal Buch, Karen Liu, al. Behavior: Benchmark for everyday household activities in virtual, interactive, and ecological environments. In Conference Robot Learning, pages 477-490. PMLR, 2022. 6
[72] Tianmin Shu, Abhishek Bhandwaldar, Chuang Gan, Kevin Smith, Shari Liu, Dan Gutfreund, Elizabeth Spelke, Joshua Tenenbaum, and Tomer Ullman. Agent: A benchmark for core psychological reasoning. In International Conference on Machine Learning, 2021. 6, 21, 22
[73] Borui Wang, Ehsan Adeli, Hsu-kuang Chiu, De-An Huang, and Juan Carlos Niebles. Imitation learning for human pose prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7124-7133, 2019. 7
[74] Guy Tevet, Sigal Raab, Brian Gordon, Yonatan Shafir, Daniel Cohen-Or, and Amit H Bermano. Human motion diffusion model. International Conference on Learning Representations, 2023.
7
[75] Edouard Leurent. An enviromment for autonomous driving (eeisiionmmaing https://gtthh:ti com/
2018. 7, 23
[76] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint
1312.6 14, 2013. 25
[77] Oier Mees, Lukas Hermann, Erick Rosete-Beat and Wolfram Burgard. Calvin: A benchmark for Ianuuage-condtiione policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters, 73))733777333, 2022. 9
[78] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.068177 2022.
[79] Corey Lynch, Ayzaan Wahid, Jonathan Tompson, Tianli Ding, James Betker, Robert Baruch, Travis Armstrong, and Pete Florence. Interactive language: Talking to robots real time. IEEE Robotics and Automation Letters, 2023,
[80] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9(11), 2008. 10
15

---

[81] Klaus Greff, Francois Belletti, Lucas Beyer, Carl Doersch, Yilun Du, Daniel Duckworth, David Fleet, Dan Gnanapragasam. Florian Golemo, Charles Hermann, al. Kubric: A scalable dataset generator. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3749-3761, 2022. 22
[82] Julian Tanke, Chintan Zaveri, and Juergen Gall. Intention- based Long-Term Term Human Motion Anticipation. International Conference 3D Vision 2021. 23
[83] Edouard Leurent. rl-agents: Implementations of reinforcement learning algorithms. https:
/ //github. com/ conm/eleurent//laggentss 2018. 23
[84] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 770-778, 2016. 25
[85] Dean Pomerleau. Alvinn:
An autonomous land vehicle in neural network. In D.S. Touretzky, editor, Proceedings of Neural Information Processing Systems, pages 305 313. Morgan Kaufmann, December 1989. 25
[86] Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. arXiv preprint arXiv:1711.05101, 5, 2017. 25
16

---

A Additional Results
Object Rearrangement. We display additional states generated by model for various concepts. We generate learned concepts that are compositions of training concepts (Figure 14), new concepts that not explicit compositions of training concepts (Figure 15), and concept composed with a training concept (Figure 16). We analyze the learned concepts qualitatively (Figure 12) and quantitatively (Table 3).
Goal-Oriented Navigation. provide accuracy for open loop evaluation in Figure 13, We further analyze the learned concepts quantitatively (Table 3).
Component 1
Component
circle right of triangle and triangle right of square
square right circle and triangle above circle
Figure 12: Object rearrangement new compositions analysis qualitative evaluation. 'circle right of triangle triangle right square' (top) each learned component corresponds single composed concept: 'circle right of triangle' (component 1) and 'triangle right of square' (component 2). In "square right of circle triangle above circle" (bottom), learned component corresponds both composed concepts and component to none.
L,00
0.75
Acturtty
0.25
0,00
Training
Training Compositon
Training Compostion
VAE
Coment
Figure 13: AGENT open loop evaluation. We plot the average and standard error over task types.
B Data Generation and Evaluation
B.1 Object Rearrangement
Training. The training dataset of ~ 11k samples consists of concepts 'A right of B' and 'A above B' where A and B are one of three objects: circle, triangle, or square. Altogether, there are 12 possible concepts (two relations and three objects where order is important). In the data generation process, A at center position (x)) E [0,5] 5|2 with radius r'A E [0.3, 1] and angle 0 E [0, 2n is considered 'right of" B at center position (x B- yB) with radius r'g if 2 Fg and ly yol r Similarly, A considered 'above' B if yA yB and IAA Fg S TA- We further verify that training objects do not overlap.
New Tasks.
The new scenarios include
Five novel compositions of training concepts (training composition Figure 6). 'triangle right of square A circle above square', 'square right of triangle A circle above triangle',
17

---

triengle right of and circle ebove
square right of triengle and circle ebove triengle
circle right square and triangle above equare
Squara right circle triangle
line: circle right of triengle and triangle right of square
Figure 14: New concepts that training concept compositions. We display successful (green
frames) and unsuccessful (red frames) states generated by our model conditioned on learned concepts that are training concept compositions.
18

---

triangle circle red.
triengle diagonal squere
circle diagonal triangle
trianglc diagonal circle
Figure 15: New concepts that are not explicit training concept compositions. We display successful (green frames) and unsuccessful (red frames) states generated by model conditioned on learned concepts that are not explicit training concept compositions.
19

---

square diagonal triangic
circle above square
squere diagonal to triangle
circle above triangle
square diagonal triangle
circle right triangle
square diagonal triangle
triangle
Figure 16: New concept composed with training concepts. We display successful (green frames) and unsuccessful (red frames) states generated by our model conditioned learned concept that is not an explicit training concept composition ('square diagonal to triangle") in composition with various training concepts.
20

---

Table 3: Analysis of learned concepts. For concepts that explicit training concept compositions, we evaluate what each learned component captures. For both concepts and 2, we report accuracy with respect to the concept when trajectories generated solely from one learned component. e.g. in AGENT (bottom), 42% of the trajectories generated by component and 48% of the trajectories generated by component target the red object.
Environment
Concept
Component
Component
Concept
Component
Component 2
ARD
0.42
0.62
OAC
0.88
0.36
DRA
0.16
0.64
OAA
0.08
0.08
Object Rearrangement
ORT
0.1
0.48
AAT
0.86
0.4
ORC
0.14
0.9
AAO
0.02
0.6
ORA
0
0.92
ARD
0.64
0.14
red
0.42
0.48
bowl
0.66
0.56
yellow
0.72
0.44
howl
0.58
0.74
AGENT
red
0.42
0.76
cube
0.44
0,96
yellow
0.4
0.32
cube
0.82
0.78
purple
0.48
0.42
cone
0.56
0.36
"circle right of square A triangle above square', 'square right of circle triangle above circle', 'line": 'circle right of triangle A triangle right of square'. There are five demonstrations of each novel composition.
Five new concepts (new concept in Figure 6). 'circle': circle, triangle, and square lie in equal intervals the circumference of circle with radius 1.67 and center E [0, 5)² 'square diagonal to triangle', 'triangle diagonal to square', 'circle diagonal to triangle", 'triangle diagonal to circle'. A is considered diagonal to B if their centers lie on f(xx "above' and 'right of" B. There are five demonstrations of each novel concept.
and if A is
Composing a new concept, 'square diagonal to triangle', with five training concepts (new and training concept composition in Figure 6): 'circle right of square', "circle above square', 'circle above triangle', 'circle right of triangle', 'triangle above circle'. When we learn weights for the new concept, the training concept weighted by 1. Otherwise, both training and learned concepts are weighted by fixed w.
We verify that objects do not overlap and that no training relations unrelated the specified new concepts exist in the demonstrations between objects.
State space. The state space a 21-tuple describing three shapes (circle, triangle, and square), each represented by 7-tuple: their center 21) position, size, angle, and one-hot shape type.
Evaluation data and metries. For training, generate 50 states conditioned uniformly sampled training concept embeddings for single relations. For each concept composition and new concept, we report accuracy on 50 generated states conditioned on the learned concept. For new concepts composed with training concepts, we report accuracy based on 50 generated states for each training concept. Results in Figure are reported for best learned w and fixed w {1.2, 1.4, 1.6, 1.8) we learn two concepts with classifier- free guidance weight w = 1.2. We report accuracy based on relaxed version of the data generation process: A is considered 'right of B if æ Tg and Iya WA-yBl Σ 2-max(ra., TB} A is considered 'above' if yA > }B and 1eA-EBl < 2- 2-maxfra. rg), A. B and C lie on 'circle' if their centers form a circle of radius where |r 1.67 < 0.3 and A is 'diagonal' to FA > BB. A > yB and they lie within a max(r TB) margin of fx) = x.
B.2 Goal-Oriented Navigation
Training. To collect demonstrations, we follow the data generation process in the AGENT bench- mark environment [72], which provides a planner for navigation given the desired target. In the provided environment, the agent's initial position ar=o (0, 0. 102, -3.00)) color (yellow) and shape (cone) are fixed. There are two objects: the target and distractor. Each object has a color of. of E (red, yellow, purple, green). shape of, of E (cone, sphere, bowl, cube) and a position of E [0, 1.66) x {0.102) x -4.355, 3.2557, of E [-1.66, 0] x {0.102) x [-4.555, 3.2557
21

---

The training dataset of ~ 900 samples consists of concepts 'go to red object' and 'go to yellow object' where of E fred, ye og E yellow) \ {of] and of oz sphnee,, and concepts "o to bowl' and "go to cube" where of E (bowl, cube), of {bowl, cube) \ {of} and of. og E ]reen).
New Tasks.
The new scenarios include
Five novel compositions of training target color and shape attributes (training composi- tion in Figure 6). 'go to red bowl' (o red, of bowl, of E sphere) or red, of E (purple, green] where "go yellow bowl' (of, bowl, of - yellow, of {purple, (pprple,green]). 'go to red cube' of red, of cube, of E sphner) or of.of of cube, of red, of (purple,green) green)). 'go yel- low cube' (of,of yellow, of cube, of E {cone, sphere) or 0f cube, yellow, of (purple, (purple,grenn)). 'go purple cone' (f = purple, of cone, of {bowl, cube} or 03 = cone, of = purple, of E {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.
of, 0f bowl,
{1,2] and j {1,2) i),
yellow, of bowl, 03 sphere) or 0f
Conditioning each novel composition concept novel initial states sampled from the novel concept distribution (training composition new initial state Figure 6).
State space. The initial state space we condition on is a 52-tuple based on the state space in Shu al. [72] describing the agent and two objects. The agent is represented by its 3D position. quaternion E R velocity E RR, angular velocity E R3 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 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; S 150 and are padded to length 150 using the final state. During training, we sample the demonstrations generate subtrajectories of length 128.
Evaluation data and metries. 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 w and fixed w E (1.2, 1.4, 1.6, 1.8] we learn two concepts with classifier-free guidance weight w 1.6. We report accuracy based on whether the agent in the generated trajectory has made progress towards the desired target (Ia==18 ourgell < ||[_ Raraal) ). Accuracy for closed loop evaluation (Figure 6) is reported based on optimally reaching target Ouaget before a distractor Odistmator ((0f=128
Gurget < 0.365 A AI < 128 loβ
Odstractoor 2 0.365).
Closed loop evaluation. The action space AA C R2 represents forces applied to the agent. We take actions a, T++s s, where s; is an observation in the environment, and Te+s future step in the open loop plan generated by G conditioned on observation 8, and learned concept č. We execute actions in the environment until reaching the target or maximum number of steps and report the evaluation metric described above. Closed loop evaluation is done in pybullet simulation [8] for efficiency.
B.3 MoCap
Training. All human actions in the CMU Graphics Lab Motion Capture Database (http://mocap.cs.coo www.occcpp....c..... www.ws..eccuuu.. 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.
22

---

New Tasks, The new scenarios include
Three new concepts. jumpping (3 demonstrations), 'chop wood' (2), and "breaststroke'
(3).
Conditioning each novel concept on novel initial states the last states in the new concept demonstrations.
Composing new learned concept jumpiig 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 14 and to remove rotation and translation. The demonstration state space is the same with horizons H,, i E [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 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 w and fixed w E {1.2, 1.4, 1.6, 1.8] - we learn two concepts with classifier-free guidance weight w = 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 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.
B.4 Autonomous Driving
Training. We use the following scenarios from the HighwayEnv driving simulation environment [75]. The training scenarios include four driving scenarios: 'highway', 'exit', 'merge', and 'intersec- tion'. In 'highway', the objective is driving on a highway 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 four-lane highway with an exit lane and 20 vehicles. The controlled vehicle 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 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 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 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 non- controlled vehicles are sampled. An episode ends if the controlled vehicle crashes time
23

---

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 IR55,, the controlled vehicle and the four closest vehicles. Each vehicle represented by 7-tuple, including whether it is present on the road, its and y positions, and y velocities, and cosine and sine heading directions. The demonstration state space observations of the controlled vehicle in R for horizons H,, i [N]. We sample trajectories to generate subtrajectories of length 8, which we train on.
Evaluation data and metries. 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 't have a success score as there is no final state to reach in this scenario. Results are reported for best learned w and fixed w E {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 the right lane, drive faster, drive slower. We assume access to a planner P that given two states plans which action take the environment via access to simulation in the environment, a = PT 8t). The planner simulates the possible actions from observed state 8; and randomly selects an action out of the ones closest to the model predicted next state Te+1 that does not result in the controlled vehicle crashing. We execute actions until reaching maximum number of steps or until the controlled vehicle crashes.
B.5 Toble-Top Manipulation
Training. We collect 214 expert demonstrations with a Franka Research 3 robot via telcop 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 fask. The new scenario includes pushing the green cirele the orange triangle on 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, 7-tuple of its 3D position and quatermion, and gripper state in E. We verify that most states are fully observable. The demonstration state space includes observations of the end effector pose and gripper state in R$ for horizons H,, i E [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 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 w 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 st, we assume access 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 maximum number of steps H 160.
C Implementation Details
Diffusion model. We represent the noise model ep 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 temporal U-Net where images are processed
24

---

by pretrained resnet18 [84] that is finetuned with the model during training. We use the same hyperparameters Ajay al. [17] except for the probability of removing conditioning information, which we set to 0.1. In Table 4 we demonstrate the effect of choosing different classifier-free guidance weights w. We are overall better or comparable to the baselines.
Table 4: Classifier-free guidance weight choice effeet. For Object Rearrangement and Goal- Oriented Navigation we report the accuracy and standard error of the mean for new concepts from new initial states, and for Driving, the success and crash rates. We compare the four baselines with our approach reported in Figures 6, 9 and 13 and Section 5.2. We report results for our approach with all w in our hyperparameter search and mark the reported Figures 6, 9 and 13 in bold font.
Domain
BC
VAR
In-Context
I Language
3
Ours
Object
0.67 1 0.15
0.08 0.05
0.2 0.05
1.2
0.82 0.09
Rearrangement
1.4
0.8 1 0.1
1.6
0.66 1 0.08
1.8
0.61 = 0.06
learned
0.77 0.09
AGENT
0.46 + 0.05
0.6 + 0.06
0.57 + 0.03
0.63 + 0.07
1.2
0.64 + 0.08
1.4
0.66 0.05
1.6
0.73 + 0.07
1.8
0.63 = 0.06
learned
0.67 = 0.05
Driving Success Rate
4%
24%
0%
1.2
18%
1.4
14%
1.6
4%
1.8
10%
learned
24%
Driving 64%
48%
16%
1.2
38%
Crash Rate
1.4
40%
1.6
48%
1.8
40%
learned
32%
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, average results over 50 different seeds. We an MLP with one hidden layer of size 512, ReLU activations, and AdamW [86] with learning rate 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 predict the next partially observable agent state (13-dim). Composing concepts implemented by adding conditions.
VAE. We implement a conditional variational autoencoder model (CVAE) [ We use an MLP for the architecture with AdamW [86] and learning rate 6- 10-1 We sample 50 trajectories per concept and average the results. Composing concepts is implemented by adding conditions.
In-Context. For the In-Context learning baseline [8, 28], architecture with 4 layers and 4 heads, where multiple demonstrations are passed into the encoder with vector separate consecutive demonstration trajectories. During training, there five prompt demonstrations, during training evaluation, we provide one demonstration, and for new concepts, we provide two five demonstrations, depending on the number of demonstrations provided in each domain. The decoder is passed a window of the previous states and predicts the next state. In AGENT and in MoCap K 2. We simplify the learning objective
a transformer encoder-decoder
and Driving, we use window size K to negative log-likelihood loss and convert the data to 20 bins for AGENT and Driving and 16 bins for MoCap. We further simplify the AGENT data by sampling the trajectories every eight states. We use learning rate 10 in all domains. We do not evaluate on Object Rearrangement since the horizon one.
based the implementation in http:///wttbboooooooooooooooo/cd
25

---

Language. The language instructions are as follows. For Object Rearrangement, training composi- tions are training concept descriptions composed with the word 'and'. New concepts are described as 'circle triangle and square in circle of radius 1.67' and 'A right and above diagonally to B' where A, BE ('circle", "triangle', 'square"). New concepts composed with training concepts are composed with the word 'and'. AGENT concepts are described based on new concept attributes, e.g. 'go to red bowl'. In MoCap, descriptions are human actions such as 'jumping jacks'.
Compute. We run all simulated experiments a single NVIDIA RTX A4000 machine. We evaluate our method on real-world table-top manipulation tasks using Franka Research robot with overhead Realsense D435I RGB camera and NVIDIA RTX 4090 machine. Concept learning can take approximately one to two hours. Since a new concept only has to learned once to generate behavior, in the domains we presented, it is reasonable to learn concept offline, and therefore, it is not prohibitive.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We describe our formulation in Section 3, the adapted method in Section 4 and empirical evaluation in Section 5,
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
27

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Described in Section 5 and Appendix and C.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer:
[No]
Justification:
We plan to release our code in the near future.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nipp cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /Codeeuomiisseeeeeeee
/nips. wwtpccccpppcccaeccccccccccccccat / assssnolicyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We discuss implementation details in Appendix and C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide errors in Figures 6, 9 and 13. Plots without errors display accuracy over a fixed set. Similarly for Tables 5.1, 2 and 3.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
29

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative emror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Described in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:://wwrrppscccpppiieeeeeeeee hcpppp/ic///hthicctttueeeeeeeee ines?
Answer: [Yes]
Justification:
We discuss limitations Section and obtain an IRB for small-scale human experiment.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: In this paper, demonstrate that effectively learning behavior (or concept) representations from few demonstrations is possible given a pretraining dataset, dependent crucially the quality of training and new concept demonstrations, While such effective few-shot learning system may have many societal implications, none are specific to this work that feel need 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 why the paper does not address societal impact.
30

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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 package should be provided. For popular datasets, paperswi wwwode.hhoe...c. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
31

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We not release any new assets. In the future, we plan to release our code and add documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [Yes]
Justification: We conduct a small scale human experiment described in Section 5, Table 5.1 and Appendix B.3.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [Yes]
Justification:
IRB approval was obtained.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
32

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "LxjJiPOD5q", "reviewer": "Reviewer_gXQZ", "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.", "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.\n\n[3] https://research.google/blog/inceptionism-going-deeper-into-neural-networks/", "limitations": "The limitations of the method are addressed.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The authors propose a novel method that learns concepts by optimizing the input concept embedding to a pre-trained generative model.\n\n2. The method also allows for composable concepts based on prior work on composable diffusion models.\n\n3. The method is evaluated in a diverse set of environments, the baselines the authors use are reasonable.", "weaknesses": "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.\n\n2. 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.\n\n3. 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.\n\nReferences:\n\n[1] Invertible Residual Networks. Behrmann et al. 2018.\n\n[2] Analyzing inverse problems with invertible neural networks. Ardizzone et al. 2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review and positive feedback.\n\n>is there something specific about the diffusion model that makes it “invertible”?\n\nNo, we state in Limitations that *“our framework is general for any parameterized generative model”*. We choose to implement our framework with a diffusion model given its success on interpolation and compositionality, both in learning compositions of concepts and in generating compositions of concepts. We describe these properties in the Introduction.\n\n>the community uses the term “invertible networks”\n\nWe refer to inverting the generative model [4-6] as opposed to inverting the model.\n\n>Equation 2 could be explained better so that the paper is more self-contained.\nThe methods section is very short; a background section could prepare readers to understand it better.\n\nWe will extend the method section to include the formulation and objectives of diffusion models [7], conditional diffusion models [8, 9], classifier-free guidance [10], and concept learning in computer vision [11], including for multiple visual concepts [12].\n\n>It is unclear if this approach learns meaningful concepts when used with an image or a video generative model.\n\nWe evaluate our method's capability to learn a novel concept for **real-world table-top manipulation with a Franka Research 3 robot conditioned on RGB images**. We train on a suite of tasks (table-top pick-and-place onto elevated surfaces, and table-top pushing scenarios) and learn a new task that requires high precision (pushing on an elevated surface) from ten demonstrations, **see Figure 15 in new pdf**. We evaluate in closed loop and achieve success rate of 0.9 on training pushing, **0.55 success on the learned new concept**, elevated pushing, and surpass a **baseline** that conditioned on the training pushing concept, achieves **0.15 success rate** in the elevated pushing setup. For more details please see general response.\n\n>related work section would benefit from further discussion of generative models and composable representations.\n\nWe will add the following discussion to the related work:\n\nGenerative Models in Decision Making. The success of diffusion policy in predicting sequences of future actions has led to 3D extensions [13], and combined with ongoing  robotic data collection efforts [14] and advanced vision and language models, has led to vision-language-action generative models [15,16].\n\nComposable representations. There has been work on obtaining composable data representations. $\\beta$-VAE [17] learns unsupervised disentangled representations for images. MONet [18] and IODINE [19] decompose visual scenes via segmentation masks and COMET [20] 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 [21-26] and robotic skills [27]. The generative process can also be altered to generate compositions of visual [28-31] and molecular [32] concepts. We aim to obtain task concepts and generate them in composition with other task concepts.\n\n>Is unconstrained gradient descent with respect to the input concept embedding the right way to learn concepts?\n\nConcept learning is not unconstrained. Unlike in deep dream, in this work we do not maximize a particular neuron, instead we optimize the input concept for generating the demonstrated new concept, therefore this optimization is supervised.\n\n\n[4] Action understanding as inverse planning. Baker et al. 2009\n\n[5] Rational quantitative attribution of beliefs, desires and percepts in human mentalizing. Baker et al. 2017\n\n[6] Online bayesian goal inference for boundedly rational planning agents. Zhi-Xuan et al. 2020\n\n[7] Denoising Diffusion Probabilistic Models, Ho et al. 2020\n\n[8] Diffusion Models Beat GANs on Image Synthesis, Dhariwal & Nichol 2021\n\n[9] Hierarchical Text-Conditional Image Generation with CLIP Latents, Ramesh et al., 2022\n\n[10] Classifier-Free Diffusion Guidance, Ho & Salimans, 2022\n\n[11] An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion, Gal et al. 2023\n\n[12] Unsupervised Compositional Concepts Discovery with Text-to-Image Generative Models. Liu et al. 2023\n\n[13] 3D Diffusion Policy. Ze et al. 2024\n\n[14] Open X-Embodiment: Robotic Learning Datasets and RT-X Models. 2024\n\n[15] 3D-VLA: A 3D Vision-Language-Action Generative World Model. Zhen et al. 2024\n\n[16] Octo: An Open-Source Generalist Robot Policy. Ghosh et al. 2024\n\n[17] $\\beta$-VAE: Learning basic visual concepts with a constrained variational framework. Higgins et al. 2017\n\n[18] MONet: Unsupervised Scene Decomposition and Representation. Burgess et al. 2019\n\n[19] Multi-Object Representation Learning with Iterative Variational Inference. Greff et al. 2020\n\n[20] Unsupervised Learning of Compositional Energy Concepts. Du et al. 2021\n\n[21] Learning to Compose Visual Relations. Liu et al. 2021\n\n[22] Compositional Visual Generation with Composable Diffusion Models. Liu et al. 2023\n\n[23] Controllable and Compositional Generation with Latent-Space Energy-Based Models. Nie et al. 2021\n\n[24] Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC. Du et al. 2023\n\n[25] Concept Algebra for (Score-Based) Text-Controlled Generative Models. Wang et al\n\n[26] Compositional Visual Generation with Energy Based Models. Du et al. 2020\n\n[27] Is Conditional Generative Modeling All You Need For Decision-Making? Ajay et al. 2023\n\n[28] Training-Free Structured Diffusion Guidance For Compositional Text-To-Image Synthesis. Feng et al. 2023\n\n[29] Exploring Compositional Visual Generation with Latent Classifier Guidance. Shi et al. 2023\n\n[30] Attribute-Centric Compositional Text-to-Image Generation. Cong et al. 2023\n\n[31] Composer: Creative and Controllable Image Synthesis with Composable Conditions. Huang et al. 2023\n\n[32] Compositional Sculpting Of Iterative Generative Processes. Garipov et al. 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VZ4GQoQ1Iv", "reviewer": "Reviewer_1rf8", "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.", "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": "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.", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. **Important Research Problem:** The paper addresses a critical problem in reinforcement learning—learning agent intents from limited data. This is a significant challenge with practical implications in various domains.\n   \n2. **Effective Visualization:** The paper provides effective visualizations, making the results and methodology easier to understand and interpret. The accompanying website further enhances the clarity and impact of the findings.\n\n3. **Generalizability:** The approach demonstrates the ability to learn novel concepts and generate corresponding agent plans or motions in unseen environments and in composition with training concepts. This suggests a promising level of generalizability.", "weaknesses": "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.\n\n2. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and suggestions for clarifying the paper.\n\n>The concept of \"task concept learning\" is not rigorously defined.\n\nIn the Formulation section we define:\n\n- task concepts as latent representations in $\\mathbb{R}^n$.\n- task concept learning as $argmax_{\\tilde{c}}{\\mathbb{E}}_{\\tau\\sim D _{\\mathrm{new}}}[\\log\\mathcal{G} _{\\theta}(\\tau|\\tilde{c},s_0)]$.\n\nWe state this in the 4th paragraph of the Introduction as well: *“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.”*\n\nIntuitively, task concept learning means inferring the intent of a demonstration and can be formally defined in various ways such as policies, rewards or trajectories, as we discuss in the 2nd paragraph of the Introduction. We will edit the 3rd paragraph of the Introduction to make clear what we mean by task concept learning earlier on in the paper.\n\n>The distinction between \"concept\" and \"context\" or \"contextual\" is unclear.\n\nConcept refers to a latent task representation that we learn. Context is **only** used in relation to *In-Context* [1], a baseline we compare with. We do not use context as a technical term in our formulation. We will change \"in context\" appearances in the paper to \"in-context\" for clarity.\n\n>The scalability of the pretrained model to diverse and complex real-world environments is not well addressed, raising concerns about its practical applicability\n\nWe evaluate our method's capability to learn a novel concept for **real-world table-top manipulation with a Franka Research 3 robot**. We train on a suite of tasks (including table-top pick-and-place onto elevated surfaces, and table-top pushing scenarios) and learn a new task that requires high precision (pushing on an elevated surface) from ten demonstrations, **see Figure 15 in the new pdf**. We evaluate in closed loop and achieve success rate of 0.9 on training pushing, **0.55 success on the learned new concept**, elevated pushing, and surpass a **baseline** that conditioned on the training pushing concept, **achieves 0.15 success rate** in the elevated pushing setup. **For more details please see the general response**.\n\n>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?\n\nThe paper compares against **four baselines**. As described in the Experiments section, **in addition to BC and VAE, we compare with In-Context and Language baselines**. The **In-Context** baseline [1] (**Figures 5,8 & 10 and https://sites.google.com/view/FTL-IGM/home**) represents tasks with demonstrations and generates behavior conditioned on those demonstrations. For new concepts, behavior is generated in a zero-shot manner by conditioning on the new demonstrations. For the **Language** baseline (**Section 5.2, *Conditioning on Language Descriptions of New Concepts*** **and https://sites.google.com/view/FTL-IGM/home**), we do not assume access to new concept demonstrations and instead represent the new task with a T5 embedding of its language description (as we do for labeled training concepts). We input this representation to our model to zero-shot produce new behavior.\n\nOur work is in the field of learning from demonstrations. As described in the Related Work section, the main approaches in this field are BC, IRL, Inverse Planning and In-Context learning. Specifically, we focus on few-shot task representation learning from demonstrations. Each baseline highlights a different aspect of our approach.\n\n- The **BC** baseline highlights the generative aspect of our approach. BC is **autoregressive** (predicts one state at a time given past states and a latent task representation) whereas our approach is generative (predicts 𝐻 future states from an initial state and latent task representation).\n- Similar to our approach, VAEs are generative models. The **VAE** baseline highlights our choice of representing tasks as T5 embeddings during training. VAE does not utilize these representations during training and instead **learns latent task representations** from demonstrations in a self-supervised manner.\n- The **In-Context** baseline [1] is **autoregressive** and **represents tasks as demonstrations**.\n- The **Language** baseline utilizes our **generative** model in an **In-Context** fashion and highlights the need for demonstrating new concepts.\n\nIRL is not relevant for few-shot learning as it learns only one task representation (reward or policy) at a time and assumes access to taking actions in the environment during training. Inverse Planning is not suitable for our case as it assumes knowledge about the task space. We also discuss why fine-tuning task conditioned BC is not possible — we assume no access to new concept task representations, only to their demonstrations.\n\n[1] Prompting decision transformer for few-shot policy generalization. Xu et al. 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UonzzCSllo", "reviewer": "Reviewer_sy18", "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.", "questions": "1. Are concept representations updated during training or are they fixed (as the T5 embeddings)? \n2. How important is the classifier-free guidance weight? Does it require task-specific tuning?", "limitations": "Limitations are present in section 6.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The proposed method presents a possible solution to the few-shot learning task, although the method may have been applied in computer vision applications [1,2]\n2. Experimental results suggest that by inverting a generative model to learn task concepts, the method works better than goal-conditioned or in-context learning-based baselines. \n3. The method is clearly presented in the manuscript. \n\n[1] Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. In International Conference on Learning Representations, 2023. 3\n[2] Nan Liu, Yilun Du, Shuang Li, Joshua B Tenenbaum, and Antonio Torralba. Unsupervised compositional concepts discovery with text-to-image generative models. In International Conference on Computer Vision, 2023. 1, 3", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review and positive feedback.\n\n>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. \n\nThe Language and In-Context baselines perform poorly in cases where the new concept is not an explicit composition of training tasks in natural language symbolic space, such as new object rearrangement concepts, new human motions and a new driving scenario. In these cases we demonstrate the benefit of providing demonstrations of a new task that require some learning but can then be used in various scenarios: (1) we keep the same concept latent representation, and change the initial state; (2) we combine the concept latent representation with other concept representations to form novel, compositional concepts. All of these can not be achieved without finding the concept representation. Since a new concept only has to be learned once to generate behavior, in the domains we presented, it is reasonable to learn a concept offline, and therefore not prohibitive.\n\n>Are concept representations updated during training or are they fixed (as the T5 embeddings)?\n\nRepresentations are fixed during training.\n\n>How important is the classifier-free guidance weight? Does it require task-specific tuning?\n\nIn Table 4 (see the new pdf) we report results for various classifier-free guidance weights compared with the baselines. Overall for various choices of $\\omega$ we are better or comparable to baselines. We find that learning the weights in all domains is a reasonable choice, and different domains may further benefit from tuning $\\omega$ to a fixed value.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their feedback and for acknowledging that learning agent intent from limited data is an important problem, and that our method is evaluated in a diverse set of environments, surpassing baseline performance and demonstrates generalizability, including compositionality. Moreover, we thank the reviewers for acknowledging that the paper is clear and provides effective visualizations. 

We describe additional experiments (please see new pdf) here to address reviewers’ comments, and address individual comments in response to each reviewer below.

# **Real-World Experiments**

We evaluate our method's capability to learn a novel concept for **real-world table-top manipulation with a Franka Research 3 robot**. We **train** on a suite of tasks (including **table-top pick-and-place onto elevated surfaces, and table-top pushing scenarios**) and learn a **new task** that requires high precision (**pushing on an elevated surface**) from ten demonstrations, see **Figure 15** in new pdf. We **evaluate in closed loop** and achieve success rate of **0.9** on **training** pushing, **0.55** success on the **learned new concept**, elevated pushing, and surpass a **baseline** that conditioned on the training pushing concept, achieves **0.15** success rate in the elevated pushing setup.

Specifically, 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). The new concept includes pushing the green circle to the orange triangle on a book. We provide ten demonstrations of this task. Demonstrations have horizons $H_i,\,i\in[N]$ which we split into subtrajectories of length 32. Our generative **model learns to predict the next 32 states**, each composed of the **end effector pose**, a 7-tuple of its 3D position and quaternion, **and gripper state** $\in\mathbb{R}^8$, **conditioned on** an overhead **RGB image** of the scene, and the **robot's current end effector pose and gripper state**. We verify that most states are fully observable.

For closed loop evaluation, given the predicted state $\tau_{t+1}$ and current end effector pose $s_t$, we assume access to a planner that linearly interpolates between them. We repeatedly make predictions and roll out the 32 predicted states in closed loop until success or upon executing a maximum number of steps $H=160$. The success rate we report above is for 20 episodes with new initial states sampled from the training and new 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 with $\omega=1.8$ and during new concept learning, we learn two concepts and their corresponding classifier-free guidance weights. The Diffusion network architecture is a temporal U-Net where images are processed by a pretrained resnet18 [1] that is finetuned with the model during training. In our manipulation experiments we use a Realsense D435I RGB camera and an NVIDIA RTX 4090 machine.

# **Classifier-Free Guidance Parameter**

In Table 4 (see new pdf) we report results for various classifier-free guidance weights compared with the baselines in the Object Rearrangement, Goal-Oriented Navigation and Driving domains. For Object Rearrangement and Goal-Oriented Navigation we report the accuracy and standard error of the mean for new concepts from new initial states, and for Driving, the success and crash rates. We compare the four baselines with our approach as reported in Figures 5, 8 and 10 and section 5.2. We report results for our approach with all $\omega$ in our hyperparameter search, and mark the reported $\omega$ in Figures 5, 8 and 10 in bold font. Overall for various choices of $\omega$ we are better or comparable to the baselines. **We find that learning the weights along with concepts in all domains is a reasonable choice, and different domains may further benefit from tuning $\omega$ to a fixed value.**

[1] Deep residual learning for image recognition. He, et al. 2016.

Meta-review: 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.
Decision: Accept (poster)