Thought Cloning: Learning to Think while Acting by Imitating Human Thinking

TL;DR

This work introduces Thought Cloning, an imitation-learning framework where agents learn both how to think and how to act by training on synchronized human thought-action data. The authors implement a two-component architecture with a Thought Generator and an Action Generator, trained using a two-part loss that aligns thoughts with actions and with ground-truth demonstrations. In a synthetic BabyAI BossLevel domain, Thought Cloning outperforms Behavioral Cloning, shows stronger generalization to out-of-distribution tasks, and enables safety and interpretability features such as the Future Action Declaration Score and Precrime Intervention. The results suggest that scaling thinking-in-language data could dramatically improve AI capabilities and safety, with potential applicability to foundation models and internet-scale datasets. The work also discusses related planning-with-language approaches and emphasizes the value of thought-data alignment for planning, debugging, and steerability.

Abstract

Language is often considered a key aspect of human thinking, providing us with exceptional abilities to generalize, explore, plan, replan, and adapt to new situations. However, Reinforcement Learning (RL) agents are far from human-level performance in any of these abilities. We hypothesize one reason for such cognitive deficiencies is that they lack the benefits of thinking in language and that we can improve AI agents by training them to think like humans do. We introduce a novel Imitation Learning framework, Thought Cloning, where the idea is to not just clone the behaviors of human demonstrators, but also the thoughts humans have as they perform these behaviors. While we expect Thought Cloning to truly shine at scale on internet-sized datasets of humans thinking out loud while acting (e.g. online videos with transcripts), here we conduct experiments in a domain where the thinking and action data are synthetically generated. Results reveal that Thought Cloning learns much faster than Behavioral Cloning and its performance advantage grows the further out of distribution test tasks are, highlighting its ability to better handle novel situations. Thought Cloning also provides important benefits for AI Safety and Interpretability, and makes it easier to debug and improve AI. Because we can observe the agent's thoughts, we can (1) more easily diagnose why things are going wrong, making it easier to fix the problem, (2) steer the agent by correcting its thinking, or (3) prevent it from doing unsafe things it plans to do. Overall, by training agents how to think as well as behave, Thought Cloning creates safer, more powerful agents.

Figures (9)

• Figure 1: Overall framework for Thought Cloning (TC). The TC agent has two components: the Thought Generator and Action Generator. At each timestep, the TC agent receives an observation, a mission, and a history of thoughts as inputs. The Thought Generator generates thoughts, and the Action Generator generates actions conditioned on these thoughts. Generated thoughts and actions are compared to the ground truth from the demonstration dataset to calculate the loss.
• Figure 2: Left: A BabyAI babyai_iclr19 environment example. The environment contains various colored items (ball, key, box, door). The agent can pick up, drop, and move objects or open and close doors, while locked doors can only be unlocked with color-matched keys. The agent can observe the $7\times 7$ grid cells in front of it, which can be blocked by walls and closed doors. Right: An example from a trained Thought Cloning agent planning and replanning. The mission requires reaching the purple box (highlighted), but a purple ball blocks the way. The agent's thoughts and actions show replanning when encountering the obstacle, removing it, and resuming the previous goal.
• Figure 3: Training progress comparison of Thought Cloning (TC), Behavioral Cloning (BC), and a TC ablation variant without the Thought Cloning loss. The BC architecture is identical to the Action Generator of TC and the TC w/o Imitating Thought has the same architecture as TC, without the TC loss. BC and the ablation variant are trained solely with the action loss (which leads to some minor architectural differences, see Section \ref{['section:setup']}.) The error bars are the 95% confidence interval from five runs of experiments. The results indicate that TC learns faster than BC and also outperforms it. Furthermore, the comparison between TC and TC w/o Imitating Thought demonstrates that the superiority of TC is not simply due to having more parameters.
• 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 \ref{['section:generalization']}). 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.
• Figure 5: (a): A heatmap illustrating the Future Action Declaration Score, a metric designed to evaluate the interpretability of Thought Cloning agents (Section \ref{['section:safety_interpretability']}). 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 \ref{['section:generalization']}). 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.