SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SWIFTSAGE: A Generative Agent with Fast and Slow Thinking for Complex Interactive Tasks
Bill Yuchen Lin  $^{1}$  Yicheng Fu  $^{4}$  Karina Yang  $^{2}$  Faeze Brahman  $^{13}$  Shiyu Huang  $^{5}$  Chandra Bhagavatula  $^{1}$  Prithviraj Ammanabrolu  $^{67}$  Yejin Choi  $^{31}$  Xiang Ren  $^{21}$
$^{1}$ Allen Institute for Artificial Intelligence  
 $^{2}$ University of Southern California  $^{3}$ University of Washington  $^{4}$ Tsinghua University  
 $^{5}$ 4Paradigm Inc.  $^{6}$ University of California, San Diego  $^{7}$ MosaicML
https://swiftsage.github.io
We introduce SWIFT SAGE, a novel agent framework inspired by the dual-process theory of human cognition, designed to excel in action planning for complex interactive reasoning tasks. SWIFT SAGE integrates the strengths of behavior cloning and prompting large language models (LLMs) to enhance task completion performance. The framework comprises two primary modules: the SWIFT module, representing fast and intuitive thinking, and the SAGE module, emulating deliberate thought processes. The SWIFT module is a small encoder-decoder LM fine-tuned on the oracle agent's action trajectories, while the SAGE module employs LLMs such as GPT-4 for subgoal planning and grounding. We develop a heuristic method to harmoniously integrate the two modules, resulting in a more efficient and robust problem-solving process. In 30 tasks from the ScienceWorld benchmark, SWIFT SAGE significantly outperforms other methods such as SayCan, ReAct, and Reflexion, demonstrating its effectiveness in solving complex interactive tasks. $^{1}$
The advancement of artificial general intelligence is largely dependent on the development of agents that are proficient in complex interactive reasoning tasks. These agents should be capable of exhibiting problem-solving abilities akin to humans within dynamic, open-world environments [26, 7]. For example, the ScienceWorld benchmark [36] features a task where an agent must determine the electrical conductivity of an unknown object. In a simulated environment, the agent must navigate to appropriate rooms, locate and acquire essential items, such as batteries and light bulbs, build a circuit, perform an experiment, and interpret the results. Tackling such a complex interactive task demands agents to exhibit long-horizon planning, long-term memorization, subgoal decomposition, spatial reasoning, exception handling, and commonsense knowledge capabilities [37].
There are three primary approaches to developing agents capable of addressing complex interactive reasoning tasks: (1) (deep) reinforcement learning (RL), (2) behavior cloning (BC) [34] through sequence-to-sequence (seq2seq) learning [33], and (3) prompting large language models (LLMs) [6]. In addition to conventional RL methods such as DRRN [14], interactive reasoning can be framed as a seq2seq task, where the input text serves as the current state description and the output text corresponds to the subsequent action [9, 3]. By leveraging numerous gold trajectories generated by oracle agents, it becomes feasible to fine-tune Transformer models [35], like T5 [25], to effectively imitate the
$^{1}$ Contact: yuchenl@allenai.org
SayCan
Figure 1: Comparing methods of prompting LLMs to build agents for interactive tasks.
ReAct
SwiftSage
behavior of these oracle agents. Recent studies have also demonstrated that generative agents based on prompting LLMs, such as GPT-4, can produce reasonable plans and actions [18, 15, 32].
Although the aforementioned methods exhibit remarkable performance in relatively simple tasks, their ability to generalize to more complex and demanding tasks is limited. Both RL-based and seq2seq-based BC approaches effectively acquire knowledge from the environment through large-scale interactions and learn general action patterns from oracle agents. However, they face difficulties in decomposing tasks into subgoals, maintaining long-term memory, generalizing to unseen tasks, and handling exceptions. In contrast, instruction-tuned LLMs [24] demonstrate the ability to generate reasonable high-level plans for complex tasks and adapt their outputs based on human feedback. Yet, grounding their outputs to executable actions in the environment remains a challenge. These procedures also lack the capability to efficiently handle environment-specific exceptions that prevent agents from adhering to the LLM's plans. Additionally, previous methods such as SAYCAN [1], REACT [41] and REFLEXION [30] require a new inference with LLMs for each time step, making them considerably costly and inefficient (see Figure 1).
Inspired by the dual process theory [39, 16], we propose a novel framework that enables agents to closely emulate how humans solve complex, open-world tasks. The dual-process theory posits that human cognition is composed of two distinct systems: System 1, characterized by rapid, intuitive, and automatic thinking; and System 2, which entails methodical, analytical, and deliberate thought processes. System 1 is reminiscent of seq2seq methods, which learn through imitation of oracle agents and primarily operate utilizing shallow action patterns. Conversely, System 2 bears resemblance to LLMs that excel in applying commonsense knowledge, engaging in step-by-step reasoning, devising subgoal strategies, and exercising self-reflection. Thus, our proposed method, SWIFTSAGE, is designed to enable both fast and slow thinking in complex interactive reasoning tasks. It effectively integrates the strengths of behavior cloning (representing System 1) and prompting LLMs (emulating System 2), resulting in significant enhancements in task completion performance and efficiency.
Specifically, SWIFT SAGE consists of two primary modules: the SWIFT module and the SAGE module. The SWIFT module is a small encoder-decoder LM, fine-tuned on a T5-large (770m) checkpoint using the searched oracle trajectories of training tasks. It encodes short-term memory components, such as previous actions, observations, visited locations, as well as the current environment state. Then, it decodes the next individual action. This module simulates the fast, intuitive thinking characteristic of System 1. The SAGE module, representing the deliberate thinking of System 2, utilizes LLMs, such as GPT-4, and is structured around two prompting stages: planning and grounding. In the planning stage, we prompt LLMs to locate necessary items, plan and track subgoals, as well as detect and fix potential exceptions and mistakes. In the grounding stage, we focus on utilizing LLMs to transform the output subgoals derived from the planning stage into a sequence of actions by demonstrating potential action templates. Unlike prior methods, where LLMs only generate the next immediate action, our procedures engage in longer-term action planning. To harmoniously integrate the SWIFT and SAGE modules, we developed a heuristic algorithm that determines when to (de)activate the SAGE module and how to combine the outputs effectively with an action buffer mechanism.
In a comprehensive evaluation on 30 task types from the ScienceWorld benchmark, SWIFTSAGE significantly outperforms other methods, achieving a state-of-the-art average score of 84.7. In comparison, SAYCAN scores 33.8, REACT obtains 36.4, and REFLEXION reaches 45.3. Moreover, SWIFTSAGE is more cost-effective and efficient, requiring much fewer tokens per action for LLM inference than previous methods. This considerable performance advantage highlights the effectiveness and efficiency of the SWIFTSAGE framework in addressing complex interactive tasks.
2.1 Complex Interactive Reasoning
We define interactive reasoning as the problems where agents are tasked with accomplishing a goal within an interactive environment, typically simulated by engines such as AI2Thor [17] and TextWorld [11]. Our focus lies on the textual environment of ScienceWorld [36] and the complex interactive tasks it supports. Simple interactive tasks, like those created in ALFWorld [31] and TWC [21], primarily involve searching for and placing objects as well as performing basic actions within a single location. Many of these simple tasks have been almost solved by recent works.
In contrast, tasks in ScienceWorld exhibit greater complexity, characterized by more challenging task planning and a significantly larger action space (encompassing 10 locations, 200 types of objects with varying states, and 25 types of actions). Furthermore, agents may encounter random, unforeseen obstacles, such as broken stoves or missing soil, which hinder the execution of planned actions. As a result, agents must adapt and re-plan accordingly, for example, by seeking alternative heat sources or using a shovel on the outside ground to get soil. These challenges demand that agents possess skills in long-horizon planning, long-term memory, subgoal decomposition, exception handling, and commonsense knowledge—capabilities that are not explicitly required for simple interactive tasks.
2.2 Reinforcement Learning and Imitation Learning Methods
DRRN. Interactive tasks can naturally be framed as partially-observable Markov decision processes (POMDPs), enabling the application of RL-based methods. Deep Reinforced Relevance Network (DRRN) [14] is a standard baseline method to learn agents within text environment. It aims to learn representations of observations and actions separately and train a policy network to select actions from candidates based on feedback from the simulated environment. CALM [40] is a reranking-based method that combines DRRN with a causal language model (LM) fine-tuned with oracle transcripts. In essence, the causal LM captures task-specific and environment-specific knowledge through imitation learning, and the DRRN learns to rerank the predictions from the LM.
The KG-A2C [2] method uses an OpenIE technique [4] to represent environment states with graph structures and dynamically update these graphs. These graphs guide policy networks by constraining the combinations of action templates and objects. This method has been shown to be effective in other domains such as for multimodal embodied agents [22].
Behavior cloning for offline imitation learning. Behavior cloning is an imitation learning method that trains a seq2seq Transformer offline with action transcripts of similar training tasks generated by oracle agents [34, 3]. During training, it uses the previous action, observation at time step  $t - 1$ , and the current observation as input and learns to output the next action. The Text Decision Transformer (TDT) is a textual variant of the Decision Transformer [9], which also employs behavior cloning and uses the same data. The primary innovation of TDT is the introduction of reward-to-go as part of the inputs, enabling the model to learn predicting actions that maximize future expected rewards.
2.3 Prompting LLMs for Action Planning.
Language models (LLMs) such as GPT-4 have shown promise for action planning in interactive tasks [18, 15, 32, 38]. In this paper, we adapt three prominent methods to complex interactive reasoning tasks in ScienceWorld: SAYCAN [1], REACT [41], and REFLEXION [30].
SAYCAN [1] is a straightforward agent that integrates an LLM with a value function of underlying policies regarding grounding affordances (i.e., the feasibility of an action in the environment). We
need to provide the history and current environment as textual inputs to LLMs for generating a ranked list of action candidates. This action list is then reranked based on a value function.
REACT [41] presents a virtual 'think' action, enabling LLMs to generate subgoals during action planning. This approach requires human annotators to supply examples of correct subgoals for each task type, employing few-shot in-context learning to teach LLMs when and how to 'think' in order to plan subsequent subgoals, in addition to providing complete action trajectories.
REFLEXION [30], a recent work building on REACT, proposes a multi-round approach enabling LLMs to use the history of previously failed rounds to refine their planning for the next round. This self-reflection mechanism helps LLMs improve after each failed attempt. However, this may not be practical in real-world applications for many tasks, as actions in failed trials can be irrecoverable.
All three methods require a new LLM inference at each time step to predict the next immediate action, resulting in inefficient and costly agents. REACT and REFLEXION require human annotations of correct subgoals for each unseen task type. Moreover, it is difficult to generalize REFLEXION to real-world situations where trial-and-error approaches can be infeasible for embodied tasks.
2.4 Dual-Process Theory
The dual-process theory [39, 16] is a cognitive psychological framework proposing the existence of a fast and a slow thinking systems in the human brain. This influential theory has found widespread applications across various fields, highlighting the critical role of both systems in shaping human cognition [5, 8, 12, 20, 23]. By integrating the complementary strengths of both systems, agents can effectively and efficiently handle diverse challenges in real-world scenarios. Inspired by this, we aim to construct a generative agent that utilizes a small seq2seq LM as System 1 for associative reasoning via behavior cloning while developing System 2 for analytical reasoning by prompting LLMs.
3 SWIFTSAGE: A Generative Agent with Fast and Slow Thinking
In this section, we first establish the problem. Then, we present the two core modules, SWIFT and SAGE, individually. Lastly, we demonstrate the integration of these two modules. , resulting in a harmonious and effective interactive reasoning process.
Environment and tasks. We focus on complex interactive reasoning tasks situated in virtual textual environments such as ScienceWorld [36]. ScienceWorld provides an optimal setting for developing and evaluating agents in complex tasks, comprising 30 distinct task types covering 10 topics in science experiments. It features 10 locations, including an art studio, workshop, kitchen, living room, bedroom, bathroom, foundry, greenhouse, outdoor area, and a connecting hallway. The environment includes  $200+$  object types with multiple states (e.g., open, activated) and supports 25 action templates, resulting in an intractable search space. The simulator can generate numerous variations of each task type, providing a rich training ground. In each variation, the agent and environment initialization, such as the locations and states of objects, will differ. A plethora of training variations encompassing all task types are available for training agents. Additionally, it provides a handcrafted oracle agent to search for successful transcripts with minimal actions for offline learning.
Evaluation is done on a set of testing variations with unseen combinations of required objects and situations, thus substantially different from the training variations. For example, a training variation may involve boiling water, while a testing variation could require boiling tin. Therefore, it is crucial to ensure the agent's compositional generalization ability for effectively handling real-world scenarios.
Interactions. Given a task variation, an agent is provided with the task description  $D$  and the initial environment state  $(t = 0)$ . The task description  $D$  is a text specifying a high-level goal, e.g., "Your task is to test if an unknown substance  $A$  is electronically conductive." At each time step  $t$ , the agent generates an action  $A_{t}$  based on a set of supported action templates (e.g., pick up X, use X on Y).  $A_{0}$  is always "look around" for showing initial environment information. Upon receiving an action from the agent, the environment produces feedback in four dimensions:
Figure 2: An example of how SWIFTSAGE works with fast and slow thinking. The SWIFT module is offline trained via imitation learning with a small LM such as T5-large (770m). When it is necessary, for example, encountering an exception, we switch to the SAGE module that prompts LLMs (e.g., GPT-4) for planning and grounding the next subgoals, resulting in an action buffer.
- Observation  $O_{t}$  provides direct feedback on the action  $A_{t}$  regarding its effects on the environment or the information queried. For example, an  $A_{t}$  of "use thermometer on the substance in metal pot" may result in an  $O_{t}$  like "The temperature is 80F."
- Environment  $E_{t}$  represents the current room in which the agent is situated and provides details about all visible objects. Object visibility is based on container states, e.g., objects within a closed fridge are not included in  $E_{t}$  until the agent performs an action like "open fridge."
- Inventory  $I_{t}$  lists objects picked up by the agent, which is particularly useful when agents collect items from different locations to complete the task.
- Score  $S_{t}$  represents the agent's cumulative score ranging from 0 to 100. When a required intermediate state is achieved, the score increases with a positive reward.
3.2 SWIFT: The Module for Intuitive and Associative Thinking via Imitation Learning
Imitation learning is used to construct an agent that learns to mimic oracle agents in various training scenarios through seq2seq learning. Previous methods, such as TDT [36], mainly employ one-hop history as input context and learn to output the subsequent action  $A_{t}$  [36]. However, these methods exhibit limitations due to their restricted context of action history and harmful biases arising from data imbalance. To address these issues, we introduce our SWIFT module, depicted in Figure 2.
Representation for longer history. We expand the conventional one-hop BC to multi-hop by incorporating a sliding window of observations and rewards for the  $K = 10$  recent actions. Additionally, we include a special field for visited rooms (without duplication). This approach aims to provide agents with a longer context and prevent unnecessary room navigation. The input format is as follows: "Task:  $D$ ; Time:  $t - 1$ ; Score:  $S_{t - 1}$ ; Action history:  $[A_{t - i}(+R_{t - i})\to O_{t - i}] / ^* i$  loops from  $K$  to  $1^{*}$ /; Current room:  $E_{t - 1}$ ; Inventory:  $I_{t - 1}$ ; Visited rooms:  $\{E_1^*,\ldots ,E_{t - 1}^*\}$ ". Here,  $R_{t} = S_{t} - S_{t - 1}$  represents the reward at  $t$ , and  $E_{t}^{*}$  is the location name at  $t$ .
Balanced imitation learning. To avoid bias caused by data imbalance for seq2seq learning, we down-sampled specific types of tasks and actions to achieve a more balanced final dataset for training. We used the T5-large with 770 million parameter and instruction-following ability [10], creating an efficient agent that we named SWIFT. Our empirical results show that the SWIFT module performs much better than TDT (11 billion) despite being 15x smaller in size.
The SwIFT module exhibits greater accuracy during initial time steps, enabling it to attain higher scores in the early stages of a complex task. However, it often fails to generalize to unseen situations. The module also has a tendency to repeat meaningless actions when its learned plans yield exceptions from the environment (e.g., the broken stove in Figure 2). This is partly due to the nature of imitation learning, which prioritizes emulating the observable actions of oracle agents rather than their intrinsic planning abilities. Besides, since the oracle trajectories contain only the shortest, correct actions, it is thus also challenging for the SwIFT to learn how to fix mistaken actions.
3.3 SAGE: The Module for Deliberate and Analytical Thinking via Prompting LLMs
While the SWIFT module acquires surface knowledge about the environment and task types through imitation learning, it lacks two key abilities essential for complex interactive reasoning: 1) generalizable planning and tracking of subgoals, and 2) robust handling of exceptions. Prior research has shown that LLMs outperform smaller LMs in these abilities. They can perform step-by-step reasoning to devise concrete plans for tasks and self-refine their outcomes. However, the performance of prior methods remains unsatisfactory in complex interactive tasks such as those in ScienceWorld.
We introduce a novel two-stage approach, named SAGE. This method initially acquires higher-level recommendations from LLMs during the planning stage, followed by their translation into specific action sequences in the grounding stage. By decoupling the planning and grounding processes, SWIFTSAGE effectively generates a series of actions for completing the planned subgoals.
Planning stage. In this stage, we leverage LLMs to plan based on the current state. Specifically, we prompt LLMs with a single prompt that includes a summarized version of the task description and action history, and asks the following five key questions:
Q1 (locate objects): "To complete the task, which objects do I need to collect? Please list them and their possible locations one by one."
Q2 tracker objects): "Are there any objects that have not been collected yet?"
$\triangleright$  Q3 (plan subgoals): "To complete the task most efficiently, what are the important sub-goals to achieve? Please list the subgoals one by one."
$\triangleright$  Q4 (track progress): "Considering these subgoals, what have I already completed? And which subgoal should I focus on right now?"
Q5 (handle exceptions): "Have I made any mistakes that might prevent me from efficiently completing the next subgoal? If any, how should I fix them?"
Before posing the five planning-related questions, we condense the entire action history  $(A_{<t}$  and  $O_{<t})$ , and the current environment information  $E_{t-1}$ . Q1 and Q2 pertain to objects, as acquiring all necessary objects serves as the foundation for effective task planning. By addressing these questions, we ensure that LLMs develop a comprehensive understanding of the current environment. Q3 prompts LLMs to engage in step-by-step planning by decomposing the task into a series of subgoals. Q4 acts as a follow-up question, allowing the agent to monitor its progress based on the action history and determine completed subgoals, subsequently focusing on the remaining tasks. Lastly, Q5 is employed to identify and address potential exceptions. These questions can be further tailored with additional environment-specific hints, thereby enhancing their adaptability.
To improve the structure of the LLMs' outputs and facilitate parsing, we incorporate additional instructions in the prompt. By utilizing a single input to obtain answers to all five questions in one output, rather than engaging in multiple rounds of interactive prompting, our approach is more efficient and cost-effective than the iterative prompting methods.
Q4 and Q5 are of primary importance, while Q1-Q3 serve as auxiliary guidance for the LLMs. If the action history indicates a mistaken action or an unachievable previous subgoal, the response to Q5 refines the answer to Q4 through self-reflection on the fly. This approach differs from the REFLEXION agent, which only prompts reflective questions at the end of a failed trial, allowing agents to improve their planning in subsequent attempts. In contrast, our method detects exceptions and errors each time the agent plans for the next subgoals, enabling earlier correction of the agent's behavior.
Grounding stage. While the answers to Q1-Q5 provide valuable guidance for agents, they are not directly executable. Converting plans into valid actions that can be accepted by the environment
remains a challenge. Previous methods using LLMs over-generate candidates, and they rely on reranking or filtering based on the action space to select the next action. However, this is inefficient and inaccurate for complex tasks with vast action spaces. Additionally, these methods generate a single action at a time, which can be both costly and ineffective for long-horizon tasks.
To tackle these issues, we first present supported action types using a formal style accompanied by remarks. For instance, the action type "pour X into Y" is introduced as "POUR(X, Y): pour object X into container Y; e.g., pour red paint into wood cup". More examples: