# Unsupervised Retrieval and Rewriting of Short Story Events for More
Interesting Chitchat Dialog Responses

Bodhisattwa Prasad Majumder Taylor Berg-Kirkpatrick Julian McAuley Harsh
Jhamtani

# Unsupervised Enrichment of Persona-grounded Dialog
with Background Stories

Bodhisattwa Prasad Majumder Taylor Berg-Kirkpatrick Julian McAuley Harsh
Jhamtani

# Unsupervised Enrichment of Persona-grounded Dialog
with Background Stories

Bodhisattwa Prasad Majumder Taylor Berg-Kirkpatrick Julian McAuley Harsh
Jhamtani

###### Abstract

Humans often refer to personal narratives, life experiences, and events to
make a conversation more engaging and rich. While persona-grounded dialog
models are able to generate responses that follow a given persona, they often
miss out on stating detailed experiences or events related to a persona, often
leaving conversations shallow and dull. In this work, we equip dialog models
with ‘background stories’ related to a persona by leveraging fictional
narratives from existing story datasets (e.g. ROCStories). Since current
dialog datasets do not contain such narratives as responses, we perform an
unsupervised adaptation of a retrieved story for generating a dialog response
using a gradient-based rewriting technique. Our proposed method encourages the
generated response to be fluent (i.e., highly likely) with the dialog history,
minimally different from the retrieved story to preserve event ordering and
consistent with the original persona. We demonstrate that our method can
generate responses that are more diverse, and are rated more engaging and
human-like by human evaluators, compared to outputs from existing dialog
models.

##  1 Introduction

Humans often rely on specific incidents and experiences while conversing in
social contexts Dunbar et al. (1997). Responses from existing chitchat dialog
agents often lack such specific details. To mitigate this, some prior work has
looked into assigning personas to dialog agents Zhang et al. (2018); Majumder
et al. (2020). However, persona descriptions are often shallow and limited in
scope, and while they lead to improvements response specificity, they still
lack the level of detail with which humans share experiences.

Figure 1: We enrich agent personas with ‘background stories’ from an existing
corpus. We propose a gradient-based technique which encourages the generated
response to be fluent with the dialog history, minimally different from the
retrieved story, and consistent with the persona. The proposed approach leads
to more specific and interesting responses.

In this work, we propose methods to enrich dialog personas with relevant
background events using fictional narratives from existing story datasets such
as ROCStories Mostafazadeh et al. (2016). For example, for a persona attribute
‘I have two children and a dog,’ we are able to identify a relevant narrative
from a story corpus (Figure 1). However, such stories may not directly fit
fluently in the dialog context. Thus, retrieved stories should be adapted to
construct a response that is fluent and relevant to the context. Since
existing datasets (such as PersonaChat Zhang et al. (2018)) do not contain
responses with such background stories, such adaptation has to be done in an
unsupervised fashion with decoders trained to generate responses conditioned
only on a dialog history and persona.

To adapt a retrieved narrative incident as a relevant background story, we use
a decoding procedure which encourages the generated response to (1) be fluent
with the dialog history, (2) be consistent with the original persona, and (3)
be minimally different from the retrieved story. While fluency with dialog
context is encouraged directly by the likelihood as per the underlying
language model the remaining two constraints are incorporated via iterative
updates to the decoder output distributions at inference time. Our inference-
time decoding method is different from the only recent effort by Su et al.
(2020) that leverages non-dialog data (forum comments, book snippets) as
distant labels to train dialog systems with supervision. Our contributions can
be summarized as follows:

  * •

We propose a novel approach to enrich dialog agent personas with relevant
backstories, relying only on existing story datasets.

  * •

We propose to use an unsupervised back-propagation based decoding
procedure111Code can be found at
https://github.com/majumderb/pabst to adapt the relevant stories such that the
resulting response is fluent with the dialog history and consistent with the
dialog agent persona. Our method works with a model trained just with dialog
data i.e. without access to story corpus at training time.

  * •

Our experiments demonstrate that the proposed approach results in much more
engaging and specific dialog outputs in a persona-grounded dialog setup. This
fills a gap in existing dialog models which often lack the capability to
generate responses about specific events and experiences relevant to persona
attributes.

##  2 Unsupervised Persona Enrichment with Background Stories

Given dialog history $h$ and persona $C$ consisting of several (typically 3-5,
example shown in Figure 1) attributes, our goal is to construct a dialog
response $x$. Our underlying model is based on the discrete persona attribute
choice model from Majumder et al. (2020). To generate a dialog utterance $x$,
we first sample a persona attribute $c\sim p(c|h)$ conditioned on the dialog
history $h$. $x$ is then generated conditioned on the dialog history and the
chosen persona attribute. The underlying dialog model’s decoder is initialized
with a pretrained GPT-2 model, and is fine-tuned on the PersonaChat dataset
Zhang et al. (2018). However, in our current setup, we also have to identify
relevant background stories and use them to construct fluent responses at
decoding time. Therefore, we propose a different decoding procedure.

To generate a response, we first sample a persona attribute $c\sim p(c|h)$.
Next we retrieve stories corresponding to the persona attribute $c$ (Section
2.1). However, the underlying dialog model is trained to generate responses
conditioned only on the dialog history and persona. To incorporate the
retrieved story in the response, we perform gradient-based inference (Section
2.2), that only assumes a left-to-right language model trained on dialog
context and responses, and the story is handled at decoding time in an
unsupervised fashion. We refer to the proposed method as Pabst (Unsupervised
PersonA enrichment with Background STories).

###  2.1 Retrieving Relevant Stories

For a persona attribute $c$, we aim to identify relevant stories from a story
corpus. Toward this goal, we rank the stories using the F1 component of BERT-
score Zhang et al. (2020) based retrieval using the persona attribute $c$ as
the query and the highest scoring story is chosen. Note that many of the
stories are written in the third person. For use as background stories, we
must first transform them to first–person. Following prior work Brahman and
Chaturvedi (2020), we identify the protagonist of such stories as the most
frequently occurring character. Thereafter, we use co-reference resolution Lee
et al. (2017) to identify all words or phrases that refer to the protagonist.
Finally, all words or phrases so identified are replaced with suitable first
person pronouns (e.g. ‘his books’ to ‘my books’).

###  2.2 Gradient-based Inference

Our underlying dialog model is not trained to condition on a retrieved story,
and cannot be directly used to construct a desirable response using $s$. To
tackle this, we consider a decoding strategy which, in addition to fluency
with history $h$, encourages response $x$ to follow two soft constraints: (1)
be minimally different from story $s$, and (2) be consistent with persona $c$.

First, we generate an initial response based only on the dialog history. Then
we perform an iterative procedure which alternates between performing a
forward pass on the language model to encourage fluency, and a backward pass
which updates the response via back-propagation to respect the two soft
constraints. However, $x$ is discrete, and cannot be directly updated using
gradients from back-propagation. Instead, we maintain and update a soft
representation $o$ of $x$, where $o_{i}$ corresponds to the last hidden state
representation for the $i^{th}$ token position, i.e.,
$p(x_{i})\sim\operatorname{softmax}(Wo_{i}/\tau)$, where $\tau$ is the
temperature parameter, $W$ is the embedding matrix, and
$Wo_{i}\in\mathcal{R}^{V}$ ($V$ is the vocabulary size). Our approach is
inspired by recent works that use gradient-based decoding for text generation
with soft constraints Dathathri et al. (2020); Qin et al. (2020). Next we
describe the backward and forward passes of the iterative procedure.

#### Backward Pass with Soft Constraints

We define the following soft constraints on response $x$:
(1) Divergence from story: We want to encourage $x$ to be _minimally
different_ from the story $s$. Following prior work Qin et al. (2020), we
compute a cross entropy loss (denoted by $\operatorname{cross-entr}$
henceforth) with story $s=\\{s_{1},\ldots,s_{T}\\}$ tokens as labels and
$Wo_{1},\ldots,Wo_{T}$ as the logits.

(2) Consistency to persona: We want $x$ to be _consistent with persona
attribute_ $c$. Consider a classifier $q_{\phi}(o,c)$ which predicts the
probability of $x$ (or rather the soft representation $o$ of $x$) entailing
$c$. The classifier $q_{\phi}(o,c)$ is a bag-of-words classification head on
decoder hidden states $o$, fine-tuned on the Dialogue-NLI dataset Welleck et
al. (2019) to predict whether pairs of persona attributes and responses are
entailed or not. The objective to maximize can be written as:

where $\lambda_{c}$ and $\lambda_{d}$ are hyper-parameters. We update $o$
through back-propagation by computing the gradient
$\nabla_{o}\mathcal{L}(c,s;o)$, while keeping the model parameters constant.
Let the resulting $o$ after the gradient-based updates be denoted by $o^{b}$.

#### Forward Pass to Encourage Fluency

Next we perform a forward pass of the underlying dialog model, with the goal
of regularizing the hidden states towards the unmodified language model
values. On computing the forward pass at the $j^{th}$ token, we mix the final
hidden states $o^{f}_{j}$ from the forward pass with $o^{b}_{j}$ computed in
the backward pass, via weighted addition to get the resulting
$o_{j}=\gamma\times o^{f}_{j}+(1-\gamma)\times o^{b}_{j}$, where
$\gamma\in(0,1)$ is a hyperparameter. The resulting $o_{j}$ is used for
computing the logits at the next time step $j+1$.

We initialize the output response by performing greedy decoding from the
underlying dialog model, conditioned on the dialog history and persona
attribute. Then we iteratively update $o$ by alternate backward and forward
passes. We sample the final response $x\sim\operatorname{softmax}(Wo/\tau)$.
In practice, we found that 5 iterations are sufficient to generate good
quality outputs.

##  3 Experiments

Method | Training | Decoding | D-1 | D-2 | ENTR
---|---|---|---|---|---
W/o Story Data
Transfero | Persona | Nucleus | $0.05$ | $0.11$ | $1.21$
DiscChoice | Persona | Nucleus | $0.15$ | $0.25$ | $1.25$
DiscChoice | cs-kb | Nucleus | $0.87$ | $1.07$ | $2.04$
With Story Data
DiscChoice | Pseudo | Nucleus | $0.91$ | $2.45$ | $2.89$
DiscChoice | MultiTask | Nucleus | $0.99$ | $2.54$ | $2.71$
DiscChoice | Persona | Retrieval | $2.56$ | $9.67$ | $3.86$
Pabst (Ours) | Persona | Grad. Inf. | $1.56$ | $3.57$ | $3.21$

Table 1: Diversity metrics on the PersonaChat test set. D-1/2 is the % of distinct uni- and bi-grams. ENTR is the geometric mean of n-gram entropy. Grad. Inf. is the unsupervised gradient-based decoding as opposed to Nucleus sampling Holtzman et al. (2020).  Pabst vs. | Transfero | DiscChoice | Retrieval | Pseudo | MultiTask | w/o DNLI | Gold
---|---|---|---|---|---|---|---
Aspect | win | loss | win | loss | win | loss | win | loss | win | loss | win | loss | win | loss
Sensible | $\mathbf{83.8}$ | $5.6$ | $\mathbf{71.4}$ | $10.1$ | $\mathbf{72.7}$ | $15.7$ | $\mathbf{94.3}$ | $3.1$ | $\mathbf{92.8}$ | $1.5$ | $\mathbf{81.7}$ | $6.8$ | $29.1$ | $\mathbf{33.6}$
Engaging | $\mathbf{96.7}$ | $1.8$ | $\mathbf{82.6}$ | $6.7$ | $\mathbf{87.2}$ | $5.3$ | $\mathbf{94.7}$ | $2.1$ | $\mathbf{94.1}$ | $2.6$ | $\mathbf{70.9}$ | $14.8$ | $\mathbf{70.2}$ | $15.6$
Table 2: Pairwise comparison (% win/loss cases, ties not reported) between
responses from Pabst and from other baselines as well as gold. All differences
are significant with $p<0.05$. Cohen’s Kappa Cohen (1960) for sensibility and
engagement were 0.79 and 0.82 respectively. Complete details of the human
evaluations are in Appendix §B.

| Dialog History
---
Agent: How was your last trip?
User: It was great. I went to Utah and saw the Canyonlands.
Retrieved Story:
I went on a trip. I went to Florida. When we got there, there
was no air conditioning. I was irritated. It was too hot.
DiscChoice: That’s great. You must be adventurous.
Retrieval: I went on a trip. I went to Florida. When we
got there, there was no air conditioning. I was irritated.
It was too hot.
Pseudo: I have a house in Florida.
MultiTask: I like Florida. I went to the beaches.
Pabst ($\lambda_{d}=5$): I went on a trip. I went to Florida. I have
a house there. We do not have air condition. It was hot.
Pabst ($\lambda_{d}=1$): That’s great. I frequently go to Florida.
I have a house there but it gets hot in summer.

| Dialog History
---
Agent: I work as stunt double. What about you?
User: I’m an engineer. I just got back from work. How’s yours?
Retrieved Story:
I was a working actor. For my new role I was a stunt double. I
started working on new stunts. When I jumped off a building
yesterday I got hurt. I had to be rushed to the hospital.
DiscChoice: That’s great. I am looking for a job.
Retrieval: I was a working actor. For my new role I was a
stunt double. I started working on new stunts. When I jumped
off a building yesterday I got hurt. I had to be rushed to the hospital.
Pseudo: I got a new job. I am an actor.
MultiTask: I started a new role. I got hurt.
Pabst ($\lambda_{d}=5$): I was working as a stunt double. I jumped
off a building. I got hurt. I had to be rushed to the hospital.
Pabst ($\lambda_{d}=1$): Yesterday while I was working, I jumped off
a building and I got hurt. I had to be taken to the hospital.

Table 3: Generations from different models. More examples are in Appendix §C.

We evaluate methods in terms of their capability to generate diverse, fluent
and engaging responses. Hyperparameters are noted in Appendix §A.

#### Datasets

We experiment with the PersonaChat dialog dataset Zhang et al. (2018)
consisting of 131,438 utterances for training, 15,602 for validation, and
15,024 for testing. For stories, we use the training split of the ROCStories
dataset Mostafazadeh et al. (2016), that consists of 78,529 stories, each
typically of 4 to 5 sentences.

#### Baselines

We consider two broad groups of models as baselines: (1) _Without access to
story corpus_ : We use finetuned GPT2 (Transfero) on PersonaChat, and the
discrete persona attribute choice model (DiscChoice) from Majumder et al.
(2020). We also consider a version of DiscChoice which enriches personas with
inferences from a commonsense knowledge base (cs-kb). (2) _Baselines using
story corpus_ : To allow DiscChoice models to generate story-like responses,
we adapt an alternative training regime (Pseudo) from Su et al. (2020), where
we randomly replace some of the target dialog responses with retrieved stories
---treating them as pseudo labels. Finally, we also consider a MultiTask
training setup from Su et al. (2020), wherein the decoder is trained on
PersonaChat as well as with a language modeling objective on ROCStories. We
additionally consider a Retrieval baseline that uses the retrieved story
verbatim as the dialog response.

###  3.1 Automatic Evaluation

We hypothesize that that the proposed approach to leverage external non-dialog
data can increase the diversity of the generated responses. Following prior
work Li et al. (2016), we report the percentage of distinct uni-grams and bi-
grams (D-1 and D-2 respectively). Note that these values do not capture the
actual frequency distribution of different word types. Therefore, we also
report the geometric mean of entropy values of empirical frequency
distributions of n-grams of words (n $\in\\{1,2,3\\}$) Jhamtani et al. (2018),
denoted by ENTR.

We observe that methods that use story data show much higher diversity
compared to methods that do not (Table 1). Among methods using story data,
gradient-based decoding (Pabst) performs better than DiscChoice trained with
Pseudo or MultiTask. Note that just using Retrieval outputs as-is leads to
even more diverse outputs than Pabst. However, they are much less sensible
with the context, as shown in human evaluations.

###  3.2 Human Evaluation

Since we do not have ground truth story-like responses in the dialog dataset,
we perform human evaluation with 150 test examples to investigate if Pabst
generates responses that are 1) sensible with the dialog history and 2)
engaging. We hired two Anglophone (Lifetime HIT acceptance % $>$ 85)
annotators for every test sample. The order of the systems present in the
interface is randomized. A snapshot of the human evaluation interface is
provided in Appendix §C. All differences in values from human evaluations are
significant with $p<0.05$ from bootstrap tests on 1000 subsets of size 50.
Cohen’s Kappa Cohen (1960) to measure inter-annotator agreement for
sensibility and engagement were 0.79 and 0.82 respectively.

From the results (shown in Table 3), we note that in comparison to responses
from baselines, responses from Pabst are more engaging and more sensible with
respect to the dialog history. We further make following observations.
Firstly, using the gradient-based decoding approach with retrieved stories
(Pabst) works significantly better than using distant supervision with stories
data (Pseudo and MultiTask). Secondly, background stories provide sufficient
detail for an engaging conversation compared to DiscChoice which expands
persona attributes using commonsense knowledge Majumder et al. (2020).
Finally, we also observe that Pabst performs worse when we do not use the
consistency constraint (w/o DNLI).

#### Choice of $\lambda_{d}$

We also experiment with different values of the weight for the divergence term
($\lambda_{d}$) in $\mathcal{L}$: High ($\lambda_{d}=5$), Moderate
($\lambda_{d}=1$), and Low ($\lambda_{d}=0.05$). We consider 100 samples for
this experiment. We attribute a high $\lambda_{d}$ to responses strictly
copying the story. We find that Pabst (moderate $\lambda_{d}$) wins wins 81.2%
and 69.1% cases against Pabst (high $\lambda_{d}$) on ‘sensible’ and
‘engaging’ response criteria respectively. Similarly, Pabst (moderate
$\lambda_{d}$) wins 93.2% and 84.7% cases against Pabst (low $\lambda_{d}$) in
terms of sensibility and engagement respectively.

#### Qualitative Analysis

Table 3 shows responses generated by different baselines. We observe that
Pabst is able to follow the retrieved story (same as output from Retrieval)
while modifying the response to be conversation-like and sensible with dialog
history. Responses from other baselines remain verbose or incoherent.
Mirroring the human evaluation, we observe that choosing a higher
$\lambda_{d}$ makes the model to almost repeat the retrieved story but a lower
value smooths the output to make it more sensible with the ongoing dialog.

##  4 Related Work

A desired impact of the proposed approach is increase in diversity of the
generated responses. To tackle the issue of diversity in dialog model outputs,
prior work has focused on decoding strategies such as diversity-promoting
sampling Holtzman et al. (2020); training strategies such as discouraging
undesirable responses via unlikelihood training Li et al. (2020); model
changes such as using stochastic variables Serban et al. (2017); and using
external data such as forum data Su et al. (2020) or external knowledge bases
Majumder et al. (2020). In contrast to these, our proposed method generates
responses with background stories using a gradient-based decoding approach.

One of the steps in our proposed approach is to retrieve relevant stories from
an external corpus. Prior work has explored using retrieval of similar dialog
instances as an initial step in improving response diversity and other human-
like desiderata in dialog Roller et al. (2020); Weston et al. (2018). Distant
supervision by using retrieved text snippets as pseudo responses has been
explored in prior work Su et al. (2020); Roller et al. (2020). We use an
external data source to improve dialog responses, a theme shared with some
efforts in other tasks such as machine translation Khandelwal et al. . The use
of narrative text in dialog has been explored in prior work, mostly as a
‘script’ or template for conversation Xu et al. (2020); Zhu et al. (2020). We
adapted a BERT-based retrieval method Zhang et al. (2020) in our case to
retrieve relevant story given dialog context and use retrieved story in the
decoding phase.

Gradient-based for text generation with soft constraints has been explored in
prior work Dathathri et al. (2020); Qin et al. (2020). Song et al. (2020)
focused on generating response which are consistent to given persona.
Differently, we use a gradient-based decoding to generate a dialog response
while honoring constraints such as consistency to persona and similarity to
retrieved story.

##  5 Conclusion

We propose a method to enrich persona-grounded dialog with background stories
at inference time only using an existing corpus of non-conversational
narratives---opening up new ways to generate enriched and engaging responses.
One of the limitations of Pabst is the assumption of the need of a background
story at every turn. As future work, we can include a decision step to decide
if we need to incorporate a background story or not, given the dialog history.
We can further explore ways to use retrieved stories over multiple turns
instead of a single turn.

## Acknowledgements

We thank anonymous reviewers for providing valuable feedback. BPM is partly
supported by a Qualcomm Innovation Fellowship and NSF Award #1750063. Findings
and observations are of the authors only and do not necessarily reflect the
views of the funding agencies.

## Impact Statement

In this work, we discuss ways to make a dialog system to generate more
engaging responses. Since we use a finetuned version of a pretrained
generative model, we inherit the general risk of generating biased or toxic
language, which should be carefully filtered. Furthermore, the generations may
incorporate biases that are already present in the dialog dataset and story
dataset due to crowd-sourced data collection. Hence, we cautiously advise any
developer who wishes to use a different story dataset for the background
stories to be aware of the biases present in the dataset. Finally, we also
note that experiments in this paper are limited only to English language.

## References

  * Brahman and Chaturvedi (2020) Faeze Brahman and Snigdha Chaturvedi. 2020.  Modeling protagonist emotions for emotion-aware storytelling.  In _EMNLP_ , pages 5277--5294.
  * Cohen (1960) Jacob Cohen. 1960.  A coefficient of agreement for nominal scales.  _Educational and psychological measurement_ , 20(1):37--46.
  * Dathathri et al. (2020) Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. 2020.  Plug and play language models: A simple approach to controlled text generation.  In _ICLR_.
  * Dunbar et al. (1997) Robin IM Dunbar, Anna Marriott, and Neil DC Duncan. 1997.  Human conversational behavior.  _Human nature_ , 8(3):231--246.
  * Holtzman et al. (2020) Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020.  The curious case of neural text degeneration.  In _ICLR_.
  * Jhamtani et al. (2018) Harsh Jhamtani, Varun Gangal, Eduard Hovy, Graham Neubig, and Taylor Berg-Kirkpatrick. 2018.  Learning to generate move-by-move commentary for chess games from large-scale social forum data.  In _ACL 2018_.
  * (7) Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis.  Nearest neighbor machine translation.  _CoRR_.
  * Lee et al. (2017) Kenton Lee, Luheng He, Mike Lewis, and Luke Zettlemoyer. 2017.  End-to-end neural coreference resolution.  In _EMNLP_.
  * Li et al. (2016) Jiwei Li, Michel Galley, Chris Brockett, Jianfeng Gao, and Bill Dolan. 2016.  A diversity-promoting objective function for neural conversation models.  In _NAACL HLT_.
  * Li et al. (2020) Margaret Li, Stephen Roller, Ilia Kulikov, Sean Welleck, Y-Lan Boureau, Kyunghyun Cho, and Jason Weston. 2020.  Don’t say that! making inconsistent dialogue unlikely with unlikelihood training.  In _ACL_.
  * Majumder et al. (2020) Bodhisattwa Prasad Majumder, Harsh Jhamtani, Taylor Berg-Kirkpatrick, and Julian J. McAuley. 2020.  Like hiking? you probably enjoy nature: Persona-grounded dialog with commonsense expansions.  In _EMNLP_.
  * Mostafazadeh et al. (2016) Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vanderwende, Pushmeet Kohli, and James F. Allen. 2016.  A corpus and evaluation framework for deeper understanding of commonsense stories.  _CoRR_ , abs/1604.01696.
  * Qin et al. (2020) Lianhui Qin, Vered Shwartz, Peter West, Chandra Bhagavatula, Jena D. Hwang, Ronan Le Bras, Antoine Bosselut, and Yejin Choi. 2020.  Back to the future: Unsupervised backprop-based decoding for counterfactual and abductive commonsense reasoning.  In _EMNLP_.
  * Roller et al. (2020) Stephen Roller, Emily Dinan, Naman Goyal, Da Ju, Mary Williamson, Yinhan Liu, Jing Xu, Myle Ott, Kurt Shuster, Eric M Smith, et al. 2020.  Recipes for building an open-domain chatbot.  _arXiv preprint arXiv:2004.13637_.
  * Serban et al. (2017) Iulian Vlad Serban, Alessandro Sordoni, Ryan Lowe, Laurent Charlin, Joelle Pineau, Aaron C. Courville, and Yoshua Bengio. 2017.  A hierarchical latent variable encoder-decoder model for generating dialogues.  In _AAAI_.
  * Song et al. (2020) Haoyu Song, Wei-Nan Zhang, Jingwen Hu, and Ting Liu. 2020.  Generating persona consistent dialogues by exploiting natural language inference.  In _AAAI_.
  * Su et al. (2020) Hui Su, Xiaoyu Shen, Sanqiang Zhao, Xiao Zhou, Pengwei Hu, Randy Zhong, Cheng Niu, and Jie Zhou. 2020.  Diversifying dialogue generation with non-conversational text.  In _ACL_.
  * Welleck et al. (2019) Sean Welleck, Jason Weston, Arthur Szlam, and Kyunghyun Cho. 2019.  Dialogue natural language inference.  In _ACL_.
  * Weston et al. (2018) Jason Weston, Emily Dinan, and Alexander H. Miller. 2018.  Retrieve and refine: Improved sequence generation models for dialogue.  In _SCAI@EMNLP_.
  * Wolf et al. (2019) Thomas Wolf, Victor Sanh, Julien Chaumond, and Clement Delangue. 2019.  Transfertransfo: A transfer learning approach for neural network based conversational agents.  _CoRR_ , abs/1901.08149.
  * Xu et al. (2020) Jun Xu, Zeyang Lei, Haifeng Wang, Zheng-Yu Niu, Hua Wu, and Wanxiang Che. 2020\.  Enhancing dialog coherence with event graph grounded content planning.  In _IJCAI_.
  * Zhang et al. (2018) Saizheng Zhang, Emily Dinan, Jack Urbanek, Arthur Szlam, Douwe Kiela, and Jason Weston. 2018.  Personalizing dialogue agents: I have a dog, do you have pets too? In _ACL_.
  * Zhang et al. (2020) Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. 2020\.  Bertscore: Evaluating text generation with BERT.  In _ICLR_.
  * Zhu et al. (2020) Yutao Zhu, Ruihua Song, Zhicheng Dou, Jian-Yun Nie, and Jin Zhou. 2020.  Scriptwriter: Narrative-guided script generation.  In _ACL_.

##  Appendix A Implementation Details

We obtain the PersonaChat dataset from ParlAI
repository222http://parl.ai/downloads/personachat/personachat.tgz. ROCStories
dataset is obtained from the repository of original
release333https://www.cs.rochester.edu/nlp/rocstories/. We adapted codes from
original PPLM Dathathri et al. (2020) repository444https://github.com/uber-
research/PPLM and modified them for our own objective function.

#### Network architecture

For the generator network, we use GPT2 (Transformer with 12 layers, 768 hidden
size, 12 heads--- gpt2-small555https://github.com/huggingface/transfer-
learning-conv-ai) following the state-of-the-art model Wolf et al. (2019) from
Conv-AI2 competition. The decoder has total of 124 Million parameters. We used
the pre-trained decoder model obtained from Majumder et al. (2020).

#### Hyperparameters

Pabst does not require any training since we perform gradient-based decoding
at the inference time. For our best method, in objective function
$\mathcal{L}$, we use $\lambda_{d}$ as 1 and $\lambda_{c}$ as 1. We keep
generation length to be 100 to encourage longer generations. We train the
consistency classifier using code from PPLM
repository666https://github.com/uber-
research/PPLM/blob/master/run_pplm_discrim_train.py. The weight $\gamma$ for
mixing forward and backward passes was set to 0.45. For Pseudo, we substitute
a random 30% of the original target responses with retrieved stories.

##  Appendix B Human Evaluations Setup

Figure 2 shows a sample question for the pairwise comparison between response
generated by Pabst and a baseline for sensibility and engagingness. A similar
setup was used to measure performance between Pabst variants with different
$\lambda_{d}$ values ($0.5,1,5$).

Figure 2: Human evaluation setup for pairwise comparison between Pabst and
another baseline

##  Appendix C Qualitative Examples

Tables 4 and 5 show generations from Pabst and baselines for sampled text
examples.

| Dialog History
---
Agent: I volunteer in shelters.
User: It is a noble act.
DiscChoice: Thank you. I like helping people.
Retrieval: I was at my house bored. I couldn’t figure out
what to do. I looked on TV for answers. I saw a commercial
about homeless people. So I decided to volunteer at a
homeless shelter.
Pseudo: Last week, I went to the local animal shelter on
the pet adaptation day.
MultiTask: Thank you! What about you?
Pabst ($\lambda_{d}=5.0$): I was bored. I wanted to help people. I
looked on TV about homeless people. So I decided to
volunteer at a homeless shelter.
Pabst ($\lambda_{d}=1.0$): Thank you. I was inspired by a TV
commercial and I volunteered at the homeless shelter.
Table 4: Generations from different models | Dialog History
---
Agent: I enjoy American sports. What sports do you like?
User: I love playing volleyball. I was in the university team.
DiscChoice: That’s great. I like adventure sports.
Retrieval: I always thought golf was for old people.
One day, my dad invited me to play golf with him.
I reluctantly accepted his invitation. On the golf range, I
saw lots of people of my age. I realized golf was a sport for
everyone to enjoy regardless of age.
Pseudo: Last week, I went to play golf. My dad invited
me. I was reluctant.
MultiTask: That’s great! I play volleyball during
weekend.
Pabst ($\lambda_{d}=5.0$): I like playing golf. I always thought it
is for old people. Recently I went to a golf range. But I saw
lots of people of my age. I realized golf is a
sports for everyone.
Pabst ($\lambda_{d}=1.0$): That’s great. I like playing golf. I
always thought it is for old people. Recently I went
to a golf range. But I saw lots of people of my age. I
realized golf is a sports for everyone.
Table 5: Generations from different models


# Learning with Language-Guided
State Abstractions

Andi Peng
MIT &Ilia Sucholutsky
Princeton &Belinda Z. Li11footnotemark: 1
MIT &Theodore R. Sumers
Princeton Thomas L. Griffiths
Princeton &Jacob Andreas
MIT &Julie A. Shah
MIT Equal contribution.

###### Abstract

We describe a framework for using natural language to design state
abstractions for imitation learning. Generalizable policy learning in high-
dimensional observation spaces is facilitated by well-designed state
representations, which can surface important features of an environment and
hide irrelevant ones. These state representations are typically manually
specified, or derived from other labor-intensive labeling procedures. Our
method, LGA (language-guided abstraction), uses a combination of natural
language supervision and background knowledge from language models (LMs) to
automatically build state representations tailored to unseen tasks. In LGA, a
user first provides a (possibly incomplete) description of a target task in
natural language; next, a pre-trained LM translates this task description into
a state abstraction function that masks out irrelevant features; finally, an
imitation policy is trained using a small number of demonstrations and LGA-
generated abstract states. Experiments on simulated robotic tasks show that
LGA yields state abstractions similar to those designed by humans, but in a
fraction of the time, and that these abstractions improve generalization and
robustness in the presence of spurious correlations and ambiguous
specifications. We illustrate the utility of the learned abstractions on
mobile manipulation tasks with a Spot robot.

##  1 Introduction

In unstructured environments with many objects, distractors, and possible
goals, learning generalizable policies from scratch using a small number of
demonstrations is challenging (Bobu et al., 2024; Correia & Alexandre, 2023).
Consider the demonstration in Fig. 1A, which shows a Spot robot executing a
specific maneuver to perform a desired task. Which task is demonstrated
here—grabbing an object, grabbing an orange, or giving an orange to a user?
Here, the observations do not provide enough evidence to meaningfully
disentangle which features are relevant for downstream learning.

In humans, _abstraction_ is essential for generalizable learning (McCarthy et
al., 2021b). When learning (and planning), humans reason over simplified
representations of environment states that hide details and distinctions not
needed for action prediction (Ho et al., 2022). Useful abstractions are task-
dependent, and a growing body of evidence supports the conclusion that humans
flexibly and dynamically construct such representations to learn new tasks (Ho
et al., 2023; Huey et al., 2023). Importantly, this process of abstraction
does not begin with a blank slate—instead, experience, common-sense knowledge,
and direct instruction provide rich sources of prior knowledge about which
features matter for which tasks (Fan et al., 2020). In Fig. 1A, learning that
the demonstrated skill involves a _fruit_ combined with prior knowledge about
which objects are fruits makes clear that the object’s identity (_orange_) is
likely important. Meanwhile, the demonstration provides complementary
information (about the desired movement speed, goal placement, etc.) that is
hard to specify in language and not encapsulated by the utterance _bring me a
fruit_ alone. In other words, actions also contain valuable information that
is necessary to complement the abstraction.

What would it take to build autonomous agents that can leverage both
demonstrations and background knowledge to reason about tasks and
representations? State abstraction has been a major topic of research from the
very earliest days of research on sequential decision-making, with significant
research devoted to both unsupervised representation learning (Coates & Ng,
2012; Bobu et al., 2023; Higgins et al., 2017; Lee et al., 2021) and human-
aided design (Abel et al., 2018; Cakmak & Thomaz, 2012; Bobu et al., 2021;
Abel et al., 2016). However, there are currently few tools for autonomously
incorporating human prior knowledge for constructing state abstractions in
unseen, unstructured environments.

In this paper, we propose to use _natural language_ as a source of information
for constructing state abstractions. Our approach, Language-Guided Abstraction
(LGA) (Fig. 1B), begins by querying humans for high-level task descriptions,
then uses a pre-trained language model (LM) to translate these descriptions
into task-relevant state abstractions. Importantly, LGA requires only natural
language annotations for state features. Unlike most recent work applying LMs
to sequential decision-making tasks, it does not depend on pre-trained skills
(Ahn et al., 2022; Huang et al., 2022b), environment interaction (Du et al.,
2023), large multitask datasets (Karamcheti et al., 2023; Shridhar et al.,
2022), or even the ability to describe behavior in language (Kwon et al.,
2023). It complements traditional supervised learning methods like behavior
cloning (BC), without relying on additional assumptions about the data
labeling process. Experiments comparing LGA to BC (and stronger variants of
BC) show that LGA-generated abstractions improve sample efficiency and
distributional robustness in both single- and multi-task settings. They match
the performance of human-designed state abstractions while requiring a
fraction of the human effort.

In summary, we (1) introduce LGA, a method for using text descriptions and
language models to build state abstractions for skill learning; (2) show that
LGA produces state abstractions similar (and similarly effective) to those
manually designed by human annotators with significantly less time; (3) show
that LGA-constructed state abstractions enable imitation learning that is more
robust to the presence of observational covariate shift and ambiguous
linguistic utterances; and (4) demonstrate LGA’s utility in real world mobile
manipulation tasks with a Spot robot.

Figure 1: A: Example demonstration in our environment, showing Spot picking up
an orange and bringing it to the user. B: Our approach, Language Guided
Abstraction (LGA), creates a state abstraction with task-relevant features
identified by an LM. The policy is learned directly over this abstracted
state.

##  2 Related Work

Language-Aided Reward Design. Large language models, trained on large amounts
of text data, contain commonsense information about object properties,
functions, and their salience/relevance to various tasks. Several works (Goyal
et al., 2019; Kwon et al., 2023; Sumers et al., 2021; Carta et al., 2022)
leverage this to shape or learn reward models by training a policy to complete
intermediate or higher-level tasks and asking the LM to annotate resulting
behavior. Language feedback can then be incorporated as a way to guide the
policy during training (Mu et al., 2022; Du et al., 2023). However, just as
with any reinforcement learning method, these works assume continuous
environment access during the training process as the policy gradually
improves with the updated reward (Shinn et al., 2023). In contrast, we instead
leverage LMs to produce state abstractions for learning skills via imitation
learning, which does not require continuous environment interaction or access
to rewards.

Language-Aided Planning. In parallel, a great deal of work has leveraged LMs
to output plans directly, i.e. generate primitives or high-level action
sequences (Sharma et al., 2022; Ahn et al., 2022; Huang et al., 2022a; b).=
These approaches use priors embedded in LMs to produce better instruction
following models, or in other words, better compose pre-existing skills to
generate more complex behavior (Zeng et al., 2023; Li et al., 2023; Ahn et
al., 2022; Wang et al., 2023). Such methods can be thought of as performing
action abstraction and require access to a library of pre-trained (often text
annotated) skills. In contrast, we use LMs to perform state abstraction for
learning better skills from scratch, and can therefore also sit upstream of
any instruction following method.

State Abstractions in Learning. There is substantial evidence to suggest much
of the flexibility of human learning and planning can be attributed to
information filtering of task-relevant features (McCarthy et al., 2021a; Huey
et al., 2023; Ho et al., 2022; Chater & Vitányi, 2003). This suggests flexibly
creating task-conditioned abstractions is important for generalizable
downstream learning, particularly in low-data regimes (Mu et al., 2022; Bobu
et al., 2024). While state abstraction has been explored in sequential
decision-making (Thomas & Barto, 2011), existing methods often assume pre-
defined hierarchies over the state-action space of the given environment or
hand defined primitives (Abel et al., 2016; Diuk et al., 2008). In this work,
we explore a method to autonomously construct state abstractions from RGB
observations with only a text specification of the task.

Perhaps the closest comparison to our work is Misra et al. (2018), which
conditions on language and raw visual observations to create a binary goal
mask, specifying goal location within an observation. However, our approach is
much more flexible (our approach can learn any demonstrated behavior, not just
goal navigation) and more personalizable (our approach allow humans to
interact with the system, enabling them to refine the representation based on
individual preference).

##  3 Problem Statement

Preliminaries. We formulate our tasks as Markov Decision Processes (MDPs)
(Puterman, 2014) defined by tuples
$\langle\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R}\rangle$ where
$\mathcal{S}$ is the state space, $\mathcal{A}$ the action space,
$\mathcal{T}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$
the transition probability distribution, and
$\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ the reward
function. A policy is denoted as
$\pi_{\psi}:\mathcal{S}\rightarrow\mathcal{A}$.

In behavioral cloning (BC), we assume access to a set of expert demonstrations
$D_{\text{train}}=\\{\tau^{i}\\}_{i=0}^{n}=\\{(s_{0}^{i},a_{0}^{i},s_{1}^{i},a_{1}^{i},...s_{t}^{i},a_{t}^{i})\\}_{i=0}^{n}$
from which we “clone” a policy $\pi_{\psi}$ (Pomerleau, 1988) by minimizing:

$\displaystyle\mathcal{L}_{\mathrm{BC}}=\mathbb{E}_{(s_{t}^{i},a_{t}^{i})~{}\sim
D_{\mathrm{train}}}[\|\pi_{\psi}(s_{t}^{i})-a^{i}_{t}\|^{2}_{2}]~{}.$ (1)

In goal-conditioned behavioral cloning (GCBC), policies additionally condition
on goals $\ell$ (Co-Reyes et al., 2018). Motivated by the idea that natural
language is a flexible, intuitive interface for humans to communicate, we
specify the goal through language instruction $\ell\in\mathcal{L}$, resulting
in:

$\displaystyle\mathcal{L}_{\mathrm{GCBC}}=\mathbb{E}_{(s_{t}^{i},a_{t}^{i},\mathbin{{\color[rgb]{1,0,0}\ell^{i}}})~{}\sim
D_{\mathrm{train}}}[\|\pi_{\psi}(s_{t}^{i},\mathbin{{\color[rgb]{1,0,0}\ell^{i}}})-a^{i}_{t}\|^{2}_{2}]~{}.$
(2)

Above we have highlighted differences from BC in red.

A (GC)BC policy $\pi(s^{i}_{t},\ell^{i})$ must generalize to novel commands
$\ell^{i}$ and contextualize them against potentially novel states
$s^{i}_{t}$. Due to difficulties with sampling representative data and expense
of collecting a large number of expert demonstrations, systematic biases may
be present in both the training demonstrations $(s^{i}_{0},a^{i}_{0},\cdots)$
for a goal, or the commands $\ell^{i}$ describing the goal. This can result in
brittle policies which fail to generalize (Ross et al., 2011).

We study two sources of covariate shift: (1) out-of-distribution observations
$s^{i}$, or (2) out-of-distribution utterances $\ell^{i}$. Training a policy
$\pi$ to be robust to both types requires a wealth of task-specific training
data that can be expensive for humans to annotate and produce (Peng et al.,
2023). Below, we offer a more efficient way to do so by constructing a task-
specific state abstraction.

##  4 Language-Guided Abstraction (LGA)

Traditional (GC)BC forces the behavioral policy to learn a joint distribution
over language, observations, and actions—effectively requiring the robot to
develop both language and perceptual scene understanding simultaneously, and
to ground language in the current observation. Our approach instead _offloads_
contextual language understanding to a LM which identifies task-relevant
features in the perceptual state. Unlike other recent methods that offload
language understanding to LMs (Ahn et al., 2022), we do not rely on libraries
of pre-trained language-annotated skills, which may be insufficient when there
are key behaviors indescribable in language. Instead, we introduce a _state
abstraction function_ that takes the raw perceptual inputs and language-
specified goal, and outputs a set of task-relevant features. Intuitively, LGA
can be seen as a form of language-guided attention (Bahdanau et al., 2015): it
conditions the robot’s observations on language, removing the burden of
language understanding from the policy. We describe our general framework
below, which we will instantiate in Section 5.

###  4.1 State Abstraction Function

Formally, we define a _state abstraction function_ $\hat{f}$ that produces
task-relevant state representations:
$\hat{f}:\mathcal{S}\times\mathcal{L}\to\hat{\mathcal{S}}$, consisting of
three steps:

Textualization ($s\to\phi$). First, similar to other LM-assisted methods
(Huang et al., 2022b; Ahn et al., 2022) the raw perceptual input $s$ is
transformed into a text-based feature set $\phi$, representing a set of
features (described in natural language) that encapsulate the agent’s full
perceptual inputs.111In the general case, this could be implemented via
segmentation (Kirillov et al., 2023) and a captioner. This text representation
may include common visual attributes of the observation like objects in the
scene, which are extractable via segmentation models (Kirillov et al., 2023).
In Figure 1B, for example, textualization transforms observations to a feature
set of object names and pixel locations.

Feature abstraction ($\phi\to\hat{\phi}$). Given a feature set $\phi$, we
achieve abstraction by using an LM to select the subset of features from
$\phi$ relevant to the task $\ell$: $(\phi,\ell)\to\hat{\phi}$. In Figure 1B,
abstraction removes distractor objects from the feature set, in this case
preserving only the target object (orange).

Instantiation ($\hat{\phi}\to\hat{s}$). As a last step, we transform the
abstracted feature set back into an (abstracted) perceptual input with only
relevant features on display: $\hat{\phi}\to\hat{s}$. This step is required
for rendering the abstracted feature set in a form usable by a policy. In
Figure 1B, for example, this step transforms the abstracted feature set into
an observation showing only the relevant object.

###  4.2 Abstraction-Conditioned Policy Learning

After receiving an abstract state $\hat{s}$ from the state abstraction
function, we learn a policy mapping from the abstract state to actions,
yielding $\pi_{\hat{\psi}}:\hat{\mathcal{S}}\to\mathcal{A}$.222LGA can also be
used to _augment_ the original observation, in which case a policy is learned
over the concatenated original and abstracted state-spaces, e.g.,
$\pi_{\hat{\psi}}:\hat{\mathcal{S}}\times\mathcal{S}\to\mathcal{A}$. This is a
more conservative approach, allowing the policy to observe both the original
and abstracted states. $\pi_{\hat{\psi}}$ is trained to minimize the loss:

$\displaystyle\mathcal{L}_{\mathrm{LGA}}=\mathbb{E}_{(s_{t}^{i},a_{t}^{i},\ell^{i})~{}\sim
D_{\mathrm{train}}}[||{\color[rgb]{1,0,0}\pi_{\hat{\psi}}}(\mathbin{{\color[rgb]{1,0,0}\hat{f}(s_{t}^{i},\ell^{i})}})-a^{i}_{t}||^{2}_{2}],$
(3)

with the differences from GCBC (Eq. 2) highlighted in red. The LGA policy
$\pi_{\hat{\psi}}$ never sees the language input $\ell$; instead, it operates
over the language-conditioned state representation, $\hat{s}$. Consequently,
it can exhibit non-trivial generalization capabilities relative to Eqs. 1 and
2, which given a single goal seen at training and test, collapse to the same
non-generalizable policy.

LGA offers several appealing properties relative to traditional (GC)BC. First,
it mitigates spurious correlations during training because all goal
information is highlighted in semantic maps rather than raw pixels. Second,
the LM can help resolve ambiguous goals present at test by converting the
language utterance and test observation into an unambiguous abstract state —
effective even if LGA has only seen a single unambiguous task at training.
Therefore, using LMs confer valuable generalization capabilities at test-time,
as LMs can “intercede” to determine only the contextually-appropriate task-
relevant features for policy input.

Figure 2:  We evaluate on three task settings in VIMA. A: Pick-and-place. B:
Rotate. C: Sweep while avoid. Red bounding boxes depict task-relevant features
which must be accounted for in the abstraction.

##  5 Experimental Setup

###  5.1 Environment and Feature Space

We generate robotic control tasks from VIMA (Jiang et al., 2022), a vision-
based manipulation environment. We study three types of tasks: pick-and-place
of an object onto a goal location, rotation of an object to a pre-specified
degree, and sweep of an object while avoiding a potential obstacle. The
environment consists of a large feature space (29 objects, e.g. bowl, and 81
colors/textures, e.g. wooden) of various semantic meaning. Observations are
top-down RGB images and actions are continuous pick-and-place or sweep poses
each consisting of a 2D coordinate and a rotation expressed as a quaternion.
Success is determined by executing the correct action on the target object
within radius $\epsilon$ of the goal.

VIMA is an appealing domain for evaluation due to (1) its combinatorially
large feature space for assessing policy generalization and (2) the presence
of tasks that require behaviors that are difficult to specify in language
(e.g. velocity, rotation, etc.). The latter means we cannot deploy an
instruction-following method with a library of language-annotated, pre-trained
skills.

The task-relevant feature set $\hat{\phi}$ is specified via a JSON containing
the set of possible task-relevant objects and colors. For pick-and-
place/rotate, this feature set includes the target object; for sweep, this
feature set parameterizes both the target object and the object to be avoided.
Because the feature space can be combinatorially large, constructing
abstractions that preserve only task-relevant features of an ambiguously
specified objective, e.g. something that holds water, is challenging.

###  5.2 LGA Implementation

Textualization. We first extract a structured feature set $\phi$ by obtaining
the ground truth state segmentation mask and object descriptions from the
simulator.

Feature abstraction. We use two versions of LGA for defining relevant feature
sets $\hat{\phi}$. LGA uses GPT4 (OpenAI, 2023) to specify a relevant feature
set.333We use the gpt-4-0613 version of GPT4. We give GPT4 a description of
the VIMA environment (including a list of all possible object types and
colors), the goal utterance, and a target object type or feature to evaluate
(full prompt and additional details can be found in section A.2). GPT4
provides a binary response indicating whether the object type or feature
should be included in the abstraction. LGA-HILL extends this by placing a
human in the loop. The human is shown the proposed abstraction and can refine
it by adding or removing features (see Section 5.4 for details).

Instantiation. As described in Section 5.1, the abstract state representation
consists of a distribution over object types and object colors, including both
target objects and relevant obstacles. Given a specified task-relevant object
type and color, we use an image editor to identify all salient objects in the
observation. In particular, the image editor produces a “goal mask”, a pixel
mask where the relevant objects are converted to ones and everything else is
zeroed out. 444Here, we produce a simple binary pixel mask highlighting
relevant objects. In general, however, it may be fruitful to experiment with
more general scalar cost maps, such as highlighting areas that should be
avoided. In our setting, the simulator provides the ground-truth scene
segmentation which we use to produce the goal mask $\hat{s}$.

Abstraction-Augmented Policy Learning We instantiate the abstraction-augmented
policy procedure, described in Section 4.2, in our environment. We assume
access to ground truth demonstrations generated by an oracle. LGA learns a
policy from scratch via imitation learning on these demonstrations, using
$\hat{s}$. We implement a CNN architecture that processes abstract states into
embeddings, which we then feed through a MLP for action prediction. See
Section A.3 for details.

###  5.3 Training and Test Tasks

Training Tasks. We are interested in studying how well different methods
perform task-relevant feature selection, i.e. specify $\hat{\phi}$ given a
single language description of the task. We evaluate the quality of
abstractions by using them to construct a training distribution for each task
(on which we perform imitation learning). We want methods which construct a
minimal abstraction $\hat{\phi}$ that contains features necessary for the
task. If these abstractions do not contain all features necessary for
completion of the task, then the policy will not succeed; if these
abstractions are not minimal, then learning will be less data-efficient. Thus,
we evaluate both data efficiency and task performance; good abstractions will
result in the policy achieving higher success with fewer demonstrations.

We construct this distribution by placing a uniform probability over the set
of task-relevant features specified in $\hat{\phi}$. To create a task, we
sample an object and texture from this distribution as the “target” object and
place it randomly in the environment. For example, for something that can hold
water, we sample an object from {pan, bowl, container} and sample a color from
the universe of all colors. In the sweep task, we perform this same process,
but for possible obstacles. We then generate a random distractor object.
Distractors, target objects, and obstacles are randomly placed in one of four
discretized state locations. Last, for pick-and-place, we place a fixed goal
object (pallet) in the state as the “place” goal. We generate corresponding
training demonstrations via an oracle.

Test Tasks. In practice, we do not have access to ground-truth distributions –
this specification must come from the end user. However, for our experiments,
the first three authors defined “true” distributions from which we sampled
starting states for each task. See details in Section A.2.

###  5.4 Comparisons

We analyze our approach (LGA and LGA-HILL) against three baselines and two
ablations.

Human (Baseline). The first baseline explores having users (instead of LMs)
manually specify task-relevant features $\hat{\phi}$ in the feature
abstraction stage, to isolate the effect of LM-aided specification. All other
components remain the same as in LGA. The features are fed to the same CNN
architecture as for LGA and the policy is again learned over the generated
state abstractions.

GCBC-DART (Baseline): Our second baseline is (goal-conditioned) behavior
cloning, described in Eq. 2. Because imitation learning is notoriously brittle
to distribution shift, we implement a stronger variant (DART) (Laskey et al.,
2017) by injecting Gaussian noise (swept from $\mu$ = 0.0 to 0.2, $\beta=0.5$)
into the collected demonstrations, sampling 5 noisy demonstrations per true
demonstration. We implement goal-conditioning by concatenating an LM embedding
of the goal utterance with a CNN embedding of the state and then learning a
policy over the joint representation. We generate language embeddings using
Sentence-BERT (Reimers & Gurevych, 2019).

GCBC-SEG (Baseline): To study the impact of object segmentation, our third
baseline implements goal-conditioning on the segmentation mask in addition to
the observation. We independently process the two observations via CNNs, and
concatenate the embeddings as a joint representation.

LGA-S (Ablation): We explore the effect of training with both the raw
observation $s$ in addition to $\hat{s}$ to study the impact of discarding the
original observation. The two observations are independently processed via
CNNs, then concatenated as a joint representation.

LGA-L (Ablation): We implement a variant of abstraction where we condition
only on the language embedding of the relevant features $\hat{\phi}$ rather
than on the state abstraction $\hat{s}$.

Human Data Collection. For comparisons that require human data (LGA-HILL and
Human), we conducted a computer-based in-person user study to assess the
ability of humans to specify task-relevant features $\hat{\phi}$ in the
feature abstraction step, both with and without an LM. We recruited 18
participants (67% male, aged 22-44) from the greater MIT community. We paid
each participant $30 for their time. Our study received institutional IRB
approval and all user data was anonymized.

We first introduce the user to the environment and full feature space $\phi$.
To additionally help participants, we provide both a text and visual
representation of features (the latter of which the LM does not see). We also
walk the user through an example task specification. In the Human condition,
we introduce task objectives (detailed in Section 5) sequentially to the user
and ask for them to specify $\hat{\phi}$ by typing their selected features
into a task design file (we provide the full feature list for easy access). In
the LGA-HILL condition, we show the same objectives, but with LM answers
prefilled as the “raw” $\hat{\phi}$. To minimize ordering bias, we randomly
assign half of our participants to begin with each condition. For both
conditions, we measure the time each user spent specifying the feature set.

##  6 Results

###  6.1 Q1: Improving ease and policy performance of task specification

We begin by evaluating both LGA and LGA-HILL against the baseline methods on
nine scenarios from pick-and-place, two from rotate, and two from sweep. Each
scenario demonstrates varied feature complexity, where the ground truth task
distributions may be combinatorially large. Each method learns a single-task
policy from scratch for each of the tasks.

Metrics. We report two metrics: task performance vs. # of demonstrations and
user time. Task performance is assessed as the success rate of the policy on
20 sampled test states. We plot this with respect to the number of training
demonstrations an agent sees. User specification time is measured as the time
(in seconds) each user spent specifying the task-relevant feature set for each
task.

Figure 3: (Q1) A: Comparing task performance (averaged over all tasks) of each
method when controlling the number of training demonstrations. B: Comparing
the amount of time (averaged over all tasks) that human users spent specifying
task-relevant features for each method. LGA outperforms baselines on task
performance while significantly reducing user time spent compared to manual
feature specification ($p<0.001$).

Task Objectives. For pick-and-place, we construct nine scenarios with a wide
range of target objects: red heart, heart, letter, tiger-colored object,
letter from the word letter, consonant with a warm color, vowel with multiple
colors, something to drink water out of, and something from a typical kitchen.
We chose these tasks to 1) test the ability of the LM in LGA to semantically
reason over which target objects and features satisfy human properties, e.g.,
to drink water out of, and 2) explore how potentially laborious performing
manual feature specification from large feature sets can be for human users.
For rotate, we construct two scenarios: rotate block, rotate something I can
put food in, to illustrate LGA’s ability to learn behaviors that are difficult
to specify in language (e.g., rotate an object 125 degrees), but easy to
specify for task-relevance (e.g., pan). For sweep, we construct two scenarios:
sweep the block without touching the pan and sweep the block without touching
the line to illustrate LGA’s ability to identify task-relevant features that
are both target objects and obstacles.

Results. We vary the number of demonstrations that each policy is trained on
from 10 to 50. We visualize the resulting policy performance in Figure 3 (A).
In Figure 3 (B), we visualize how much time users spent specifying the task-
relevant features for each method.555While we assign zero time spent
specifying features for GCBC in our experiments, in practice, user time would
be instead spent specifying initial state configurations for generating every
demonstration, which only disadvantages naive GCBC even more when compared to
methods that specify features. From these results, it is clear that the
feature-specification methods (i.e. LGA, LGA-HILL, and Human) are more sample-
efficient than baselines with consistently higher performance at each number
of training demonstrations. Furthermore, LGA methods require significantly
less user time than the Human baseline ($p<0.001$ using a paired t-test for
all tasks). We note that despite comparable task performance between LGA and
LGA-HILL, the latter remains a valuable instantiation for tasks that require
human input, such as high-stakes decision-making scenarios or personalization
to individual user preferences.

Summary. LGA requires significantly less human time than manually specifying
features, but still leads to more sample-efficient performance compared to
baselines.

###  6.2 Q2: Improving policy robustness to observational covariate shift

We have shown that LGA enables faster feature specification when compared to
hand-engineering features from scratch, and also leads to more sample-
efficient policies when compared to traditional imitation learning. We now ask
how the different design choices taken in LGA handle state covariate shift
compared to GCBC+DART, which is a stronger imitation learning baseline
designed to help mitigate covariate shift by synthetically broadening the
demonstrations data with injected noise.

Task Objectives. We evaluate on four previously defined scenarios in Q1: heart
(pick-and-place), letter (pick-and-place), bowl (rotate), and block (sweep
while avoid). For each task, we now construct state covariate shifts designed
to test policy robustness.

For all tasks, we define a training distribution where the task-relevant
features contain some subset of the feature list, e.g., red and blue. At test
time, we evaluate by sampling tasks from a distribution where either 1) task-
relevant textures change (e.g., objects are now pink) or 2) a sampled
distractor is added to the scene. These tasks are intended to evaluate LGA’s
flexibility in including (and excluding) appropriate task-relevant features
relative to imitation learning over raw observations.

Figure 4: (Q2) A: Results on state covariate shifts. LGA variants that
condition on state abstractions (with or without the original observation)
outperform LGA-L (which conditions on the language abstraction only) and
GCBC+DART (which attempts to use noise injection to handle covariate shift).
(Q3) B: Results on multi-task ambiguity. We observe the same trends, with LGA
variants that condition on state abstractions able to resolve previously
unseen linguistic commands.

Results. For each method, we train on 50 demonstrations from the train
distribution then evaluate on 20 tasks from the test distribution. As shown in
Fig. 4A, policies trained with either of the LGA methods are more robust to
state covariate shifts than those trained with GCBC+DART. Notably, we observe
that LGA-L (our ablation where we provide the language abstraction only)
underperforms the variants that operate over state abstractions. This makes
sense considering observations (even if masked) contain information that
directly grounds task-relevant information for the policy to operate over into
the environment. This confirms our hypothesis that providing a non-ambiguous
state abstraction that already precludes non-relevant task features from the
state result in policies that are less brittle to spuriously correlated state
features.

Summary. Policies trained with LGA state abstractions (with or without raw
observations) are more robust to observational shift than GCBC+DART.

###  6.3 Q3: Improving multi-task policy robustness to linguistic ambiguity

The previous experiments have focused on evaluating LGA in single-task
settings where each policy corresponds to a singular language specification.
We now evaluate how LGA performs in the multi-task case where test time
environments allow for multiple task specifications, and conditioning on
language is necessary to resolve the ambiguity.

Task Objectives. We now define three new tasks: pick up fruit (pick-and-
place), put food in (rotate), and avoid utensil (sweep). For each task, we
train a multi-task policy. We define our training distribution as states that
contain a single object from each specification (e.g., a tomato or apple). At
test time, we present the policy with states that now contain both objects. We
condition the policy on both the linguistic utterances it has seen before:
e.g., “Bring me a tomato.” or “Bring me an apple”, and as well as one it has
not: e.g., “Bring me a fruit.” We now evaluate LGA’s zero-shot generalization
performance on both seen and unseen linguistic goal utterances.

Results. For each method, we train on 50 demonstrations sampled from the train
distribution for each task, then evaluate on 20 task instances from the test
distribution. We visualize the averaged results in Figure 4B. These results
demonstrate LGA’s ability to flexibly construct state abstractions even if the
language utterance was previously unseen.

Summary. Multi-task policies trained with LGA can better resolve task
ambiguity compared to GCBC+DART as they can flexibly adapt new language
specifications to the observation.

##  7 Real-World Robotics Tasks

We highlight LGA’s real world applicability on a Spot robot performing mobile
manipulation tasks.

Robot Platform. Spot 666Our Spot’s name is Moana. is a mobile manipulation
legged robot equipped with six RGB-D cameras (one in gripper, two in front,
one on each side, one in back), each producing an observation of size
480$\times$640\. For our demonstrations, we only use observations taken from
the robot’s front camera.

Tasks and Data Collection. We collected demonstrations of a human operator
teleoperating the robot while performing two mobile manipulation tasks with
various household objects: bring fruit and throw away object. The manipulation
action space consists of the following three actions along with their
parameters: (xy, grasp), (xy, move), (drop) while the navigation action space
consists of a SE(3) group denoting navigation waypoints. For bring fruit, the
robot is tasked with picking up fruit on a table, bringing it to a user at a
specified location, and dropping it in their hand. For throw away drink, the
robot is tasked with picking up a drink from a table, moving to a specified
location, and dropping it into a recycling bin. Both tasks include possible
distractors like brushes and drills on the table. For each task, we collected
a single demonstration of the robot successfully performing the task (e.g.,
throwing away a soda can). The full action sequence is recorded for training.

Training and Test Procedure. First, we extract a segmented image using Segment
Anything (Kirillov et al., 2023) and captioner Dedic (Zhou et al., 2022);
second, we query the feature list along with the linguistic specification
(e.g., throw away drink) to construct our abstract feature list; last, we map
the abstract feature list back into the original observation as an abstracted
state. We train a policy to map from the abstracted state into the action
sequence.777For latency purposes and ease of data generation, we perform
imitation learning over the full trajectory rather than each state (i.e.
predict a sequence of actions from an initial observation only). For test, we
then change the scene (e.g., with a new target object like a water bottle
instead of a soda can, along with distractors) and predict the action sequence
from the new observation.

Takeaway. LGA produced policies capable of successfully completing both tasks
consistently. The failures we did observe were largely due to captioning
errors (e.g., the segmentation model detected the object but was unable to
produce a good text description).

##  8 Discussion & Conclusion

We proposed LGA to leverage LMs to learn generalizable policies on state
abstractions produced for task-relevant features. Our results confirmed that
LGA produces state abstractions similarly effective to those manually designed
by human annotators while requiring significantly less user time than manual
specification. Furthermore, we showed that these abstractions yield policies
robust to observational covariate shift and ambiguous language in both single-
task and multi-task settings.

Limitations. We assume state abstractions for training good policies can be
captured by visual features present in individual scenes. An exciting
direction for future work would be to learn how to build state abstractions
for trajectory-relevant features. We also assume features are a) segmentable
and b) expressible in language, i.e. a LM (or human) can understand how each
feature may or may not relate to the task. In more complex scenarios, many
real-world tasks may be dependent on features that are less expressible in
text, e.g. relative position, and are difficult to encapsulate in an
abstraction. Although there is evidence to suggest even ungrounded LMs can
learn grounded concepts such as spatial relations (Patel & Pavlick, 2021), we
did not rigorously test this. We are excited to explore additional feature
representations (perhaps multimodal visualization interfaces in addition to
language representations) that can be jointly used to construct state
abstractions.

##  9 Acknowledgements

We thank the Boston Dynamics AI Institute for providing robotic hardware
resources used in this project, as well as Nishanth Kumar for help in running
(late-night) demonstration videos. We thank members of the Language and
Intelligence and Interactive Robotics Groups for helpful feedback and
discussions. Andi Peng is supported by the NSF Graduate Research Fellowship
and Open Philanthropy. Belinda Li and Theodore Sumers are supported by
National Defense Science and Engineering Graduate Fellowships. Ilia
Sucholutsky is supported by a NSERC Fellowship (567554-2022).

## References

  * Abel et al. (2016) David Abel, David Hershkowitz, and Michael Littman.  Near optimal behavior via approximate state abstraction.  In _International Conference on Machine Learning_ , pp. 2915–2923, 2016.
  * Abel et al. (2018) David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman.  State abstractions for lifelong reinforcement learning.  In _International Conference on Machine Learning_ , pp. 10–19, 2018.
  * Ahn et al. (2022) Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al.  Do as I can, not as I say: Grounding language in robotic affordances.  _Conference on Robot Learning_ , 2022.
  * Bahdanau et al. (2015) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  In _International Conference on Learning Representations_ , 2015.
  * Bobu et al. (2021) Andreea Bobu, Marius Wiggert, Claire Tomlin, and Anca D. Dragan.  Feature expansive reward learning: Rethinking human input.  In _International Conference on Human-Robot Interaction_ , pp. 216–224, 2021.
  * Bobu et al. (2023) Andreea Bobu, Yi Liu, Rohin Shah, Daniel S. Brown, and Anca D. Dragan.  SIRL: similarity-based implicit representation learning.  _International Conference on Human-Robot Interaction_ , 2023.
  * Bobu et al. (2024) Andreea Bobu, Andi Peng, Pulkit Agrawal, Julie Shah, and Anca D Dragan.  Aligning robot and human representations.  _International Conference on Human-Robot Interaction_ , 2024.
  * Cakmak & Thomaz (2012) Maya Cakmak and Andrea L Thomaz.  Designing robot learners that ask good questions.  In _International conference on Human-Robot Interaction_ , pp. 17–24, 2012.
  * Carta et al. (2022) Thomas Carta, Pierre-Yves Oudeyer, Olivier Sigaud, and Sylvain Lamprier.  Eager: Asking and answering questions for automatic reward shaping in language-guided rl.  _Advances in Neural Information Processing Systems_ , 35:12478–12490, 2022.
  * Chater & Vitányi (2003) Nick Chater and Paul Vitányi.  Simplicity: a unifying principle in cognitive science?  _Trends in cognitive sciences_ , 7(1):19–22, 2003.
  * Co-Reyes et al. (2018) John D Co-Reyes, Abhishek Gupta, Suvansh Sanjeev, Nick Altieri, Jacob Andreas, John DeNero, Pieter Abbeel, and Sergey Levine.  Guiding policies with language via meta-learning.  _arXiv preprint arXiv:1811.07882_ , 2018.
  * Coates & Ng (2012) Adam Coates and Andrew Y Ng.  Learning feature representations with k-means.  In _Neural Networks: Tricks of the Trade: Second Edition_ , pp. 561–580. Springer, 2012.
  * Correia & Alexandre (2023) André Correia and Luís A Alexandre.  A survey of demonstration learning.  _arXiv preprint arXiv:2303.11191_ , 2023.
  * Diuk et al. (2008) Carlos Diuk, Andre Cohen, and Michael L Littman.  An object-oriented representation for efficient reinforcement learning.  In _International Conference on Machine Learning_ , pp. 240–247, 2008.
  * Du et al. (2023) Yuqing Du, Olivia Watkins, Zihan Wang, Cédric Colas, Trevor Darrell, Pieter Abbeel, Abhishek Gupta, and Jacob Andreas.  Guiding pretraining in reinforcement learning with large language models.  In _International Conference on Machine Learning_ , 2023.
  * Fan et al. (2020) Judith E Fan, Robert D Hawkins, Mike Wu, and Noah D Goodman.  Pragmatic inference and visual abstraction enable contextual flexibility during visual communication.  _Computational Brain & Behavior_, 3:86–101, 2020.
  * Goyal et al. (2019) Prasoon Goyal, Scott Niekum, and Raymond J Mooney.  Using natural language for reward shaping in reinforcement learning.  _International Joint Conference on Artificial Intelligence_ , 2019.
  * Higgins et al. (2017) Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  beta-vae: Learning basic visual concepts with a constrained variational framework.  In _International Conference on Learning Representations_ , 2017.
  * Ho et al. (2022) Mark K Ho, David Abel, Carlos G Correa, Michael L Littman, Jonathan D Cohen, and Thomas L Griffiths.  People construct simplified mental representations to plan.  _Nature_ , 606(7912):129–136, 2022.
  * Ho et al. (2023) Mark K Ho, Jonathan D Cohen, and Thomas L Griffiths.  Rational simplification and rigidity in human planning.  _PsyArXiv_ , Mar 2023.
  * Huang et al. (2022a) Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch.  Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.  _arXiv preprint arXiv:2201.07207_ , 2022a.
  * Huang et al. (2022b) Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, et al.  Inner monologue: Embodied reasoning through planning with language models.  _arXiv preprint arXiv:2207.05608_ , 2022b.
  * Huey et al. (2023) Holly Huey, Xuanchen Lu, Caren M Walker, and Judith E Fan.  Visual explanations prioritize functional properties at the expense of visual fidelity.  _Cognition_ , 236:105414, 2023.
  * Jiang et al. (2022) Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan.  Vima: General robot manipulation with multimodal prompts.  _International Conference on Machine Learning_ , 2022.
  * Karamcheti et al. (2023) Siddharth Karamcheti, Suraj Nair, Annie S Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang.  Language-driven representation learning for robotics.  _Robotics: Science and Systems_ , 2023.
  * Kirillov et al. (2023) Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al.  Segment anything.  _arXiv preprint arXiv:2304.02643_ , 2023.
  * Kwon et al. (2023) Minae Kwon, Sang Michael Xie, Kalesha Bullard, and Dorsa Sadigh.  Reward design with language models.  _International Conference on Learning Representations_ , 2023.
  * Laskey et al. (2017) Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg.  Dart: Noise injection for robust imitation learning.  In _Conference on Robot Learning_ , pp. 143–156, 2017.
  * Lee et al. (2021) Kimin Lee, Laura Smith, and Pieter Abbeel.  Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.  _arXiv preprint arXiv:2106.05091_ , 2021.
  * Li et al. (2023) Belinda Z Li, William Chen, Pratyusha Sharma, and Jacob Andreas.  Lampp: Language models as probabilistic priors for perception and action.  _arXiv e-prints_ , pp. arXiv–2302, 2023.
  * McCarthy et al. (2021a) William P McCarthy, Robert D Hawkins, Haoliang Wang, Cameron Holdaway, and Judith E Fan.  Learning to communicate about shared procedural abstractions.  _arXiv preprint arXiv:2107.00077_ , 2021a.
  * McCarthy et al. (2021b) William P McCarthy, Marcelo G Mattar, David Kirsh, and Judith E Fan.  Connecting perceptual and procedural abstractions in physical construction.  In _Proceedings of the Annual Meeting of the Cognitive Science Society_ , volume 43, 2021b.
  * Misra et al. (2018) Dipendra Misra, Andrew Bennett, Valts Blukis, Eyvind Niklasson, Max Shatkhin, and Yoav Artzi.  Mapping instructions to actions in 3D environments with visual goal prediction.  In _Conference on Empirical Methods in Natural Language Processing_ , pp. 2667–2678, 2018.
  * Mu et al. (2022) Jesse Mu, Victor Zhong, Roberta Raileanu, Minqi Jiang, Noah Goodman, Tim Rocktäschel, and Edward Grefenstette.  Improving intrinsic exploration with language abstractions.  _Advances in Neural Information Processing Systems_ , 2022.
  * OpenAI (2023) OpenAI.  GPT-4 technical report, 2023.
  * Patel & Pavlick (2021) Roma Patel and Ellie Pavlick.  Mapping language models to grounded conceptual spaces.  In _International Conference on Learning Representations_ , 2021.
  * Peng et al. (2023) Andi Peng, Aviv Netanyahu, Mark K Ho, Tianmin Shu, Andreea Bobu, Julie Shah, and Pulkit Agrawal.  Diagnosis, feedback, adaptation: A human-in-the-loop framework for test-time policy adaptation.  2023\.
  * Pomerleau (1988) Dean A Pomerleau.  Alvinn: An autonomous land vehicle in a neural network.  _Advances in Neural Information Processing Systems_ , 1, 1988.
  * Puterman (2014) Martin L Puterman.  _Markov decision processes: discrete stochastic dynamic programming_.  John Wiley & Sons, 2014.
  * Reimers & Gurevych (2019) Nils Reimers and Iryna Gurevych.  Sentence-bert: Sentence embeddings using siamese bert-networks.  _arXiv preprint arXiv:1908.10084_ , 2019.
  * Ross et al. (2011) Stéphane Ross, Geoffrey Gordon, and Drew Bagnell.  A reduction of imitation learning and structured prediction to no-regret online learning.  In _Conference on Artificial Intelligence and Statistics_ , pp. 627–635, 2011.
  * Sharma et al. (2022) Pratyusha Sharma, Antonio Torralba, and Jacob Andreas.  Skill induction and planning with latent language.  In _Association for Computational Linguistics_ , pp. 1713–1726, 2022.
  * Shinn et al. (2023) Noah Shinn, Beck Labash, and Ashwin Gopinath.  Reflexion: an autonomous agent with dynamic memory and self-reflection.  _arXiv preprint arXiv:2303.11366_ , 2023.
  * Shridhar et al. (2022) Mohit Shridhar, Lucas Manuelli, and Dieter Fox.  Cliport: What and where pathways for robotic manipulation.  In _Conference on Robot Learning_ , pp. 894–906, 2022.
  * Sumers et al. (2021) Theodore R Sumers, Mark K Ho, Robert D Hawkins, Karthik Narasimhan, and Thomas L Griffiths.  Learning rewards from linguistic feedback.  In _AAAI Conference on Artificial Intelligence_ , volume 35, pp. 6002–6010, 2021.
  * Thomas & Barto (2011) Philip S Thomas and Andrew G Barto.  Conjugate markov decision processes.  In _International Conference on Machine Learning_ , pp. 137–144, 2011.
  * Wang et al. (2023) Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar.  Voyager: An open-ended embodied agent with large language models.  _arXiv preprint arXiv:2305.16291_ , 2023.
  * Zeng et al. (2023) Andy Zeng, Maria Attarian, brian ichter, Krzysztof Marcin Choromanski, Adrian Wong, Stefan Welker, Federico Tombari, Aveek Purohit, Michael S Ryoo, Vikas Sindhwani, Johnny Lee, Vincent Vanhoucke, and Pete Florence.  Socratic models: Composing zero-shot multimodal reasoning with language.  In _International Conference on Learning Representations_ , 2023.
  * Zhou et al. (2022) Xingyi Zhou, Rohit Girdhar, Armand Joulin, Philipp Krähenbühl, and Ishan Misra.  Detecting twenty-thousand classes using image-level supervision.  In _European Conference on Computer Vision_ , pp. 350–368, 2022.

##  Appendix A Appendix

###  A.1 Environment Details

States are fully-observable images and represent a RGB fixed-camera topdown
view of the scene. The action space is continuous of dimension 4 and consists
of high-level pick-and-place actions parameterized by the pose of the end
effector. There are 4 possible objects that can be spawned: the target
(manipulated) object, goal object, and (up to) two possible distractors. There
are two types of target features (object type and texture) and 29 possible
instantiations of object type and 81 instantiations of texture. Visual
depictions of the features can be found in the user study below.

###  A.2 Task Details

We provide details regarding all tasks, including ground truth task
distributions and full LM prompt.

Pick-and-place:

red heart:

  * •

Task prompt: Bring me the red heart.

  * •

True distribution: {objects: heart}, {textures: red, dark red, dark red swirl,
red paisley}

heart:

  * •

Task prompt: Bring me the heart.

  * •

True distribution: {objects: heart}, {textures: ALL}

tiger-colored object:

  * •

Task prompt: Bring me the tiger-colored object.

  * •

True distribution: {objects: ALL}, {textures: tiger}

letter from the word letter:

  * •

Task prompt: Bring me a letter from the word ‘letter’.

  * •

True distribution: {objects: letter E, letter R, letter T}, {textures: ALL}

consonant with warm-color:

  * •

Task prompt: Bring me a consonant with a warm color on it.

  * •

True distribution: {objects: letter G, letter M, letter R, letter T, letter
V}, {textures: dark {red—yellow—pink—orange}, dark {red—yellow—pink—orange}
and * stripe, {red—yellow—pink—orange} and * polka dot, dark
{red—yellow—pink—orange} and * polka dot, {red—yellow—pink} swirl, dark
{red—yellow—pink} swirl, {red—yellow—pink} paisley, tiger, magma, wooden,
rainbow, tiles, brick}

vowel with multiple colors:

  * •

Task prompt: Bring me a vowel with multiple colors on it.

  * •

True distribution: {objects: letter A, letter E}, {textures: polka dot, tiles,
checkerboard, plastic, tiger, magma, rainbow, * and * stripe, * and * polka
dot, * swirl, * paisley}

something to drink water out of:

  * •

Task prompt: Bring me something to drink water out of.

  * •

True distribution: {objects: bowl, pan, container}, {textures: ALL}

something from a typical kitchen:

  * •

Task prompt: Bring me something from a typical kitchen.

  * •

True distribution: {objects: bowl, pan, container}, {textures: ALL}

Rotate:

block:

  * •

Task prompt: Rotate the block.

  * •

True distribution: {objects: block, small block, shorter block, L shaped
block}, {textures: ALL}

something to drink water out of:

  * •

Task prompt: Rotate something to drink water out of.

  * •

True distribution: {objects: container, bowl, pan}, {textures: red, dark red,
dark red swirl, red paisley}

Sweep While Avoid:

block, pan:

  * •

Task prompt: Sweep the block without touching the pan.

  * •

True distribution: {objects: block, small block, shorter block, L shaped
block, pan}, {textures: ALL}

block, line:

  * •

Task prompt: Sweep the block without touching the line.

  * •

True distribution: {objects: block, small block, shorter block, L shaped
block, line}, {textures: ALL}

####  A.2.1 Full Prompt

ChatGPT models (including GPT4) can take in both system prompts and user
prompts. We split our prompt into these two parts as follows.

System prompt where {object_list} is replaced by the list of all object types
in the environment and {object_colors} by the list of all colors and textures:

> You are interfacing with a robotics environment that has a robotic arm
> learning to manipulate objects based on some linguistic command (e.g. “pick
> up the red bowl”). At each interaction, the researcher will specify the
> command that you need to teach the robot. In order to teach the robot, you
> will need to help design the training distribution by specifying what
> properties task-relevant objects can have based on the given command.
> Objects in this environment have two properties: object type, object color.
> Any object type can be paired with any color, but an object can only take on
> exactly one object type and exactly one color.
> Object types:
> {object_list}
> Object colors:
> {object_colors}

User prompt where {rule} is replaced by one of the task prompts listed above,
{group} is replaced by “object color” or “object type”, and {candidate} is
replaced by each candidate object color or type that we would like the LM to
evaluate:

> The command is “{rule}”. In an instantiation of the environment that
> contains only some subset of the object types and colors, could the target
> object have {group} “{candidate}”? Think step-by-step and then finish with a
> new line that says “Final answer:” followed by “yes” or “no”.

###  A.3 Architecture and Training Details

Architecture.

For GCBC, Sentence-BERT processes a goal utterance in natural language into an
embedding of size 384, which we additionally process through a linear MLP of
output size 100. We process the state using a standard Conv2D block consisting
of 3 stacked Conv2D layers of output channel sizes 32, 64, 32, and strides 4,
2, and 1. Each output layer is processed by BatchNorm2D as well as a ReLU
activation. After the last Conv2D layer, we flatten the output and concatenate
with the output of the goal utterance. We feed the concatenated output through
a last linear layer for action prediction.

For LGA, we process both the state and the state abstraction through dual
Conv2Ds blocks like above. We concatenate the output and feed through a last
linear layer for action prediction.

Training.

We train all networks to convergence for a maximum of 750 epochs. All
computation was done on two NVIDIA GeForce RTX 3090 GPUs. Rollouts were
rendered locally using PyBullet on commercial hardware (a Macbook Pro).

###  A.4 User Study

In the following pages, we have included the full user study shared with
participants. Following standard user study procedure, we initial briefed
users by telling them how long the study was (roughly 30 mins) and that they
were free to leave at anytime. Demographic information was collected in
person. We showed participants the first pdf during the familiarization phase
to introduce them to the environment and full feature list. We then randomly
chose to begin with either the human-only task (second pdf) or the LM-aided
task (third pdf). After the study, users were debriefed and given the email of
the study designer to contact if they had any questions.

See pages - of figures/dictionary.pdf See pages - of figures/human.pdf See
pages - of figures/llm.pdf