SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This section compares the generalization abilities of the TC and BC agents by testing them on environments that are increasingly out of distribution. We define the distribution of environments with two difficulty dimensions: Behavioral Difficulty and Cognitive Difficulty. Behavioral Difficulty is based on the length of the action sequence required to solve the environment (provided by Oracle Solver, see Section 3.1). The simplest environments require about 20 steps, while the most challenging environments require more than 500 steps. Cognitive Difficulty reflects the complexity of the mission, with more difficult environments requiring stronger planning abilities to complete complex tasks. The definition of Cognitive Difficulty is adapted from the maxStep parameter in BabyAI environments [26] and is given by (# of {GoTo, PickUp, OpenDoor} + 2 × # of {PutNextTo} + # of ordering constraints). The PutNextTo task is assigned a higher weight because it involves a combination of picking up, navigating, and dropping, making it the most challenging task among the four. The range of cognitive difficulty spans from 1 (simplest) to 9 (most difficult). In the training distribution, the environments exhibit means and standard deviations of Behavioral and Cognitive Difficulties of  $84.2 \pm 68.8$  and  $2.7 \pm 1.6$ , respectively. In this paper, we define out-of-distribution (OOD) environments as those with a Behavioral Difficulty > 175 or a Cognitive Difficulty  $\geq 4$ , each being approximately more than one standard deviation away from the mean. The furthest OOD environments, with a Behavioral Difficulty greater than 425 or a Cognitive Difficulty of 9, had
(b) Success rates after fine-tuning agents on out-of-distribution environments
(a) Zero-shot success rates of agents on environments that are increasingly out of distribution.
Figure 4: The zero-shot and fine-tuning success rate of Thought Cloning (TC) and Behavioral Cloning (BC) agents on environments that are increasingly out of distribution. Behavioral and Cognitive Difficulties are defined by the length of the solutions to environments and the mission complexity of environments respectively (Section 3.4). The error bars in bar and line plots are the  $95\%$  confidence interval from five runs of experiments. (a): The gray region indicates the training distribution. The Oracle Thought + TC Learned Control refers to the TC agent with oracle high-level thoughts. The results demonstrate TC generalizes much better than BC. They also illustrate that with a more powerful Thought Generator trained from vast human thought data, the agent should become drastically more capable. (b): TC is much better at adapting to novel situations than BC.
less than  $5.7 \times 10^{-5}$  and  $1.6 \times 10^{-4}$  probability of being sampled during training (calculated with rejection sampling). For testing both in-distribution and out-of-distribution environments, we sample various sets of environments that extend away from the distribution in terms of both Behavioral and Cognitive Difficulty, and then evaluate agents on these sets. For Cognitive Difficulty, we sample sets of environments across the full range of Cognitive Difficulty levels 1-9. For Behavioral Difficulty, we sample sets of environments within intervals of 50 (e.g., 125-175, 175-225, etc.), starting from 25. Environments with a Behavioral Difficulty  $>425$  are grouped into one set.
First, we test the zero-shot performance of TC and BC agents in OOD environments. The results show that the TC agent substantially outperforms the BC agent with environments being increasingly out of distribution (Fig. 4a), and the results are statistically significant across all testing difficulties (Mann-Whitney U test  $p < 0.05$ ), thereby supporting our hypothesis that language utilization can enhance agents' generalization capabilities. Moreover, we observe that the Oracle Thoughts + TC Learned Control achieves near-optimal performance even on the most challenging environments. This indicates that the current limitation of TC performance lies in high-level thinking. As we scale our approach to leverage internet-sized datasets of human thinking, the high-level thinking capability is expected to improve substantially, thereby enhancing the power of the TC agent.
Next, we investigate how well the agents adapt to new situations by fine-tuning them on OOD environments. We fine-tune the TC and BC agents on the corresponding environments for 15 epochs, with the same settings described in Section 3.2. The results demonstrate that the TC agent is better at adapting to OOD environments (Fig. 4b). The superiority of TC over BC is statistically significant across all testing difficulties, as supported by the Mann-Whitney U test  $p < 0.05$ , with the exception of Cognitive Difficulty 4, where both methods already achieve near-perfect performance. The results support our argument that language can better assist agents in adapting to novel situations.
3.5 AI Safety and Interpretability
The ability to observe the agent's thought process gives our model a high degree of interpretability. To empirically assess the interpretability of TC, we introduce a metric named the Future Action
(a) Interpretability of Thought Cloning agents
(b) Effectiveness of the Precrime Intervention
Figure 5: (a): A heatmap illustrating the Future Action Declaration Score, a metric designed to evaluate the interpretability of Thought Cloning agents (Section 3.5). The  $x$  and  $y$  axes denote various levels of difficulty. Each cell represents a region of sampled environments, with the color intensity reflecting the mean score. Brighter cells indicate a higher degree of match between the agent's declared thoughts and subsequent actions. The green square denotes the training distribution, while the rest of the regions are out of distribution (Section 3.4). The results illustrate the robust and consistent interpretability of Thought Cloning agents. (b): A bar chart demonstrating the effectiveness of the Precrime Intervention mechanism, which is to halt the Thought Cloning agents upon detecting dangerous plans in their thoughts and thus prevent unsafe behaviors. We show three tests ( $x$  axis) where (1) touching red items, (2) picking up balls, and (3) picking up requested items were declared unsafe. We report the fraction of episodes where unsafe behaviors occurred ( $y$  axis). The error bars are the  $95\%$  confidence interval from five runs of experiments. The results show that Precrime Intervention effectively eliminates almost all unsafe behaviors.
Declaration Score. This metric quantifies the fraction of times when an agent, preparing to execute an action other than navigation, declares this impending action in its thoughts beforehand. In the training distribution, TC agents performed exceptionally well (green square in Fig. 5a). Interestingly, TC agents also scored near-perfectly across all out-of-distribution environments (rest of Fig. 5a), demonstrating the robust and consistent interpretability of our model even under novel, out-of-distribution situations, which is an important property for AI safety and interpretability.
Due to its high degree of interpretability, Thought Cloning allows a simple method that can considerably enhance AI safety. We call it Precrime Intervention. In practical settings, an agent might employ dangerous or undesirable strategies to accomplish challenging tasks. However, because Thought Cloning features such strong interpretability, we can simply halt the agent upon detecting dangerous thoughts and thereby prevent the unsafe behavior it was planning to conduct. Additionally and importantly, Precrime Intervention does not require any changes to the weights of the model. If we learn or decide after training that a certain behavior is unsafe or undesirable, Precrime Intervention can still prevent it. The same flexibility allows different definitions of what is allowable and unsafe behavior in different settings (e.g. in the presence of adults vs. children or customized to the preferences of different countries with different regulations). To demonstrate this flexibility and test to what extent Precrime Intervention works, we conducted three separate tests, where we declared three different behaviors as unsafe (1) touching any red item, (2) picking up any ball, and (3) picking up the object the agent is being asked to pick up in its mission. The last one is particularly interesting because the agent has a strong prior to want to perform that action, which Precrime Intervention has to combat. We report the fraction of episodes where such unsafe behaviors occurred with and without Precrime Intervention (Fig. 5b). Remarkably, Precrime Intervention almost entirely eliminates all unsafe behaviors, thereby demonstrating the promising potential of TC agents in advancing AI safety.
Moreover, the interpretability of the model also greatly aids in diagnosing problems, thus simplifying the development of more capable and safer AI. This feature actually proved beneficial during the development of this paper. Initially in our development, the TC agent showed promising performance in training, but frequently failed during testing, repetitively oscillating between incorrect thoughts (plans) without actively exploring new ideas. This observation helped us to recognize that, because we had trained with teacher forcing throughout with oracle (i.e. perfect) thoughts, the agent had never practiced having incorrect thoughts, and thus had never practiced recovering from them by trying alternate ideas. Thus, at inference time when it had to generate its own thoughts, which are sometimes incorrect, it did not know how to recover. We then instead tested an immediate transition from teacher-forcing to  $100\%$  auto-regressive sampling and training (i.e. from  $100\%$  teacher-forcing on
one training step to  $0\%$  on the next), but the agent generated too many nonsensical thoughts. Thanks to the model's interpretability, we were able to recognize the situation and try an alternative strategy that worked well, and is the one we report results for in this paper (Section 3.3): we gradually decay the teacher-forcing rate (fraction) during training, which dramatically improved performance. Supplementary Material D contains more details about this example.
Lastly, TC enables steerability, including helping agents when they are stuck. That is because the actions of TC agents are conditioned on their thoughts, and we can manually inject alternate thoughts to have the agents do what we wish. The TC agent, when provided with Oracle high-level thoughts, is capable of near-perfect performance across almost all environments (Fig. 4a), which provides evidence that steering agents is possible and effective.
4.1 Planning in RL with Language
Recent work leverages the reasoning capability, the ability to flexibly combine abstractions, and the interpretability offered by natural language to address high-level planning challenges in real-world domains. We augment this approach by enabling agents to think in language, facilitating the capability of agents, AI Safety, and Interpretability. There are two major categories of works in the literature that enable language planning. The first involves Hierarchical RL methods, where the language represents the hierarchy [36-40]. However, the planning space in these works is constrained to a pre-defined subgoal set, limiting their generalization to novel scenarios and preventing them from utilizing the reasoning and powerful commonsense found in pre-trained LLMs [29, 41, 42]. The second category of work involves pre-trained LLMs that generate plans in language for RL systems. Earlier works [16, 17] allow the LLM to predict step-by-step plans for a specific task. However, these works are open-loop methods, as the LLM cannot perceive the environment while acting, and thus cannot adapt and change once things do not go according to plan, which is a crucial capability in complex environments. Some recent approaches have developed closed-loop methods to provide LLMs with dynamic information for planning [17, 43, 44]. While these works show exciting performance in different environments, their closed-loop feedback mechanisms for the LLMs either rely on an oracle from the environment or complicated captioning models. The work most relevant to our vision is PALM-E [31], in which a pre-trained Vision-Language Model is adopted as the planner, allowing it to recognize patterns from observations directly. However, the works mentioned above were not trained on data of humans thinking and acting, meaning they do not benefit from learning from human thought demonstrations how to do things like plan, replan, create high-level goals and the subgoals required to achieve them, and the many other benefits of thinking intelligently during acting.
Apart from embodied domains, many studies utilize LLMs as agents [45] operating in text domains, such as coding [46, 47] or using API tools [48]. Many LLM agents benefit from the reasoning and planning capabilities of LLMs [48, 49]. Some even allow LLMs to replan based on new information [46, 47]. However, none of the methods described above directly learn to think while acting by imitating human thought data. Thus, unlike Thought Cloning, they do not benefit from learning from human thought demonstrations. Within LLM agent research, ReAct [50] is the most relevant to our approach. It prompts LLMs to first generate reasoning and then conditions actions based on such reasoning text. The paper also includes an experiment that fine-tunes LLMs with demonstrations that contain both reasoning and actions. However, obtaining such demonstrations is challenging, as highlighted in the ReAct paper [50], and the paper does not provide insights on overcoming this data bottleneck to scale the method.
4.2 Learning from Dataset Aligning Action and Language
Several studies have recognized the value of datasets that align action with language. DIAL [51] employs such a dataset to train language-conditioned agents with Behavioral Cloning, achieving impressive results in real-world robotic tasks. However, it is limited by a pre-defined instruction set. Another work,  $(\mathrm{SL})^3$  [52], generates a hierarchical dataset for agents to learn from, demonstrating superiority in a challenging 3D simulation domain, but has the drawback discussed in the previous section of being open-loop. Finally, in the study most similar to our own, Hu et al. [53] collected a dataset from two human players collaborating on an RTS game. However, the agent in [53] is not language conditioned, which limits its potential to learn to do any task (e.g. arbitrary tasks
requested of it in natural language). Similarly, a work concurrent to ours constructed a dataset with BabyAI oracle plans [54]. However, their architecture, unlike ours, is not compatible with most pre-trained models, making ours more able to harness new, powerful foundation models to tackle increasingly complex challenges. Additionally, although previously mentioned two methods [53, 54] employ learning frameworks similar to ours, they do not explore the full potential of learning from datasets that align action with language, particularly in terms of the resulting benefits in terms of generalization, AI Safety, and Interpretability.
Our research findings are focused on two main areas. First, our Thought Cloning (TC) agent demonstrated superior performance compared to Behavioral Cloning (BC), effectively showcasing its capabilities in generalization, exploration, planning, replanning, and adaptation to various situations. Second, we presented empirical evidence underscoring the benefits Thought Cloning provides in AI Safety and Interpretability. The robust interpretability of the TC agent not only help developers in diagnosing AI systems, but also contributes to AI safety, as evidenced by mechanisms such as Precrime Intervention. Our empirical results on steerability further spotlight the potential of TC agents in effectively collaborating with humans to tackle complex tasks.
We utilized a synthetic dataset and trained a model from scratch as a proof of concept. However, the full vision for the TC framework will be when Thought Cloning agents are trained on internet-scale datasets of humans thinking out loud while acting, such as YouTube videos and their narration [19, 18] (whether in closed caption transcripts or directly from audio). Consider the prospect of an agent that has both learned to think and act like humans in a huge variety of settings. Much like the thoughts of human children are guided by teachers, our agents could become skilled at planning, replanning, reasoning, and explaining their thinking to us (either via their outputs or because we have the unique ability to observe the thoughts in their minds). The vision for utilizing internet-scale datasets is also supported by the empirical results in Section 3.4, which suggest that the current bottleneck in agent capability is its high-level thinking, a skill that could be enhanced by scaling to vast online data. Additionally, we trained all models from scratch in this paper. In more complex domains, we anticipate that the utilization of pre-trained Transformer-based models could be beneficial.
Moreover, Thought Cloning can also improve Foundation Models (FMs) by enabling a separate "thought channel" where they can output thoughts that get fed back in when they are planning and answering. Recent studies indicate that FMs can benefit from contexts they self-generate, such as reasoning steps [11, 50]. Thought Cloning can amplify this advantage by training FMs with Thought Cloning data which includes not only the answers or solutions (action data) but also the reasoning behind them (thought data). For example, a human programmer needs the ability to think when developing software, but in a way where those thoughts are not part of the final output. However, current LLMs are not trained this way. In addition to performance gains, adding such a thought channel to FMs allows all of the AI Safety and Interpretability advantages of Thought Cloning, including Precrime Intervention, to be added to the world-changing FM technology.
Of course, there are also risks associated with such agents. As occurs with LLM pretraining and other forms of Behavioral Cloning, Thought Cloning could inadvertently inherit undesirable human flaws, such as speaking falsehoods, providing false yet persuasive rationalizations, or being biased. Alignment techniques are being constantly improved to address these challenges [55]. However, even improving AI agent safety up to the level of a (flawed) human would be a major advance, even if the resulting system is not perfect. Additionally, a distinguishing feature of Thought Cloning is it provides the ability to interpret and prevent these flaws from culminating into actions, making TC a more favorable approach in this regard. (See the experimental analyses on the effectiveness of Precrime Intervention Section 3.5). Our intent was to inspire researchers to advance this method and or test it in real-world scenarios, but the method itself should not be considered as a comprehensive safety solution for real-world systems.
In conclusion, this paper has introduced Thought Cloning, where agents not only simply learn to act from human demonstrations, as in Behavioral Cloning, but also learn to think from demonstrations where humans think out loud while acting. Through Thought Cloning, we have illustrated how an agent can become more capable, interpretable, and safe by imitating human thinking. This work facilitates the training of increasingly powerful agents and opens up numerous avenues for future scientific investigation in Artificial General Intelligence, AI Safety, and Interpretability.
This work was supported by the Vector Institute, the Canada CIFAR AI Chairs program, grants from Schmidt Futures and Open Philanthropy, an NSERC Discovery Grant, and a generous donation from Rafael Cosman. We thank Aaron Dharna, Ben Norman, and Jenny Zhang (sorted alphabetically) in our lab at the University of British Columbia for insightful discussions, and Michiel van de Panne for helpful commentary.
[1] Robert C Berwick and Noam Chomsky. Why Only Us: Language and Evolution. MIT press, 2016.
[2] David Premack. Is language the key to human intelligence? Science, 303(5656):318-320, 2004.
[3] Noam Chomsky et al. Language and Mind. Cambridge University Press, 2006.
[4] Steven Pinker. The Language Instinct: How the Mind Creates Language. Penguin UK, 2003.
[5] Mohit Shridhar, Lucas Manuelli, and Dieter Fox. CIoR: What and where pathways for robotic manipulation. In Conference on Robot Learning, 2021.
[6] Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, Pierre Sermanet, Noah Brown, Tomas Jackson, Linda Luu, Sergey Levine, Karol Hausman, and Brian Ichter. Inner Monologue: Embodied reasoning through planning with language models. In arXiv preprint arXiv:2207.05608, 2022.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
[8] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, et al. Open-ended learning leads to generally capable agents. arXiv preprint arXiv:2107.12808, 2021.
[9] Ben Alderson-Day and Charles Fernyhough. Inner speech: development, cognitive functions, phenomenology, and neurobiology. *Psychological bulletin*, 141(5):931, 2015.
[10] Romain Grandchamp, Lucile Rapin, Marcela Perrone-Bertolotti, Cédric Pichat, Célise Haldin, Emilie Cousin, Jean-Philippe Lachaux, Marion Dohen, Pascal Perrier, Maëva Garnier, et al. The condialint model: Condensation, dialogality, and intentionality dimensions of inner speech within a hierarchical predictive control framework. Frontiers in Psychology, 10:2019, 2019.
[11] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed H Chi, Quoc V Le, Denny Zhou, et al. Chain-of-Thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
[12] Gary Marcus. The algebraic mind, 2001.
[13] Lev Vygotsky et al. Interaction Between Learning and Development. Linköpings universitet, 2011.
[14] John D Bransford, Ann L Brown, Rodney R Cocking, et al. How People Learn, volume 11. Washington, DC: National academy press, 2000.
[15] Steve Olusegun Bada and Steve Olusegun. Constructivism Learning Theory: A paradigm for teaching and learning. Journal of Research & Method in Education, 5(6):66-70, 2015.
[16] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang,
A Architecture and Training Details