# Transformers are Adaptable Task Planners

Vidhi Jain1,2 &Yixin Lin1 &Eric Undersander1 &Yonatan Bisk2 &Akshara Rai1
1 Meta AI, 2 Carnegie Mellon University

###### Abstract

Every home is different, and every person likes things done in their
particular way. Therefore, home robots of the future need to both reason about
the sequential nature of day-to-day tasks and generalize to user’s
preferences. To this end, we propose a Transformer Task Planner (TTP) that
learns high-level actions from demonstrations by leveraging object attribute-
based representations. TTP can be pre-trained on multiple preferences and
shows generalization to unseen preferences using a single demonstration as a
prompt in a simulated dishwasher loading task. Further, we demonstrate real-
world dish rearrangement using TTP with a Franka Panda robotic arm, prompted
using a single human demonstration. Code:
https://anonymous.4open.science/r/temporal_task_planner-Paper148/

> Keywords: Task Planning, Prompt, Preferences, Object-centric Representation

##  1 Introduction

Consider a robot tasked with loading a dishwasher. Such a robot has to account
for task constraints (e.g. only an open dishwasher rack can be loaded), and
dynamic environments (e.g. more dishes may arrive once the robot starts
loading the dishwasher). Dishwasher loading is a canonical example of personal
preferences, where everyone has a different approach which the robot should
adapt to. Classical task planning deals with task constraints through symbolic
task description, but such descriptions are difficult to design and modify for
new preferences in complex tasks. Building easily adaptable long-horizon task
plans, under constraints and uncertainty, is an open problem in robotics.

Machine learning (ML) enables learning complex tasks without extensive expert
intervention: robotic navigation [1, 2, 3], in-hand manipulation [4, 5, 6, 7,
8], and planning [9, 10, 11, 12, 13]. Within task planning, ML is used to
replace user-defined symbolic descriptions [14], deal with uncertainty [15],
and adapt to preferences [16]. Recent work [17] has shown Transformer networks
[18] can learn temporally-consistent representations, and exhibit
generalization to new scenarios [19, 20, 21]. Our central question is: Can a
Transformer network learn task structure, adapt to user preferences, and
achieve complex long-horizon tasks using no symbolic task representations?

We hypothesize that task structure and preference are implicitly encoded in
demonstrations. When loading a dishwasher, a user pulls out a rack before
loading it, inherently encoding a structural constraint. They may place mugs
on the top rack and plates on the bottom, encoding their preference. Learning
user preferences from long-horizon demonstrations requires policies with
temporal context. For example, a user might prefer to load the top rack before
the bottom tray. The policy needs to consider the sequence of actions
demonstrated, rather than individual actions. Transformers are well-suited to
this problem, as they have been shown to learn long-range relationships [22],
although not in temporal robotic tasks. We propose Transformer Task Planner
(TTP) - an adaptation of a classic transformer architecture that includes
temporal-, pose- and category-specific embeddings to learn object-oriented
relationships over space and time. TTP generalizes beyond what was seen in
demonstrations – to variable numbers of objects and dynamic environments. By
pre-training TTP on multiple preferences, we build temporal representations
that can be generalized to new preferences.

Figure 1: (Left-to-Right, Top-to-bottom) A Franka Emika Panda arm organizing 4
dishes into 2 drawers, following preference shown in a demonstration. The
robot opens a top drawer, places objects in the top drawer, closes it, and
does the same for the bottom drawer. The high-level policy that makes
decisions about when to open drawers, what objects to pick, etc is learned in
simulation and transferred zero-shot to the real-world.

The main contributions of our work are: (1) Introduce transformers as a
promising architecture for learning task plans from demonstrations, using
object-centric embeddings (2) Demonstrate that preference conditioned pre-
training generalizes at test time to new, unseen preferences with a single
user demonstration. Our experiments use a complex high-dimensional dishwasher
loading environment (Fig. 4) with several challenges: complex task structure,
dynamically appearing objects and human-specific preferences. TTP successfully
learns this task from seven preferences with 80 demonstrations each, and
generalizes to unseen scenes and preferences and outperforms competitive
baselines [23, 24]. Finally, we transfer TTP to a rearrangement problem in the
real-world, where a Franka arm places dishes in two drawers, using a single
human demonstration (Fig. 1).

##  2 Transformer Task Planner (TTP)

We introduce TTP, a Transformer-based policy architecture for learning
sequential manipulation tasks. We assume low-level ‘generalized’ pick-place
primitive actions that apply to both objects like plates, bowls, and also to
dishwasher door, racks, etc. TTP learns a high-level policy for pick-place in
accordance with the task structure and preference shown in demonstrations. The
following sections are described with dishwasher-loading as an example, but
our setup is applicable to most long-horizon manipulation tasks with
‘generalized’ pick-place.

#### State-Action Representations

We consider a high-level policy that interacts with the environment at
discrete time steps. At every timestep $t$, we receive observation
$\boldsymbol{o}_{t}$ from the environment which is passed through a perception
pipeline to produce a set of rigid-body instances $\\{x_{i}\\}^{n}_{i=1}$,
corresponding to the $n$ objects currently visible. We express an instance as:
$x_{i}=\\{p_{i},c_{i},t\\}$, where $p_{i}$ is the pose of the object, $c_{i}$
is its category, and $t$ is the timestep while recording $\boldsymbol{o}_{t}$
(Fig. 2(a)). For example, for a bowl at the start of an episode, $p$ is its
current location w.r.t a global frame, $c$ is its category (bowl), and $t=0$.
The pick state $S^{pick}_{t}$ is described as the set of instances visible in
$\boldsymbol{o}_{t}:$ $S^{pick}_{t}=\\{x_{0},x_{1},\cdots,x_{n}\\}$.
$S^{pick}_{t}$ is passed to a learned policy $\pi$ to predict a pick action
$\boldsymbol{a}_{t}$. We describe the action $\boldsymbol{a}_{t}$ in terms of
what to pick and where to place it. Specifically, the pick action chooses one
instance from $x_{i}$ observed in $\boldsymbol{o}_{t}$:
$\pi(S^{pick})\rightarrow x_{target},\text{ where
}x_{target}\in\\{x_{i}\\}^{n}_{i=1}.$

Once a pick object is chosen, a similar procedure determines where to place
it. We pre-compile a list of discrete placement poses corresponding to viable
place locations for each object-category. These poses densely cover the
dishwasher racks, and are created by randomly placing objects in the
dishwasher and measuring the final pose they land in. All possible placement
locations for the picked object category, whether free or occupied, are used
to create a set of place instances $\\{g_{j}\\}_{j=1}^{l}$. Similar to
$x_{i}$, $g_{j}=\\{p_{j},c_{j},t\\}$, consists of a pose, category and
timestep. A boolean value $\mathit{r}$ in the attribute set distinguishes the
two instance types. The place state is
$S^{place}=S^{pick}\cup\\{g_{j}\\}_{j=1}^{l}.$ The same policy $\pi$ then
chooses where to place the object (Fig.2). Note that the input for predicting
place includes both objects and place instances, since the objects determine
whether a place instance is free to place or not. $x_{target}$ and
$g_{target}$ together make action $\boldsymbol{a}$, sent to a low-level pick-
place policy. The policy $\pi$ is modeled using a Transformer [18].

(a) Scene to Instance embeddings

(b) Instances to prediction

Figure 2: (Left) Architecture overview of how a scene is converted to a set of
instances. Each instance is comprised of attributes, i.e. pose, category,
timestep, and whether it is an object or place instance. (Right) Instance
attributes (gray) are passed to the encoder, which returns instance embeddings
for pickable objects (red), placeble locations (green), and $<$ACT$>$
embeddings (blue). The transformer outputs a chosen pick (red) and place
(green) instance embedding.

#### Instance Encoder

Each object and goal instance $x_{i}$, $g_{j}$ is projected into a higher-
dimensional vector space (Fig. 2(a)). Such embeddings improve the performance
of learned policies [18, 25]. For the pose embedding $\Phi_{p}$, we use a
positional encoding scheme similar to NeRF [25] to encode the 3D positional
coordinates and 4D quaternion rotation of an instance. For category, we use
the dimensions of the 3D bounding box of the object to build a continuous
space of object types and process this through an MLP $\Phi_{c}$. For each
discrete 1D timestep, we model $\Phi_{t}$ as a learnable embedding in a lookup
table, similar to the positional encodings in BERT [26]. To indicate whether
an instance is an object or placement location, we add a 1D boolean value
vectorized using a learnable embedding function $\Phi_{r}$. The concatenated
$d$-dimensional embedding for an instance at timestep $t$ is represented as
$f_{e}(x_{i})=\Phi_{t}||\Phi_{c}||\Phi_{p}||\Phi_{r}=e_{i}$. The encoded state
at time $t$ can be represented in terms of instances as:
$S^{enc}_{t}=[e_{0},e_{1},\cdots e_{N}]$. We drop $()^{enc}$ for brevity.

#### Demonstrations

are state-action sequences
$\mathcal{C}=\\{(S_{0},a_{0}),(S_{1},a_{1}),\cdots,(S_{T-1},a_{T-1}),(S_{T})\\}$.
Here $S_{i}$ is the set of object and place instances and $a_{i}$ are the
pick-place actions chosen by expert at time $i$. At every time step, we record
the state of the objects, the pick instance chosen by the expert, place
instances for the corresponding category, and the place instance chosen by the
expert. Expert actions are assumed to belong to the set of visible instances
in $S_{i}$. However, different experts can exhibit different preferences over
$a_{i}$. For example, one expert might choose to load bowls first in the top
rack, while another expert might load bowls last in the bottom rack. In the
training dataset, we assume labels for which preference each demonstration
belongs to, based on the expert used for collecting the demonstration. Given
$K$ demonstrations per preference $m\in\mathcal{M}$, we have a dataset for
preference $m$:
$\mathcal{D}_{m}=\\{\mathcal{C}_{1},\cdots,\mathcal{C}_{K}\\}$. The complete
training dataset consists of demonstrations from all preferences:
$\mathcal{D}=\bigcup_{m=1}^{M}\mathcal{D}_{m}$. During training, we learn a
policy that can reproduce all demonstrations in our dataset. This is
challenging, since the actions taken by different experts are different for
the same input, and the policy needs to disambiguate the preference. At test
time, we generalize the policy to both unseen scenes and unseen preferences.

###  2.1 Prompt-Situation Transformer

We use Transformers [18], a deep neural network that operates on sequences of
data, for learning a high-level pick-place policy $\pi$. The input to the
encoder is a $d$-dimensional token111Terminology borrowed from natural
language processing where tokens are words; here, they are instances. per
instance $e_{i}\in[1,..N]$ in the state $S$. In addition to instances, we
introduce special $<$ACT$>$ tokens222Similar to $<CLS>$ tokens used for
sentence classification., a zero vector for all attributes, to demarcate the
end of one state and start of the next. These tokens help maintain the
temporal structure; all instances between two $<$ACT$>$ tokens in a sequence
represent one observed state. A trajectory $\tau$, without actions, is:
$\tau=\big{[}S_{t=0},<$ACT$>,\cdots,S_{t=T-1},<$ACT$>,S_{t=T}\big{]}.$ To
learn a common policy for multiple preferences, we propose a prompt-situation
architecture (Fig. 3). The prompt encoder receives one demonstration
trajectory as input, and outputs a learned representation of the preference.
These output prompt tokens are input to a situation decoder, which also
receives the current state as input. The decoder $\pi$ is trained to predict
the action chosen by the expert for the situation, given a prompt
demonstration.

Figure 3: Prompt-Situation Architecture. The left is the prompt encoder which
takes as input a prompt demonstration and outputs a learned preference
embedding. The right half is the situation decoder, which conditioned on
preference embedding from prompt encoder, acts on the current state.

The left half is prompt encoder $\psi$ and the right half is the situation
decoder or policy $\pi$ acting on given state. The prompt encoder
$\psi:f_{slot}\circ f_{te}\circ f_{e}$ consists of an instance encoder
$f_{e}$, transformer encoder $f_{te}$, and a slot-attention layer $f_{slot}$
[27]. $\psi$ takes the whole demonstration trajectory $\tau_{\text{prompt}}$
as input and returns a fixed and reduced preference embedding
$\gamma=\psi(\tau_{\text{prompt}})$ of sequence length $H$. Slot attention is
an information bottleneck, which learns semantically meaningful
representations of the prompt.

The situation decoder is a policy $\pi:f_{td}\circ f_{e}$ that receives as
input $N$ instance tokens from the current scene $S$, consisting of objects,
as well as, placement instances separated by $<$ACT$>$ tokens (Fig. 2(b)). The
policy architecture is a transformer decoder [18] with self-attention layers
over the $N$ input tokens. This is followed by a cross-attention layer with
preference embedding $\gamma$ ($H$ tokens) from the prompt encoder. We select
the output of the situation decoder at the $<$ACT$>$ token and calculate dot-
product similarity with the $N$ input tokens $e_{i}$. The token with the
maximum dot-product is chosen as the predicted instance:
$x_{pred}=\max_{e_{i},i\in\\{1,N\\}}\big{(}\hat{x}_{<\textsc{ACT}>}\cdot
e_{i}\big{)}$. The training target is extracted from the demonstration dataset
$\mathcal{D}$, and the policy trained with cross-entropy to maximize the
similarity of output latent with the expert’s chosen instance embedding.

###  2.2 Multi-Preference Task Learning

We adopt a prompt-situation architecture for multi-preference learning. This
design (1) enables multi-preference training by disambiguating preferences,
(2) learns task-level rules shared between preferences (e.g. dishwasher should
be open before placing objects), (3) can generalize to unseen preferences at
test time, without fine-tuning. Given a ‘prompt’ demonstration of preference
$m$, our policy semantically imitates it in a different ‘situation’ (i.e. a
different initialization of the scene). To this end, we learn a representation
of the prompt $\gamma^{m}$ which conditions $\pi$ to imitate the expert.

$\displaystyle\psi(\tau_{\text{prompt}}^{m})$
$\displaystyle\rightarrow\gamma^{m}$ (1)
$\displaystyle\pi(S_{\text{situation}}|\gamma^{m})$ $\displaystyle\rightarrow
a^{m}_{t}=\\{x_{\text{pred}}^{m},g_{\text{pred}}^{m}\\}$ (2)

We train neural networks $\psi$ and $\pi$ together to minimize the total
prediction loss over all preferences using the multi-preference training
dataset $\mathcal{D}$. The overall objective becomes:

$\min_{m\sim
M,\tau\sim\mathcal{D}_{m},(S,a)\sim\mathcal{D}_{m}}\mathcal{L}_{CE}(a,\pi(S,\psi(\tau)))$
(3)

For every preference $m$ in the dataset, we sample a demonstration from
$\mathcal{D}_{m}$ and use it as a prompt for all state-action pairs $(S,a)$ in
$\mathcal{D}_{m}$. This includes the state-action pairs from
$\tau_{\text{prompt}}$ and creates a combinatorially large training dataset.
At test time, we record one prompt demo from a seen or unseen preference and
use it to condition $\pi$ and $\psi$ : $a=\pi(S,\psi(\tau_{\text{prompt}}))$.
All policy weights are kept fixed during testing, and generalization to new
preferences is zero-shot using the learned preference representation $\gamma$.
Unlike [24], $\gamma$ captures not just the final state, but a temporal
representation of the whole demonstration. Building a temporal representation
is crucial to encode demonstration preferences like order of loading racks and
objects. Even though the final state is the same for two preferences that only
differ in which rack is loaded first, our approach is able to distinguish
between them using the temporal information in $\tau_{\text{prompt}}$. To the
best of our knowledge, our approach is the first to temporally encode
preferences inferred from a demonstration in learned task planners.

##  3 Experiments

Figure 4: Dishwasher Loading demonstration in AI Habitat Kitchen Arrange
Simulator. Objects dynamically appear on the counter-top (ii-iv), and need to
be placed in the dishwasher. If the dishwasher racks are full, they land in
sink (v)

We present the “Replica Synthetic Apartment 0 Kitchen”333“Replica Synthetic
Apartment 0 Kitchen” was created with the consent of and compensation to
artists, and will be shared under a Creative Commons license for non-
commercial use with attribution (CC-BY-NC). (see figure 4, appendix and
video), an artist-authored interactive recreation of the kitchen of the
“Apartment 0” space from the Replica dataset [28]. We use selected objects
from the ReplicaCAD [29] dataset, including seven types of dishes, and
generate dishwasher loading demonstrations using an expert-designed data
generation script (see Appendix B). Given 7 categories of dishes and two
choices in which rack to load first, the hypothesis space of possible
preferences is $2\times 7!$. Our dataset consists of 12 preferences (7 train,
5 held-out test) with 100 sessions per preference. In a session,
$n\in\\{3,...,10\\}$ instances are loaded in each rack. The training data
consists of sessions with 6 or 7 objects allowed per rack. The held-out test
set contains 5 unseen preferences and settings for $\\{3,4,5,8,9,10\\}$
objects per rack. Additionally, to simulate a dynamic environment, we randomly
initialize new objects mid-session on the kitchen counter. This simulates
situations where the policy does not have full information of every object to
be loaded at the start of the session, and has to learn to be reactive to new
information.

We train a 2-head 2-layer Transformer encoder-decoder with 256 input and 512
hidden dimensions, and 50 slots and 3 iterations for Slot Attention (more
details in Appendix C). We test both in- and out-of-distribution performance
in simulation. For in-distribution evaluation, 10 sessions are held-out for
testing for each training preference. For out-of-distribution evaluation, we
create sessions with unseen preferences and unseen number of objects. We
evaluate trained policies on ‘rollouts’ in the simulation, a more complex
setting than accuracy of prediction. Rollouts require repeated decisions in
the environment, without any resets. A mistake made early on in a rollout
session can be catastrophic, and result in poor performance, even if the
prediction accuracy is high. For example, if a policy mistakenly fails to open
a dishwasher rack, the rollout performance will be poor, despite good
prediction accuracy. To measure success, we rollout the policy from an initial
state and compare the performance of the policy with an expert demonstration
from the same initial state. Note that the policy does not have access to the
expert demonstration, and the demonstration is only used for evaluation.
Specifically, we measure (1) how well is the final state packed and (2) how
much did the policy deviate from the expert demonstration?

Packing efficiency: We compare the number of objects placed in the dishwasher
by the policy to that in the expert’s demonstration. Let $a_{i}$ be number of
objects in top rack and $b_{i}$ be objects on bottom rack placed by the expert
in the $i^{th}$ demonstration. If the policy adheres to the preference and
places $\hat{a}_{i}$ and $\hat{b}_{i}$ on the top and bottom respectively,
then the $\textit{packing efficiency
(PE)}=\sum_{i}\Big{(}\frac{\hat{a}_{i}}{max(\hat{a}_{i},a_{i})}+\frac{\hat{b}_{i}}{max(\hat{b}_{i},b_{i})}\Big{)}$.
Packing efficiency is between 0 to 1, and higher is better. Note that if the
policy follows the wrong preference, then PE is 0, even if the dishwasher is
full.

Inverse Edit distance: We also calculate is the inverse edit distance between
the sequence of actions taken by the expert versus the learned policy. We
compute the Levenshtein distance444 Levenshtein distance=
textdistance(learned_seq, expert_seq) / len(expert_seq)  (LD) between the
policy’s and expert’s sequence of pick and place instances. Inverse edit
distance is defined as $ED=1-LD$; higher is better. This measures the temporal
deviation from the expert, instead of just the final state. If the expert
preference was to load the top rack first and the learned policy loads the
bottom first, $PE$ would be perfect, but inverse edit distance would be low.

###  3.1 Evaluation on simulated dishwasher loading

#### Baselines:

We compare our approach against Graph Neural Network (GNN) based preference
learning from [23] and [24]. Neither of these works are directly suitable for
our task, so we combine them to create a stronger baseline. We use ground-
truth preference labels, represented as a categorical distribution and add
them to the input features of the GNN, similar to [24]. For unseen
preferences, this is privileged information that our approach does not have
access to. We use the output of the GNN to make sequential predictions, as in
[23]. Thus, by combining the two works and adding privileged information about
ground-truth preference category, we create a GNN baseline which can act
according to preference in a dishwasher loading scenario. We train this policy
using imitation learning (IL), and reinforcement learning (RL), following
[23]. GNN-IL is trained from the same set of demonstrations as TTP using
behavior cloning. For GNN-RL, we use Proximal Policy Optimization [30] from
[31]. GNN-RL learns from scratch by directly interacting with the dishwasher
environment and obtaining a dense reward. For details, see Appendix C.1.

(a) Packing efficiency

(b) Inverse Edit Distance

Figure 5:  Comparisons of TTP, GNN-IL/RL and a Random policy in simulation
across two metrics. TTP shows good performance at the task of dishwasher
loading on seen and unseen preferences, and outperforms the GNN and random
baselines. However TTP’s generalization to unseen # objects if worse, but
still close to GNN-IL, and better than GNN-RL and random.

GNN-IL does not reach the same performance as TTP for in-distribution tasks
(PE of $0.34$ for GNN-IL vs $0.62$ for TTP). Note that unseen preferences are
also in-distribution for GNN-IL since we provide ground-truth preference
labels to the GNN. Hence, there isn’t a drop in performance for unseen
preferences for GNN, unlike TTP. Despite having no privileged information, TTP
outperforms GNN-IL in unseen preferences, and performs comparably on unseen
#objects. Due to the significantly long time horizons per session (more than
30 steps), GNN-RL fails to learn a meaningful policy even after a large budget
of 32,000 environment interactions and a dense reward (PE $0.017\pm 0.002$
using GNN-RL). Lastly, we find that the random policy RP is not able to solve
the task at all due to the large state and action space. TTP is able to solve
dishwasher loading using unseen preferences well (PE $0.54$). In contrast,
classical task planners like [32] need to be adapted per new preference. This
experiment shows that Transformers make adaptable task planners, using our
proposed prompt-situation architecture. However, TTP’s performance on unseen
#objects deteriorates ($0.62$ for seen versus $0.34$ on unseen #objects), and
we look more closely at that next.

Generalization to unseen # objects: Fig. 7(a) examines PE on out-of-
distribution sessions with lesser i.e. 3-5 or more i.e. 8-10 objects per rack.
The training set consists of demonstrations with 6 or 7 objects per rack. The
policy performs well on 5-7 objects, but poorer as we go further away from the
training distribution. Poorer PE is expected for larger number of objects, as
the action space of the policy increases, and the policy is more likely to
pick the wrong object type for a given preference. Poorer performance on 3-4
objects is caused by the policy closing the dishwasher early, as it has never
seen this state during training. Training with richer datasets, and adding
more randomization in the form of object masking might improve out-of-
distribution performance of TTP.

###  3.2 Real-world dish-rearrangement experiments

Figure 6: Pipeline for Real Hardware Experiments

We zero-shot transfer our policy trained in simulation to robotic hardware, by
assuming low-level controllers. We use a Franka Panda equipped with a Robotiq
2F-85 gripper, controlled using the Polymetis control framework [33]. For
perception, we find the extrinsics of three Intel Realsense D435 RGBD cameras
[34] using ARTags [35]. The camera output, extrinsics, and intrinsics are
combined using Open3D [36] and fed into a segmentation pipeline [37] to
generate object categories. For low-level pick, we use a grasp candidate
generator [38] applied to the object point cloud, and use it to grasp the
target object. Place is approximated as a ‘drop’ action in a pre-defined
location. Our hardware setup mirrors our simulation, with different categories
of dishware (bowls, cups, plates) on a table, a “dishwasher” (cabinet with two
drawers). The objective is to select an object to pick and place it into a
drawer (rack) (see Fig. 1).

We use a policy trained in simulation and apply it to a scene with four
objects (2 bowls, 1 cup, 1 plate) through the hardware pipeline described
above. We start by collecting a prompt human demonstration (more details in
Appendix A)., The learned policy, conditioned a prompt demonstration, is
applied to two variations of the same scene, and the predicted actions
executed. The policy was successful once with 100% success rate, and once with
75%, shown in Figure 1, bottom. The failure case was caused due to a
perception error – a bowl was classified as a plate. This demonstrates that
such TTP can be trained in simulation and applied directly to hardware. The
policy is robust to minor hardware errors, such as if a bowl grasp fails, it
just repeats the grasping action (see video and Appendix A). However, it
relies on accurate perception of the state. In the future, we would like to
further evaluate our approach on more diverse real-world settings and measure
its sensitivity to the different hardware components, informing future choices
for learning robust policies.

###  3.3 Ablations

(a)  Performance decays for OOD number of objects, i.e 3-5 & 8-10

(b)  Performance improves with the # of unique sessions in training

(c)  Performance improves with the # of unique preferences in training

Figure 7: (a) Out-of-distribution generalization to #objects. (b-c) ablation
experiments.

We study the sensitivity of our approach to training hyperparameters. First,
we vary the number of training sessions per preference and study
generalization to unseen scenarios of the same preference. Figure 7(b) shows
that the performance of TTP improves as the number of demonstrations increase,
indicating that our model is not overfitting to the training set, and might
benefit from further training samples. Next, we vary the number of training
preferences and evaluate generalization performance to unseen preferences.
Figure 7(c) shows that the benefits of adding additional preferences beyond 4
are minor, and similar performance is observed when training from 4-7
preferences. This is an interesting result, since one would assume that more
preferences improve generalization to unseen preferences. But for the kinds of
preferences considered in our problem, 4-7 preference types are enough for
generalization.

Table 1: Attribute Ablations Pose | Cat | Time | PE
---|---|---|---
$\times$ | $\times$ | $\times$ | 0.0
$\times$ | $\times$ | ✓ | 0.0
$\times$ | ✓ | $\times$ | 0.027
$\times$ | ✓ | ✓ | 0.142
✓ | $\times$ | $\times$ | 0.411
✓ | $\times$ | ✓ | 0.419
✓ | ✓ | $\times$ | 0.517
✓ | ✓ | ✓ | 0.606

Finally, we analyze which instance attributes are the most important for
learning in a object-centric sequential decision making task. We mask
different instance attributes to remove sources of information from the
instance tokens. As seen in Table 1, all components of our instance tokens
play significant role ($0.606$ with all, versus the next highest of $0.517$).
The most important attribute is pose of the objects (without the pose, top PE
is $0.142$), followed by the category. The timestep is the least important,
but the best PE comes from combining all three. (more details in Appendix
D.1).

##  4 Prior Work

#### Object-centric representations for sequential manipulation

Several works build object-centric pick-place representations using off-the-
shelf perception methods [39, 40, 41, 42, 43, 44]. Once estimated, object
states are used by a task and motion planner for sequential manipulation [32,
45], but objects in the scene are known. [46] combine motor learning with
object-centric representation, but transfer of policies is challenging.
Transporter Nets [47] use a visual encoder-decoder for table-top manipulation
tasks; SORNet [48] extracts object-centric representations from RGB images and
demonstrates generalization in sequential manipulation task. Inspired from
these, we learn policies for dishwasher loading, choosing from visible objects
to make pick-place decisions.

#### Transformers for sequence modeling

Transformers in NLP [18] and vision [49] have focused on self-supervised
pretraining due to abundant unsupervised data available. Recent works have
repurposed transformers for other sequence modeling tasks [50, 51, 52, 53, 54,
55]. Prompt Decision Transformer [56] considers a single model to encode the
prompt and the successive sequence of state. We consider state as variable
number of instance PlaTe [51] proposes planning from videos, while [53, 54,
55] model a sequential decision-making task. We consider long-horizon tasks
with partially observable state features, and user-specific preferences.

#### Preferences and prompt training

In literature, there are several ways of encoding preferences. [24] propose
VAE to learn user preferences for spatial arrangement based on just the final
state, while our approach models temporal preference from demonstrations.
Preference-based RL learns rewards based on human preferences [57, 58, 59,
60], but do not generalize to unseen preferences. On complex long-horizon
tasks, modeling human preferences enables faster learning than RL, even with
carefully designed rewards [61]. We show generalization to unseen preferences
by using prompts. Large language and vision models have shown generalization
through prompting [62, 63]. Prompting can also be used to guide a model to
quickly switch between multiple task objectives [64, 65, 66]. Specifically,
language models learn representations that be easy transferred to new tasks in
a few-shot setting [67, 68, 19, 62]. Our approach similarly utilizes prompts
for preference generalization in sequential decision-making robotic tasks.

##  5 Conclusions and Limitations

We present Tranformer Task Planner (TTP): a high-level, sequential,
preference-based policy from a single demonstration using a prompt-situation
architecture. We introduced a simulated dishwasher loading dataset with
demonstrations that adhere to varying preferences. TTP can solve a complex,
long-horizon dishwasher-loading task in simulation and transfer to the real
world.

We have demonstrated the TTP’s strong performance in the dishwasher setting.
This environment is both complex by virtue of its strict sequential nature and
yet incomplete as we assume doors and drawers can be easily opened and
perception is perfect. In real settings, the policy needs to learn how to
recover from its low-level execution mistakes. More complex preferences may
depend on differences in visual or textural patterns on objects, for which the
instance encoder would require modifications to encode such attributes. An
important question to address is how more complex motion plans interact with
or hinder the learning objective, especially due to different human and robot
affordances. Finally, prompts are only presented via demonstration, while
language might be a more natural interface for users.

## References

  * Kadian et al. [2020] A. Kadian, J. Truong, A. Gokaslan, A. Clegg, E. Wijmans, S. Lee, M. Savva, S. Chernova, and D. Batra.  Sim2real predictivity: Does evaluation in simulation predict real-world performance?  _IEEE Robotics and Automation Letters (RA-L)_ , 2020.
  * Truong et al. [2020] J. Truong, D. Yarats, T. Li, F. Meier, S. Chernova, D. Batra, and A. Rai.  Learning navigation skills for legged robots with learned robot embeddings.  _International Conference on Intelligent Robots and Systems (IROS)_ , 2020.
  * Kumar et al. [2021] A. Kumar, Z. Fu, D. Pathak, and J. Malik.  Rma: Rapid motor adaptation for legged robots.  _Robotics: Science and Systems (RSS)_ , 2021.
  * Kalashnikov et al. [2018] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke, et al.  Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation (2018).  _arXiv preprint arXiv:1806.10293_ , 2018.
  * Nagabandi et al. [2020] A. Nagabandi, K. Konolige, S. Levine, and V. Kumar.  Deep dynamics models for learning dexterous manipulation.  In _Conference on Robot Learning_ , pages 1101–1112. PMLR, 2020.
  * Wirnshofer et al. [2020] F. Wirnshofer, P. S. Schmitt, G. von Wichert, and W. Burgard.  Controlling contact-rich manipulation under partial observability.  In _Robotics: Science and Systems_ , 2020.
  * Qin et al. [2020] Z. Qin, K. Fang, Y. Zhu, L. Fei-Fei, and S. Savarese.  Keto: Learning keypoint representations for tool manipulation.  In _2020 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 7278–7285. IEEE, 2020.
  * Simeonov et al. [2020] A. Simeonov, Y. Du, B. Kim, F. R. Hogan, J. Tenenbaum, P. Agrawal, and A. Rodriguez.  A long horizon planning framework for manipulating rigid pointcloud objects.  _arXiv preprint arXiv:2011.08177_ , 2020.
  * Yang et al. [2020] G. Yang, A. Zhang, A. Morcos, J. Pineau, P. Abbeel, and R. Calandra.  Plan2vec: Unsupervised representation learning by latent plans.  In _Learning for Dynamics and Control_ , pages 935–946. PMLR, 2020\.
  * Pertsch et al. [2020] K. Pertsch, Y. Lee, and J. J. Lim.  Accelerating reinforcement learning with learned skill priors.  _arXiv preprint arXiv:2010.11944_ , 2020.
  * Singh et al. [2020] A. Singh, H. Liu, G. Zhou, A. Yu, N. Rhinehart, and S. Levine.  Parrot: Data-driven behavioral priors for reinforcement learning.  _arXiv preprint arXiv:2011.10024_ , 2020.
  * Driess et al. [2020] D. Driess, J.-S. Ha, and M. Toussaint.  Deep visual reasoning: Learning to predict action sequences for task and motion planning from an initial scene image.  _arXiv preprint arXiv:2006.05398_ , 2020.
  * Andreas et al. [2017] J. Andreas, D. Klein, and S. Levine.  Modular multitask reinforcement learning with policy sketches.  In _International Conference on Machine Learning_ , pages 166–175. PMLR, 2017.
  * Silver et al. [2022] T. Silver, R. Chitnis, N. Kumar, W. McClinton, T. Lozano-Perez, L. P. Kaelbling, and J. Tenenbaum.  Inventing relational state and action abstractions for effective and efficient bilevel planning.  _arXiv preprint arXiv:2203.09634_ , 2022.
  * Gordon et al. [2019] D. Gordon, D. Fox, and A. Farhadi.  What should i do now? marrying reinforcement learning and symbolic planning.  _arXiv preprint arXiv:1901.01492_ , 2019.
  * Kaushik et al. [2020] R. Kaushik, T. Anne, and J.-B. Mouret.  Fast online adaptation in robotics through meta-learning embeddings of simulated priors.  In _2020 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 5269–5276. IEEE, 2020.
  * Kaplan et al. [2017] R. Kaplan, C. Sauer, and A. Sosa.  Beating atari with natural language guided reinforcement learning.  _arXiv preprint arXiv:1704.05539_ , 2017.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Brown et al. [2020] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Sanh et al. [2021] V. Sanh, A. Webson, C. Raffel, S. H. Bach, L. Sutawika, Z. Alyafeai, A. Chaffin, A. Stiegler, T. L. Scao, A. Raja, et al.  Multitask prompted training enables zero-shot task generalization.  _arXiv preprint arXiv:2110.08207_ , 2021.
  * Liu et al. [2021] X. Liu, K. Ji, Y. Fu, Z. Du, Z. Yang, and J. Tang.  P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks.  _arXiv preprint arXiv:2110.07602_ , 2021.
  * Chaplot et al. [2021] D. S. Chaplot, D. Pathak, and J. Malik.  Differentiable spatial planning using transformers.  In _International Conference on Machine Learning_ , pages 1484–1495. PMLR, 2021.
  * Lin et al. [2022] Y. Lin, A. S. Wang, E. Undersander, and A. Rai.  Efficient and interpretable robot manipulation with graph neural networks.  _IEEE Robotics and Automation Letters_ , 2022.
  * Kapelyukh and Johns [2022] I. Kapelyukh and E. Johns.  My house, my rules: Learning tidying preferences with graph neural networks.  In _Conference on Robot Learning_ , pages 740–749. PMLR, 2022.
  * Mildenhall et al. [2020] B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng.  Nerf: Representing scenes as neural radiance fields for view synthesis.  In _European conference on computer vision_ , pages 405–421. Springer, 2020.
  * Devlin et al. [2018] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Locatello et al. [2020] F. Locatello, D. Weissenborn, T. Unterthiner, A. Mahendran, G. Heigold, J. Uszkoreit, A. Dosovitskiy, and T. Kipf.  Object-centric learning with slot attention.  _Advances in Neural Information Processing Systems_ , 33:11525–11538, 2020.
  * Straub et al. [2019] J. Straub, T. Whelan, L. Ma, Y. Chen, E. Wijmans, S. Green, J. J. Engel, R. Mur-Artal, C. Ren, S. Verma, A. Clarkson, M. Yan, B. Budge, Y. Yan, X. Pan, J. Yon, Y. Zou, K. Leon, N. Carter, J. Briales, T. Gillingham, E. Mueggler, L. Pesqueira, M. Savva, D. Batra, H. M. Strasdat, R. D. Nardi, M. Goesele, S. Lovegrove, and R. Newcombe.  The Replica dataset: A digital replica of indoor spaces.  _arXiv preprint arXiv:1906.05797_ , 2019.
  * Szot et al. [2021] A. Szot, A. Clegg, E. Undersander, E. Wijmans, Y. Zhao, J. Turner, N. Maestre, M. Mukadam, D. Chaplot, O. Maksymets, A. Gokaslan, V. Vondrus, S. Dharur, F. Meier, W. Galuba, A. Chang, Z. Kira, V. Koltun, J. Malik, M. Savva, and D. Batra.  Habitat 2.0: Training home assistants to rearrange their habitat.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2021.
  * Schulman et al. [2017] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Raffin et al. [2021] A. Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann.  Stable-baselines3: Reliable reinforcement learning implementations.  _Journal of Machine Learning Research_ , 2021.
  * Garrett et al. [2020] C. R. Garrett, R. Chitnis, R. Holladay, B. Kim, T. Silver, L. P. Kaelbling, and T. Lozano-Pérez.  Integrated task and motion planning.  _arXiv preprint arXiv:2010.01083_ , 2020.
  * Lin et al. [2021] Y. Lin, A. S. Wang, G. Sutanto, A. Rai, and F. Meier.  Polymetis.  https://facebookresearch.github.io/fairo/polymetis/, 2021.
  * Keselman et al. [2017] L. Keselman, J. Iselin Woodfill, A. Grunnet-Jepsen, and A. Bhowmik.  Intel realsense stereoscopic depth cameras.  In _Proceedings of the IEEE conference on computer vision and pattern recognition workshops_ , pages 1–10, 2017.
  * Fiala [2005] M. Fiala.  Artag, a fiducial marker system using digital techniques.  In _2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05)_ , volume 2, pages 590–596. IEEE, 2005.
  * Zhou et al. [2018] Q.-Y. Zhou, J. Park, and V. Koltun.  Open3d: A modern library for 3d data processing.  _arXiv preprint arXiv:1801.09847_ , 2018.
  * Xiang et al. [2020] Y. Xiang, C. Xie, A. Mousavian, and D. Fox.  Learning rgb-d feature embeddings for unseen object instance segmentation.  _arXiv preprint arXiv:2007.15157_ , 2020.
  * Fang et al. [2020] H.-S. Fang, C. Wang, M. Gou, and C. Lu.  Graspnet-1billion: A large-scale benchmark for general object grasping.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 11444–11453, 2020.
  * Wang et al. [2019] H. Wang, S. Sridhar, J. Huang, J. Valentin, S. Song, and L. J. Guibas.  Normalized object coordinate space for category-level 6d object pose and size estimation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 2642–2651, 2019.
  * Deng et al. [2020] X. Deng, Y. Xiang, A. Mousavian, C. Eppner, T. Bretl, and D. Fox.  Self-supervised 6d object pose estimation for robot manipulation.  In _2020 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 3665–3671. IEEE, 2020.
  * Zeng et al. [2017] A. Zeng, K.-T. Yu, S. Song, D. Suo, E. Walker, A. Rodriguez, and J. Xiao.  Multi-view self-supervised deep learning for 6d pose estimation in the amazon picking challenge.  In _2017 IEEE international conference on robotics and automation (ICRA)_ , pages 1386–1383. IEEE, 2017.
  * Zhu et al. [2014] M. Zhu, K. G. Derpanis, Y. Yang, S. Brahmbhatt, M. Zhang, C. Phillips, M. Lecce, and K. Daniilidis.  Single image 3d object detection and pose estimation for grasping.  In _2014 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 3936–3943. IEEE, 2014.
  * Yoon et al. [2003] Y. Yoon, G. N. DeSouza, and A. C. Kak.  Real-time tracking and pose estimation for industrial objects using geometric features.  In _2003 IEEE International conference on robotics and automation (cat. no. 03CH37422)_ , volume 3, pages 3473–3478. IEEE, 2003.
  * Jain et al. [2021] V. Jain, P. Agarwal, S. Patil, and K. Sycara.  Learning embeddings that capture spatial semantics for indoor navigation.  _arXiv preprint arXiv:2108.00159_ , 2021.
  * Paxton et al. [2019] C. Paxton, N. Ratliff, C. Eppner, and D. Fox.  Representing robot task plans as robust logical-dynamical systems.  In _2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 5588–5595. IEEE, 2019.
  * Florence et al. [2019] P. Florence, L. Manuelli, and R. Tedrake.  Self-supervised correspondence in visuomot or policy learning.  _IEEE Robotics and Automation Letters_ , 5(2):492–499, 2019.
  * Zeng et al. [2020] A. Zeng, P. Florence, J. Tompson, S. Welker, J. Chien, M. Attarian, T. Armstrong, I. Krasin, D. Duong, V. Sindhwani, and J. Lee.  Transporter networks: Rearranging the visual world for robotic manipulation.  _Conference on Robot Learning (CoRL)_ , 2020.
  * Yuan et al. [2021] W. Yuan, C. Paxton, K. Desingh, and D. Fox.  SORNet: Spatial object-centric representations for sequential manipulation.  In _5th Annual Conference on Robot Learning_ , 2021.  URL https://openreview.net/forum?id=mOLu2rODIJF.
  * He et al. [2022] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16000–16009, 2022.
  * Liu et al. [2021] W. Liu, C. Paxton, T. Hermans, and D. Fox.  Structformer: Learning spatial structure for language-guided semantic rearrangement of novel objects.  _ArXiv_ , abs/2110.10189, 2021.
  * Sun et al. [2022] J. Sun, D.-A. Huang, B. Lu, Y.-H. Liu, B. Zhou, and A. Garg.  Plate: Visually-grounded planning with transformers in procedural tasks.  _IEEE Robotics and Automation Letters_ , 7(2):4924–4930, 2022.
  * Jain et al. [2020] V. Jain, R. Jena, H. Li, T. Gupta, D. Hughes, M. Lewis, and K. Sycara.  Predicting human strategies in simulated search and rescue task.  _arXiv preprint arXiv:2011.07656_ , 2020.
  * Chen et al. [2021] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  _Advances in neural information processing systems_ , 34, 2021.
  * Janner et al. [2021] M. Janner, Q. Li, and S. Levine.  Offline reinforcement learning as one big sequence modeling problem.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Putterman et al. [2022] A. L. Putterman, K. Lu, I. Mordatch, and P. Abbeel.  Pretraining for language conditioned imitation with transformers, 2022\.  URL https://openreview.net/forum?id=eCPCn25gat.
  * Xu et al. [2022] M. Xu, Y. Shen, S. Zhang, Y. Lu, D. Zhao, B. J. Tenenbaum, and C. Gan.  Prompting decision transformer for few-shot policy generalization.  In _Thirty-ninth International Conference on Machine Learning_ , 2022\.
  * Wang et al. [2022] X. Wang, K. Lee, K. Hakhamaneshi, P. Abbeel, and M. Laskin.  Skill preferences: Learning to extract and execute robotic skills from human feedback.  In _Conference on Robot Learning_ , pages 1259–1268. PMLR, 2022.
  * Lee et al. [2021] K. Lee, L. Smith, A. Dragan, and P. Abbeel.  B-pref: Benchmarking preference-based reinforcement learning.  _arXiv preprint arXiv:2111.03026_ , 2021.
  * Liang et al. [2022] X. Liang, K. Shu, K. Lee, and P. Abbeel.  Reward uncertainty for exploration in preference-based reinforcement learning.  _arXiv preprint arXiv:2205.12401_ , 2022.
  * Knox et al. [2022] W. B. Knox, S. Hatgis-Kessell, S. Booth, S. Niekum, P. Stone, and A. Allievi.  Models of human preference for learning reward functions.  _arXiv preprint arXiv:2206.02231_ , 2022.
  * Christiano et al. [2017] P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei.  Deep reinforcement learning from human preferences.  _Advances in neural information processing systems_ , 30, 2017.
  * Chen et al. [2022] Y. Chen, Y. Liu, L. Dong, S. Wang, C. Zhu, M. Zeng, and Y. Zhang.  Adaprompt: Adaptive model training for prompt-based nlp.  _arXiv preprint arXiv:2202.04824_ , 2022.
  * Ramesh et al. [2021] A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever.  Zero-shot text-to-image generation.  In _International Conference on Machine Learning_ , pages 8821–8831. PMLR, 2021.
  * Raffel et al. [2020] C. Raffel, N. M. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _ArXiv_ , abs/1910.10683, 2020.
  * Lewis et al. [2020] M. Lewis, Y. Liu, N. Goyal, M. Ghazvininejad, A. Mohamed, O. Levy, V. Stoyanov, and L. Zettlemoyer.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _ACL_ , 2020.
  * Song et al. [2019] K. Song, X. Tan, T. Qin, J. Lu, and T.-Y. Liu.  Mass: Masked sequence to sequence pre-training for language generation.  In _ICML_ , 2019.
  * Schick and Schütze [2020] T. Schick and H. Schütze.  It’s not just size that matters: Small language models are also few-shot learners.  _arXiv preprint arXiv:2009.07118_ , 2020.
  * Liu et al. [2021] P. Liu, W. Yuan, J. Fu, Z. Jiang, H. Hayashi, and G. Neubig.  Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing.  _arXiv preprint arXiv:2107.13586_ , 2021.
  * He et al. [2017] K. He, G. Gkioxari, P. Dollár, and R. Girshick.  Mask r-cnn.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2961–2969, 2017.
  * Anderson et al. [2018] P. Anderson, A. Chang, D. S. Chaplot, A. Dosovitskiy, S. Gupta, V. Koltun, J. Kosecka, J. Malik, R. Mottaghi, M. Savva, et al.  On evaluation of embodied navigation agents.  _arXiv preprint arXiv:1807.06757_ , 2018.

##  Appendix A Hardware Experiments

###  A.1 Real-world prompt demonstration

Here we describe how we collected and processed a visual, human demonstration
in the real-world to treat as a prompt for the trained TTP policy (Fig. 8).
Essentially, we collect demonstration pointcloud sequences and manually
segment them into different pick-place segments, followed by extracting object
states. At each high-level step, we measure the state using three RealSense
RGBD cameras[34], which are calibrated to the robot frame of reference using
ARTags [35]. The camera output, extrinsics, and intrinsics are combined using
Open3D [36] to generate a combined pointcloud. This pointcloud is segmented
and clustered to give objects’ pose and category using the algorithm from [37]
and DBScan. For each object point cloud cluster, we identify the object pose
based on the mean of the point cloud. For category information we use median
RGB value of the pointcloud, and map it to apriori known set of objects. In
the future this can be replaced by more advanced techniques like MaskRCNN
[69]. Placement poses are approximated as a fixed, known location, as the
place action on hardware is a fixed ‘drop’ position and orientation. The per
step state of the objects is used to create the input prompt tokens used to
condition the policy rollout in the real-world, as described in Section 3.2.

Figure 8: Human demonstration of real-world rearrangement of household dishes.

###  A.2 Hardware policy rollout

We zero-shot transfer our policy $\pi$ trained in simulation to robotic
hardware, by assuming low-level controllers. We use a Franka Panda equipped
with a Robotiq 2F-85 gripper, controlled using the Polymetis control framework
[33]. Our hardware setup mirrors our simulation, with different categories of
dishware (bowls, cups, plates) on a table, a “dishwasher” (cabinet with two
drawers). The objective is to select an object to pick and place it into a
drawer (rack) (see Fig. 8).

Once we collect the human prompt demonstration tokens, we can use them to
condition the learned policy $\pi$ from simulation. Converting the hardware
state to tokens input to $\pi$ follows the same pipeline as the ones used for
collecting human demonstrations. At each step, the scene is captured using 3
Realsense cameras, and the combined pointcound is segmented and clustered to
get object poses and categories. This information along with the timestep is
used to generate instance tokens as described in Section 2 for all objects
visible to the cameras. For visible already placed objects, the place pose is
approximated as a fixed location. The policy $\pi$, conditioned on the human
demo, reasons about the state of the environment, and chooses which object to
pick. Next, we use a grasp generator from [38] that operates on point clouds
to generate candidate grasp locations on the chosen object. We filter out
grasp locations that are kinematically not reachable by the robot, as well as
grasp locations located on points that intersect with other objects in the
scene. Next, we select the top 5 most confident grasps, as estimated by the
grasp generator, and choose the most top-down grasp. We design an pre-grasp
approach pose for the robot which is the same final orientation as the grasp,
located higher on the grasping plane. The robot moves to the approach pose
following a minimum-jerk trajectory, and then follows a straight line path
along the approach axes to grasp the object. Once grasped, the object is moved
to the pre-defined place pose and dropped in a drawer. The primitives for
opening and closing the drawers are manually designed on hardware.

The learned policy, conditioned on prompt demonstrations, is applied to two
variations of the same scene, and the predicted pick actions are executed.
Fig.9 shows the captured image from one of the three cameras, the merged point
cloud and the chosen object to pick and selected grasp for the same. The
policy was successful once with 100% success rate, and once with 75%, shown in
Fig.1. The failure case was caused due to a perception error – a bowl was
classified as a plate. This demonstrates that our approach (TTP) can be
trained in simulation and applied directly to hardware. The policy is robust
to minor hardware errors like a failed grasp; it just measures the new state
of the environment and chooses the next object to grasp. For example, if the
robot fails to grasp a bowl, and slightly shifts the bowl, the cameras measure
the new pose of the bowl, which is sent to the policy. However, TTP relies on
accurate perception of the state. If an object is incorrectly classified, the
policy might choose to pick the wrong object, deviating from the demonstration
preference. In the future, we would like to further evaluate our approach on
more diverse real-world settings and measure its sensitivity to the different
hardware components, informing future choices for learning robust policies.

Figure 9: Point cloud and grasps for different objects during policy rollout.

###  A.3 Transforming hardware to simulation data distribution

The policy trained in simulation applies zero-shot to real-world scenarios,
but it requires a coordinate transform. Fig. 10 shows the coordinate frame of
reference in simulation and real world setting. Since our instance embedding
uses the poses of objects, it is dependant on the coordinate frame that the
training data was collected in. Since hardware and simulation are
significantly different, this coordinate frame is not the same between sim and
real. We build a transformation that converts hardware measured poses to the
simulation frame of reference, which is then used to create the instance
tokens. This ensures that there is no sim-to-real gap in object positions,
reducing the challenges involved in applying such a simulation trained policy
to hardware. In this section we describe how we convert the real world
coordinates to simulation frame coordinates for running the trained TTP policy
on a Franka arm.

Figure 10: Coordinate Frame of reference in simulation (left) and real world
setting (right). Red is x-axis, green is y-axis and blue is z-axis.

We use the semantic work area in simulation and hardware to transform the
hardware position coordinates to simulation position coordinates. We measure
the extremes of the real workspace by manually moving the robot to record
positions and orientations that define the extents of the workspace for the
table. The extents of the drawers are measured by placing ARTag markers. We
build 3 real-to-sim transformations using the extents for counter, top rack
and bottom rack: Let $X\in\mathbb{R}^{3\times N}$ contain homogeneous $xz-$
coordinates of a work area, along its column, as follows:

$X=\begin{bmatrix}x^{(1)}&x^{(2)}&\cdots\\\ z^{(1)}&z^{(2)}&\cdots\\\
1&1&\cdots\\\
\end{bmatrix}=\begin{bmatrix}\boldsymbol{x}^{(1)}&\boldsymbol{x}^{(2)}&\cdots\end{bmatrix}$
(4)

As the required transformation from real to simulation involves scaling and
translation only, we have 4 unknowns, namely,
$\boldsymbol{a}=[\alpha_{x},\alpha_{y},x_{trans},z_{trans}]$. Here
$\alpha_{x},\alpha_{z}$ are scaling factors and $x_{trans},z_{trans}$ are
translation offset for $x$ and $z$ axis respectively. To solve
$X_{sim}=AX_{hw}$, we need to find the transformation matrix
$A=\hat{\boldsymbol{a}}=\begin{bmatrix}\alpha_{x}&0&x_{trans}\\\
0&\alpha_{z}&z_{trans}\\\ 0&0&1\\\ \end{bmatrix}$.

$\displaystyle X_{sim}$ $\displaystyle=\hat{\boldsymbol{a}}X_{hw}$ (5)
Rewriting the system of linear equations, (6)
$\displaystyle\implies\begin{bmatrix}x^{(1)}_{sim}\\\ z^{(1)}_{sim}\\\
x^{(2)}_{sim}\\\ z^{(2)}_{sim}\\\ \vdots\\\ \end{bmatrix}$
$\displaystyle=\begin{bmatrix}x^{(1)}_{hw}&0&1&0\\\ 0&z^{(1)}_{hw}&0&1\\\
x^{(2)}_{hw}&0&1&0\\\ 0&z^{(2)}_{hw}&0&1\\\ \vdots&\vdots&\vdots&\vdots\\\
\end{bmatrix}\boldsymbol{a}^{T}$ (7)

Let the above equation be expressed as $Y_{sim}=Z_{hw}a^{T}$ where
$Y_{sim}\in\mathbb{R}^{2N\times 1}$, $Z_{hw}\in\mathbb{R}^{2N\times 4}$, and
$a^{T}\in\mathbb{R}^{4\times 1}$. Assuming we have sufficient number of pairs
of corresponding points in simulation and real world, we can solve for
$\boldsymbol{a}$ by least squares
$a=(Z_{hw}^{T}Z_{hw})^{-1}Z_{hw}^{T}Y_{sim}$. The height $y_{sim}$ is chosen
from a look-up table based on $y_{hw}$. Once we compute the transformation
$A$, we store it for later to process arbitrary coordinates from real to sim,
as shown below.

⬇

def get_simulation_coordinates(xyz_hw: List[float], A: np.array) -> List:

xz_hw = [xyz_hw[0], xyz_hw[2]]

X_hw = get_homogenous_coordinates(xz_hw)

X_sim_homo = np.matmul(A, X_hw)

y_sim = process_height(xyz_hw[1])

X_sim = [X_sim_homo[0]/X_sim_homo[2], y_sim, X_sim_homo[1]/X_sim_homo[2]]

return X_sim

The objects used in simulation training are different from hardware objects,
even though they belong to the same categories. For example, while both sim
and real have a small plate, the sizes of these plates are different. We can
estimate the size of the objects based on actual bounding box from the
segmentation pipeline. However, it is significantly out-of-distribution from
the training data, due to object mismatch. So, we map each detected object to
the nearest matching object in simulation and use the simulation size as the
input to the policy. This is non-ideal, as the placing might differ for sim
versus real objects. In the future, we would like to train with rich
variations of object bounding box size in simulation so that the policy can
generalize to unseen object shapes in the real world.

##  Appendix B Simulation Setup

Figure 11: Human demonstration with point and click in simulation

#### Dataset

“Replica Synthetic Apartment 0 Kitchen” consists of a fully-interactive
dishwasher with a door and two sliding racks, an adjacent counter with a sink,
and a “stage” with walls, floors, and ceiling. We use selected objects from
the ReplicaCAD [29] dataset, including seven types of dishes (cups, glasses,
trays, small bowls, big bowls, small plates, big plates) which are loaded into
the dishwasher. Fig. 11 shows a human demonstration recorded in simulation by
pointing and clicking on desired object to pick and place. We initialize every
scene with an empty dishwasher and random objects placed on the counter. Next,
we generate dishwasher loading demonstrations, adhering to a given preference,
using an expert-designed data generation script. Expert actions include,
opening/closing dishwasher/racks, picking/placing objects in feasible
locations or the sink if there are no feasible locations left. Experts differ
in their preferences, and might choose different object arrangements in the
dishwasher.

Table 2: Three example preferences for dishwasher loading. Rack order and their respective contents (ordered by preference). First? | Top | Bottom
---|---|---
Top |  | 1\. cups
---
2\. glasses
3\. small bowl
| 1\. big plates
---
2\. small plates
3\. trays
4\. big bowls
Bottom |  | 1\. cups
---
2\. glasses
3\. small bowl
| 1\. big plates
---
2\. small plates
3\. trays
4\. big bowl
Bottom |  | 1\. small plate
---
2\. glasses
3\. cups
| 1\. big bowls
---
2\. trays
3\. big plates
4\. small bowl

#### Expert Preferences

We define a preference in terms of expert demonstration ‘properties’, like
which rack is loaded first with what objects? There are combinatorially many
preferences possible, depending on how many objects we use in the training
set. For example, Table 2 describes the preferences of dishwasher loading in
terms of three properties - first loaded tray, objects in top and bottom tray.
Each preference specifies properties such as which rack to load first and
their contents. In Table 2, Preferences 1 & 2 vary in the order of which rack
is loaded first, while 2 & 3 both load the bottom rack first with similar
categories on top and bottom but with different orderings for these
categories. Other preferences can have different combinations of objects
loaded per rack.

To describe a preference, let there be $k$ properties, where each can take
$m_{k}$ values respectively. For example, a property to describe preference
can be which rack is loaded first, and this can take two values; either top or
bottom rack. The total number of possible preferences is
$G=\prod_{i=1}^{k}m_{i}.$

In our demonstration dataset, we have 100 unique sessions per preference. Each
session can act as a prompt to indicate preference as well as provide
situation for the policy. Each session is about $\sim 30$ steps long. With $7$
preferences, this leads to $70,000\times 30=2,100,000$ $\sim 2$ million total
training samples, creating a relatively large training dataset from only 100
unique demonstrations per preference. Individual task preferences differ in
the sequence of expert actions, but collectively, preferences share the
underlying task semantics.

#### Dynamically appearing objects

To add additional complexity to our simulation environment, we simulate a
setting with dynamically appearing objects later in the episode. During each
session, the scene is initialized with $p\%$ of maximum objects allowed. The
policy/expert starts filling a dishwasher using these initialized objects.
After all the initial objects are loaded and both racks are closed, new
objects are initialized one-per-timestep to the policy. The goal is to
simulate an environment where the policy does not have perfect knowledge of
the scene, and needs to reactively reason about new information. The policy
reasons on both object configurations in the racks, and the new object type to
decide whether to ‘open a rack and place the utensil’ or ‘drop the object in
the sink’.

##  Appendix C Training

In this Section we describe details of the different components of our
learning pipeline.

###  C.1 Baseline: GNN

#### Architecture

We use GNN with attention. The input consists of 12 dimensional attribute
inputs (1D-timestep, 3D-category bounding box extents, 7D-pose, 1D-is object
or not?) and 12 dimensional one-hot encoding for the preference.

⬇

input_dim: 24

hidden_dim: 128

epochs: 200

batch_size: 32

#### Optimizer

: Adam with $lr=0.01$ and weight_decay$=1e-3$.

#### Reward function for GNN-RL

Reward function for the RL policy is defined in terms of preference. The
policy gets a reward of +1 every time it predicts the instance to pick that
has the category according to the preference order and whether it is placed on
the preferred rack.

###  C.2 Our proposed approach: TTP

#### Architecture

We use a 2-layer 2-head Transformer network for encoder and decoder. The input
dimension of instance embedding is 256 and the hidden layer dimension is 512.
The attributes contribute to the instance embedding as follows:

⬇

C_embed: 16

category_embed_size: 64

pose_embed_size: 128

temporal_embed_size: 32

marker_embed_size: 32

For the slot attention layer at the head of Transformer encoder, we use:

⬇

num_slots: 50

slot_iters: 3

#### Optimizer

We use a batch-size of 64 sequences. Within each batch, we use pad the inputs
with 0 upto the max sequence length. Our optimizer of choice is SGD with
momentum 0.9, weight decay 0.0001 and dampening 0.1. The initial learning rate
is 0.01, with exponential decay of 0.9995 per 10 gradient updates. We used
early stopping with patience 100.

(a) Category level accuracy grouped by batch size for prompt-situation
training.

(b) TE (SPL) metric for held-out test settings.

###  C.3 Metrics

In Section 3, we presented packing efficiency (PE) and edit distance (ED)
metrics collected on a policy rollout. We present additional metrics about
training progress and rollout here.

Category-token Accuracy indicates how well the policy can mimic the expert’s
action, given the current state. We monitor training progress by matching the
predicted instance to the target chosen in demonstration (Fig. 12(a)). We see
that TTP is able to predict the same category object to pick perfectly
(accuracy close to $1.0$). However, this is a simpler setting that sequential
decision making. During rollout, any error in a state could create a setting
that is out-of-distribution for the policy. Thus, category token accuracy sets
an upper bound for rollout performance, that is, while having high category
token accuracy is necessary, it is not sufficient for high packing efficiency
and inverse edit distance.

Temporal efficiency: Just like SPL [70] for navigation agents, we define the
efficiency of temporal tasks in policy rollout, in order to study how
efficient the agent was at achieving the task. For episode $i\in[1,..N]$, let
the agent take $p_{i}$ number of high-level interactions to execute the task,
and the demonstration consists of $l_{i}$ interactions for the initial state.
We scale the packing efficiency $PE_{i}$ of the policy by the ratio of steps
taken by expert versus policy. Temporal efficiency is defined between 0 to 1,
and higher is better. This value will be equal to or lower than the packing
efficiency. This especially penalizes policies that present a ‘looping’
behavior, such as repeatedly open/close dishwasher trays, over policies that
reach a low PE in shorter episodes (for example, by placing most objects in
the sink). Fig 12(b) shows the temporal efficiency or SPL over our 4 main
held-out test settings.

##  Appendix D Additional Ablation Experiments

In Section 3.3, we presented ablation experiments over number of
demonstrations per preference used for training, and the number of unique
preferences used. In this Section, we present additional ablation experiments
over the design of instance encodings in TTP. Additionally, we also present
results where we increase the temporal context of TTP and study its effect on
performance.

###  D.1 Design of Instance Encoding

(c) Temporal encoding

(d) Category encoding

(e) Pose encoding

Figure 13: [Left-to-Right] Comparing different design choices of attribute
encoders in terms of category token accuracy on held-out test prompt-situation
session pairs.

#### How much does temporal encoding design matter?

Fig. 13 shows that learning an embedding per timestep or expanding it as
fourier transformed vector of sufficient size achieves high success. On the
other hand, having no timestep input shows slightly lower performance.
Timestep helps in encoding the order of the prompt states. The notion of
timestep is also incorportated by autoregressive masking in both the encoder
and the decoder.

#### How much does category encoding design matter?

In our work, we represent category as the extents of an objects’ bounding box.
An alternative would be to denote the category as a discrete set of
categorical labels. Intuitively, bounding box extents captures shape
similarity between objects and their placement implicitly, which discrete
category labels do not. Fig. 13 shows that fourier transform of the bounding
box achieves better performance than discrete labels, which exceeds the
performance with no category input.

#### How much does pose encoding design matter?

We encode pose as a 7-dim vector that includes 3d position and 4d quaternion.
Fig. 13 shows that the fourier transform of the pose encoding performs better
than feeding the 7 dim through MLP. Fourier transform of the pose performs
better because such a vector encodes the fine and coarse nuances
appropriately, which otherwise either require careful scaling or can be lost
during SGD training.

###  D.2 Markov assumption on the current state in partial visibility
scenarios

Figure 14: Category level accuracy for single preference training with varying
context windows.

Dynamic settings, as used in our simulation, can be partially observable. For
example, when the rack is closed, the policy doesn’t know whether it is full
or not from just the current state. If a new object arrives, the policy needs
to decide between opening the rack if there is space, or dropping the object
in sink if the rack is full. In such partially observed settings, the current
state may or may not contain all the information needed to reason about the
next action. However, given information from states in previous timesteps, the
policy can decide what action to take (whether to open the rack or directly
place the object in the sink). To this end, we train a single preference pick
only policy for different context history. As shown in Fig. 15, context window
of size $k$ processes the current state as well as $k$ predecessor states,
that is, in total $k+1$ states. While larger context window size learns
faster, the asymptotic performance for all context windows converges in our
setting.

Let context history $k$ refer to the number of previous states included in the
input. Then the input is a sequence of previous $k$ states’ instances
(including the current state), as shown in Fig. 15. Fig 14 shows that TTP gets
$>90\%$ category level prediction accuracy in validation for all context
windows. While larger context windows result in faster learning at the start
of the training, the asymptotic performance of all contexts is the same. This
points to the dataset being largely visible, and a single context window
capturing the required information. In the future, we would like to experiment
with more complex settings like mobile robots, which might require a longer
context.

Figure 15: Processing with previous Context History $k$

##  Appendix E Limitations and Future scope

In Section 6, we briefly discussed the limitations and risks. Here we enlist
more details and highlight future directions.

#### Pick grasping depends on accurate segmentation and edge detection

Grasping policy depends on quality of segmentation and edge detection of the
selected object. Due to noise in calibration, shadows and reflections, there
are errors in detecting the correct edge to successfully grasp the object. For
example, it is hard to grasp a plate in real setting. Plate is very close to
the ground and the depth cameras cannot detect a clean edge for grasping.
Therefore, in our work, we place the plate on an elevated stand for easy
grasping. Grasping success also depends on the size and kind of gripper used.

#### Placement in real setting

For placement, the orientation of final pose is often different from initial
pose and may require re-grasping. The placement pose at final settlement is
different from the robot’s end-effector pose while releasing the object from
its grasp. Similar to picking, placement accuracy will largely depend on
approperiate size and shape of gripper used. Due to these reasons, placement
in real world is an open challenging problem and we hope to address this
future work.

#### Hardware pipeline issues due to calibration

The resulting point cloud is generated noisy due to two reasons. First,
incorrect depth estimation due to camera hardware, lighting conditions,
shadows and reflections. Second, any small movements among cameras that
affects calibration. If we have a noisy point cloud, it is more likely to have
errors in subsequent segmentation and edge detection for grasp policy. Having
sufficient coverage of the workspace with cameras is important to mitigate
issues due to occlusions and incomplete point clouds.

#### Incomplete information in prompt

The prompt session may not contain all the information to execute on the
situation. For example, in a prompt session there might be no large plates
seen, which in incomplete or ambiguous information for the policy. This can be
mitigated by ensuring complete information in prompt demo or having multiple
prompts in slightly different initialization.