SlowGuess/NAACL_2022_OCR huggingface dataset:
Charlie Snell
Yi Su
{csnell22,sherryy,justinfu,suyi,svlevine}@berkeley.edu
Goal-oriented dialogue systems face a tradeoff between fluent language generation and task-specific control. While supervised learning with large language models is capable of producing realistic text, how to steer such responses towards completing a specific task without sacrificing language quality remains an open question. In this work, we formulate goal-oriented dialogue as a partially observed Markov decision process, interpreting the language model as a representation of both the dynamics and the policy. This view allows us to extend techniques from learning-based control, such as task relabeling, to derive a simple and effective method to finetune language models in a goal-aware way, leading to significantly improved task performance. We additionally introduce a number of training strategies that serve to better focus the model on the task at hand. We evaluate our method, Context-Aware Language Models (CALM), on a practical flight-booking task using AirDialogue. Empirically, CALM outperforms the state-of-the-art method by  $7\%$  in terms of task success, matching human-level task performance.
Dialogue systems have typically approached the problem of generating realistic dialogue from the perspective of supervised learning (Dusek and Jurcicek, 2016; Eric and Manning, 2017; Mei et al., 2017; Chen et al., 2019; Wu et al., 2019a; Hosseini-Asl et al., 2020; Peng et al., 2020; Adiwardana et al., 2020). However, dialogue can also be viewed as a sequential decision making process, which is well-suited to planning and reinforcement learning (RL) algorithms. A challenge with the classical RL approach to dialogue is the requirement for active interaction with humans (Gašić et al., 2011). Training such a system with active human-in-the-loop
Figure 1: CALM is an end-to-end language model for goal oriented dialogue. CALM's training objective teaches the model to better pay attention to the dialogue task context, yielding a  $\sim 50\%$  improvement in task success over standard LM training on a flight booking task.
interaction quickly becomes expensive and cumbersome, making it desirable to develop techniques for goal-directed training of dialogue systems that can effectively leverage offline data.
While many dialogue generation techniques based on RL and learned control have been proposed (Eckert et al., 1997; Levin et al., 2000; Chung, 2004; Georgila et al., 2006; Schatzmann et al., 2007; Heeman, 2009; Georgila and Traum, 2011), most such systems take a pipelined approach, where an abstract representation of states and actions is designed by hand and then combined with RL to train a "dialogue management" system, rather than generating dialogue end-to-end. These pipelined approaches rely on a manually designed decomposition of the dialogue task, which may be domain-specific and, more importantly, may not enjoy all of the benefits of tightly integrating low-level text generation with the overall goals of the
Code at https://sea-snell.github.io/CALM_LM_site/
Findings of the Association for Computational Linguistics: NAACL 2022, pages 2351 - 2366
task. In this work, we instead ask: how can we scalably and effectively introduce the mechanisms of goal-directed decision making into end-to-end language models, to directly steer language generation toward completing specific dialogue tasks rather than simply generating probable responses?
To this end, rather than utilizing a pipelined approach, we aim to directly finetune language models in a task-aware manner such that they can maximize a given utility function. We observe that large language models can already be formulated within a Markov decision processes (MDP) as capturing both the dynamics and policy for a decision-making task, where dialogue history serves as state, and the agent's utterances serve as actions. We could utilize this observation by finetuning the models directly with online RL, but the need for human-in-the-loop training makes this difficult. Offline RL methods (Levine et al., 2020; Fujimoto et al., 2019; Wu et al., 2019b; Wang et al., 2020b) provide an alternative approach, but typically require value function estimation, which is not straightforward to perform with a language model. Instead, we propose a conditional imitation learning strategy coupled with a novel task relabeling approach that can finetune language models from offline data, such that the model still represents the joint distribution over dialogues, but tilts this distribution toward dialogues with a high reward. This amounts to a task-aware finetuning strategy that integrates task information into the model. The main contribution of our work is CALM (Context-Aware Language Modeling), a framework for end-to-end goal-directed dialogue generation. CALM unifies the traditional language modeling objective with task-specific supervision, where a language model is interpreted as a joint representation of dynamics and policies in an MDP, and the finetuning process utilizes a conditional imitation learning objective with a novel task relabeling strategy that teaches the model how to generate high-utility dialogues (see Figures 1 and 2). Because CALM interprets the language model as both a dynamics model and a policy, it can be used as either a model-free method, where the dynamics are discarded and the policy component is used to greedily generate responses, or as a model-based method, where the dynamics component can be used to plan at test-time. We empirically evaluate CALM on AirDialogue (Wei et al., 2018), the largest dataset for goal-oriented dialogue based on a flight-booking task. CALM
improves the task success by  $10\%$  over the previous state-of-the-art method (Chen et al., 2020) following the evaluation protocol proposed by Wei et al. (2018), achieving the first-ever human-level performance on this dataset.
Our goal is to enable end-to-end training of goal-directed dialogue agents. In these settings, an agent aims to complete a particular task with its utterances (Smith and Hipp, 1994). Goal-directed agents have been explored in contexts such as personal assistants (McTear, 2002; Budzianowski et al., 2018; Williams et al., 2014), recommendation systems (Liu et al., 2010; Kang et al., 2019), education (Yuan et al., 2008), and negotiation (He et al., 2018; Lewis et al., 2017). While there are multiple approaches to constructing dialogue agents, in this work we frame the problem of generating dialogue as a sequential decision making problem within a (partially observed) Markov Decision Process (MDP) (Singh et al., 1999; Young et al., 2013). Prior works that utilize such an MDP formulation typically aim to train a dialogue management system (Singh et al., 2002), in which the agent reasons about higher-level abstractions of the state of the conversation, and language generation is performed using a downstream procedure. Dialogue management systems have been trained using techniques such as online reinforcement learning via policy gradients (Gašić et al., 2011; He et al., 2018), off-policy reinforcement learning (Pietquin et al., 2011; Yu et al., 2016) or actor-critic methods (Su et al., 2017). Our method differs from dialogue management systems in that CALM is an end-to-end system optimized for successful task completion, and performs both high-level decision making and language generation.
Recent advancements in language models, such as recurrent neural networks (Sundermeyer et al., 2012; Asri et al., 2016; Su et al., 2016; Zhao et al., 2019; Wang et al., 2020a; Zhang et al., 2020) and attention-based architectures (Vaswani et al., 2017; Liu et al., 2019; Devlin et al., 2018; Brown et al., 2020), have spurred increasing interest in such end-to-end dialogue systems (Hosseini-Asl et al., 2020; Peng et al., 2020; Adiwardana et al., 2020). Model-based approaches, in which a learned agent is substituted for a human, allow learning to be done entirely within simulation without human intervention (Li et al., 2016; He et al., 2018; Kang et al., 2019; Lewis et al., 2017; Liu et al., 2018). In con
Figure 2: A visual outline of CALM. We apply Task Relabeling to our static offline dataset, by swapping out the task context — in this case a flight table — such that the attached dialogue becomes an optimal example of task completion. When fine-tuning on this relabeled dataset, we then apply a Task Specific Auxiliary Loss on top of the standard language modeling objective; this helps the model learn to use the task context. Once trained, CALM can consistently solve goal-directed dialogue tasks.
trast to these approaches, CALM augments the traditional language modeling objective with task-specific rewards in order to finetune a model that is more aware of task goals, which significantly improves performance over a naïve language model without the need for simulating human responses in an interactive training loop. Jaques et al. (2019) recently proposed a model-free, offline approach to undirected dialogue, or dialogue without a specific task goal. Our method differs in that we aim to solve goal-oriented dialogue which allows us to optimize task-specific objectives, and that we take a model-based RL approach which enables us to leverage fine-tuned language models.
In this section, we review our notation and problem formulation for casting dialogue within a sequential decision making framework.
POMDP formulation. We formulate dialogue generation as a partially observable Markov decision process (POMDP) (Kaelbling et al., 1998), with a state that consists of known and unknown context information about the task. Let  $c_{h} \in \mathcal{C}^{(h)}$  denote the hidden context for the task, and let  $c_{o} \in \mathcal{C}^{(o)}$  denote the observed context. For instance, in a flight booking task, a table of available flights might correspond to  $c_{o}$ , while the particular flight that the human wants to book, which is unknown to the agent, corresponds to  $c_{h}$ . Note that the reward, which requires booking the right flight, depends on both hidden and observed contexts. We can define such an environment as a POMDP  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{O},\mathcal{T},\mathcal{Z},\mu_0,\mathcal{R},\gamma)$ . We denote a conversation  $\tau$  as  $\tau := \{a_0,e_0,\dots,a_T\}$ , where  $T$  denotes the number of turns in a conversation and  $a_{t}$  and  $e_t$  represent utterances (strings of tokens) from the dialogue agent  $(a_{t})$  and the human  $(e_t)$  at the  $t$ -th turn, respectively. We additionally use  $\tau_{< t}$  to denote
conversation history up to the  $t$ -th turn. We can represent the underlying POMDP state  $s_t \in S$  as the concatenation of both of the contexts and the previous conversation history  $s_t := \{c_h, c_o, \tau_{< t}\} = \{c_h, c_o, a_0, e_0, \dots, a_{t-1}, e_{t-1}\}$ . However, we only observe the last two elements of the state tuple, such that our observation  $o_t \in \mathcal{O}$  at the  $t$ -th conversation turn is  $o_t = \{c_o, \tau_{< t}\}$ . An action  $a_t \in \mathcal{A}$  is the agent's response to the current state  $s_t$ . Given our definition of the state, the full conversation in a dialogue can be conveniently represented by the last observation and action,  $\{o_T, a_T\}$ . An agent  $\pi: \mathcal{O} \to \mathcal{P}(\mathcal{A})$  maps observations to sets of probability measures over the action space  $\mathcal{P}(\cdot)$ . A transition function  $\mathcal{T}(\cdot | s_t, a_t)$ , represents a distribution over the human's utterances, returning  $s_{t+1}$  as the state at turn  $t+1$ . We only consider the sparse reward setting with  $r_T = R(s_T, a_T) \in \{0, 1\}$  denoting task completion, and  $r_t = 0, \forall t < T$ . Our final reward is therefore dependent on both the context and the dialogue:  $R(s_T, a_T) = R(\tau, c_h, c_o)$ , where the context  $\{c_o, c_h\}$  is randomly sampled for each dialogue from some initial distribution  $\mu_0$ .
Goal-oriented dialogue. Goal-oriented dialogue systems aim to maximize the expected reward of the above POMDP
$$
\mathbb {E} _ {\left\{c _ {o}, c _ {h} \right\} \sim \mu_ {0}, \pi , \mathcal {T}} \left[ \sum_ {t = 0} ^ {T} \gamma^ {t} R \left(s _ {t}, a _ {t}\right) \right], \tag {1}
$$
where  $\{c_h, c_o\}$  is sampled from distribution  $\mu_0$ . On-policy RL algorithms optimize this objective via environment interaction, which is represented by a real human. However, because human-in-the-loop training is expensive, we pursue an offline learning approach where we are given a fixed dataset and there is no further interaction with the human in the learning process. This dataset is composed of  $n$  trajectories with  $\mathcal{D}^{\mathrm{off}} = \{c_h^{(i)}, c_o^{(i)}, \tau^{(i)}, r^{(i)}\}_{i=1}^n$  with each  $\tau^{(i)} = \{a_0^{(i)}, e_0^{(i)}, \dots, a_T^{(i)}\}$  and its corresponding final reward for task completion  $r^{(i)}$ . Our
goal is to learn the policy  $\pi(a|o)$  which improves the dialog agent's ability in achieving the highest task reward defined in Equation 1.
Language models. While conventionally a language model is seen simply as a sequence model over tokens of the form  $\prod_{t=1}^{T} p(x_{t+1} | x_{1:t})$ , when the sequence  $x_{1:T}$  corresponds to a dialogue trajectory  $\tau$ , we can also interpret a language model as learning the distribution over  $\tau$ . This distribution can be factored into the product of the policy  $\pi(a_t | \tau_{<t})$  and the dynamics  $\mathcal{T}(\tau_{<t+1} | \tau_{<t}, a_t)$ , and so we can say that a language model also represents the policy and the dynamics. Therefore, the maximum likelihood objective for training or finetuning a language model on a dialogue dataset  $\mathcal{D}^{\mathrm{off}}$  consisting of dialogue trajectories  $\tau$  can be written as
$$
\begin{array}{l} \mathcal {L} _ {L M} (\theta) = \max _ {\theta} \underset {\tau \sim \mathcal {D} ^ {\mathrm {o f f}}} {\mathbf {E}} \sum_ {t = 1} ^ {T} \Bigg (\log \pi_ {\theta} (a _ {t} | \tau_ {<   t}) \\ \left. + \log \mathcal {T} _ {\theta} \left(\tau_ {<   t + 1} \mid \tau_ {<   t}, a _ {t}\right)\right), \tag {2} \\ \end{array}
$$
where  $\pi_{\theta}(a_t|o_t)$  represents a policy that generates new dialogue based on the observed context and dialogue history, and  $\mathcal{T}_{\theta}(\tau_{< t + 1}|\tau_{< t},a_t)$  represents the observed dynamics characterizing human responses, and  $\theta$  denotes parameters in  $\pi$  and  $\mathcal{T}$ . Note that  $\tau_{< t}$  consists only of the conversation history, and does not contain any task-specific context. A naive approach to train dialogue systems is to jointly parameterize both  $\pi$  and  $\mathcal{T}$  as one language model, and optimize Equation 2 on pre-collected conversations  $\mathcal{D}^{\mathrm{off}}$ . This method corresponds to behavioral cloning (BC) (Pomerleau, 1989).
Context conditioning. While an agent trained using Equation 2 can learn policies and dynamics that imitate human conversations, this objective does not incorporate the task goal, and may not produce a policy that is more performant than the dataset  $\mathcal{D}^{\mathrm{off}}$ . While it is possible to input  $c_{o}$  into the language model to maximize the conditional probability of  $P(\tau | c_{o})$  using a conditional version of the language modeling objective,  $\mathcal{L}_{CTX}(\theta)$ ,
$$
\begin{array}{l} \mathcal {L} _ {C T X} (\theta) = \max _ {\theta} \underset {(\tau , c _ {o}) \sim \mathcal {D} ^ {\mathrm {o f f}}} {\mathbf {E}} \sum_ {t = 1} ^ {T} \Bigl (\log \pi_ {\theta} (a _ {t} | \tau_ {<   t}, c _ {o}) \\ \left. + \log \mathcal {T} _ {\theta} \left(o _ {t + 1} \mid \tau_ {<   t}, a _ {t}, c _ {o}\right)\right), \tag {3} \\ \end{array}
$$
contexts with particular task structures (e.g., a set of entries in a table) may not be simply processed as a sequence similarly to  $\tau$ . Additionally, the language model is not pretrained to read structured
context, and oftentimes the recent dialogue history is much more predictive of the next utterance than the task context is. As a result, language models can ignore the task context and only learn  $P(\tau)$  despite being conditioned on  $c_{o}$ . Our approach builds on this conditional modeling approach, but makes a number of improvements that allow it to be more aware of the context information, which attains significantly better results in our experiments.
4 Context-Aware Language Modeling
In this section, we present our method for goal-oriented dialogue systems, Context-Aware Language Modeling (CALM). CALM interprets a language model as a combination of a policy and a dynamics model in the POMDP formulation of a dialogue task, as described in Section 3. Under this interpretation, naive supervised finetuning on the dialogue dataset can be viewed as behavioral cloning (BC) (Pomerleau, 1989). However, BC only imitates data and does not necessarily produce a good policy in terms of completing tasks. We propose to improve the policy by utilizing a task relabeling strategy (described in Section 4.1), analogous to prior task relabeling approaches (Kaelbling, 1993; Andrychowicz et al., 2017; Pong et al., 2018; Savinov et al., 2018; Ghosh et al., 2019; Lynch et al., 2020; Eysenbach et al., 2020). This relabeling procedure augments the data with examples of near-optimal utterances, making the language model more task-aware. However, we find several shortcomings with this approach alone and propose the following improvements. First, an expressive language model is liable to ignore the task context, which we address by proposing an auxiliary loss (Section 4.2) that forces the model to utilize this information. Second, learning from structured task information is difficult and can result in models that fail to capture complex task structure, so we propose a task pre-training procedure to improve the learnability (Section 4.3). Finally, to further improve performance we use a model-based planning procedure (Section 4.4) on top of the proposed method that samples multiple dialogues in parallel and selects the most promising candidates.
4.1 Dialogue Task Relabeling
$\mathcal{L}_{CTX}(\theta)$  defines a context-conditional maximum likelihood objective for training an expert imitation policy in conjunction with a dynamics model. However, simply imitating all the dialogue data does not necessarily produce the best possible policy. We would like to learn a policy that pro
duces dialogue that is more optimal, in the sense of better maximizing the task utility, than the average dialogue in the dataset. Task relabeling enables us to learn from optimal trajectories without simply filtering the dataset for high-reward trajectories, which would unnecessarily discard potentially informative data. In the case of dialogue, we can perform task relabeling by considering the context  $\{c_{o}, c_{h}\}$  as defining the task. While a given dialogue may be unsuccessful for the context for which it was collected, it could be considered successful under a different context. In this case, we can simply swap out  $\{c_{o}, c_{h}\}$  to create optimal task examples from the many sub-optimal examples provided by  $\mathcal{D}^{\mathrm{off}}$ . Since our reward  $R(c_{h}, c_{o}, \tau)$  is a function of the dialogue and context, we can modify the reward for a given dialogue just by changing the given observed context  $c_{o}$ . Using this observation, we can relabel unsuccessful dialogues with successful ones, and even for already successful dialogues there may be multiple  $c_{o}$  corresponding to task success, allowing us to augment the number of successful  $(c_{h}, c_{o}, \tau)$  tuples.
Formally, since our POMDP includes a prior distribution over contexts  $\{c_h, c_o\} \sim \mu_0$ , there exists a posterior  $q(c_o|\tau, c_h)$  over observed contexts that correspond to optimal task completion under a given  $\tau$ . We can then re-label  $\tau$  to be optimal under its context by sampling a new  $c_o$  from  $q(c_o|\tau, c_h)$ . In practice, this sampling is performed by rejection sampling from either  $\mu_0$  or some  $P(c_o|c_h)$ ; the latter, lower entropy distribution, can be preferred if there is a low probability of sampling valid, high-reward contexts under  $\mu_0$ . Now, given any  $\tau$  from an offline dataset of dialogues, we can learn from the full distribution of contexts corresponding to optimal task completion under this dialogue.
In order for this relabeling procedure not to bias our policy towards behavior that is overly optimistic about the user's responses, it is necessary that the distribution of these responses in our dataset does not depend on the portion of the context that is relabeled. For example, relabeling the table of available flights for a flight booking task should generally be reasonable, because the user is usually unaware of the flight table. On the other hand, relabeling the desired flight would not make sense, since the user's utterances are strongly depend on this. To provide another example, in a bargaining task (Lewis et al., 2017), the agent might fail to obtain the desired item and instead get an
item of lesser value. But relabeling with a context that assigns a higher value to the item received would not lead to a reasonable example, since the agent mainly received this item as a result of the user's responses rather than as a result of their own bargaining skill.
Methods based on similar principles have previously been proposed in the deep RL community for simple parametric tasks, such as goal-reaching or linearly-parameterized reward functions (Kaelbling, 1993; Andrychowicz et al., 2017; Eysenbach et al., 2020). However, the dialogue task relabeling that we employ is particularly effective in our setting, since there may be exponentially many contexts that are optimal for a given dialogue (e.g., many different flight tables for a flight booking task), in contrast to the simpler task parameterizations used in prior work, where for example only one goal might be optimal for a given trajectory (the one that is reached). As a result, this technique not only allows us to turn sub-optimal task data into optimal data, but it also allows us to greatly increase the number of optimal task examples from which we can learn, which we will show leads to a large performance improvement.
4.2 Task-Specific Auxiliary Loss
Goal-oriented dialogue generation can be viewed as learning the conditional distribution  $P(\tau | c_o)$ , where  $\tau$  represents the generated dialogue given a specific context  $c_o$ . However when trained naively, language models are liable to ignore this conditioning context, instead focusing purely on the previous utterances in the dialogue. In this case, the model is effectively only learning  $P(\tau)$  despite having both the capacity and the context to learn the lower-entropy conditional distribution  $P(\tau | c_o)$ .
While dialogue tasks are by definition carried out through natural language, there is often an abstract high-level action  $\alpha_{h} \in \mathbf{A}$  that essentially determines the success of the task. In the case of the information retrieval task that we consider in this paper, these high-level actions correspond to deciding which database entity to retrieve for the user (e.g., suggesting a flight to the customer that meets all of their needs). While these high-level actions are theoretically learnable from correlations between the dialogue and the given context, in general, we find that learning these correlations corresponds to a relatively small decrease in dialogue entropy under the model. As a result, the model is less incentivized to learn these correlations rele
want to the task than the form of the dialogue. To address this issue, we incorporate an auxiliary objective into our training, which trains the model directly to predict the abstract high-level actions taken in the present dialogue. This objective effectively up-weights gradients relevant for learning the high-level actions, which further helps the model to utilize the context to solve the high-level task through dialogue.
For a given dialogue-context pair  $(\tau, \{c_h, c_o\})$  and high-level action,  $\alpha_h$ , our auxiliary objective is then simply to maximize the likelihood of the high-level actions taken in the dialogue: