SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shengran Hu<sup>1,2</sup>
srhu@cs.ubc.ca
Jeff Clune $^{1,2,3}$
jclune@gmail.com
$^{1}$ Department of Computer Science, University of British Columbia
$^{2}$ Vector Institute
3Canada CIFAR AI Chair
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. $^{1}$
Language may be the key to what separates humans from all other animals, endowing us with an amazing level of general intelligence [1-4]. Crucially, the benefits of language are not confined to improving our ability to communicate with others: language also helps us think better [2-4]. We first describe the benefits of agents that can understand language (a common topic in AI) before moving to the benefits of agents that think in language (a topic that has received far less attention).
There are many benefits that arise if our agents can understand language. Doing so is crucial for agents to generalize to new tasks we want them to perform. This is because it is drastically more sample efficient if one can tell an agent what the task is, rather than requiring the agent to figure out the task through trial and error [5, 6]. Moreover, agents that can understand language allow us to define new tasks at test time without having to anticipate every wish we might eventually have for
The source code, model weights, and dataset are available at https://github.com/ShengranHu/ Thought-Cloning.
our trained agents [7]. That is in contrast to conventional hand-designed task descriptions, which can be vast, but still place constraints on what we can ask an agent to perform [8].
While the benefits of agents that can understand language are commonly discussed, there has been relatively little discussion in AI, especially in Reinforcement Learning (RL), regarding the many benefits of agents that think in language. While it remains debated whether humans exclusively think in language [9, 10], many scholars believe that natural language is intricately woven into our thought processes, conferring distinct cognitive advantages: thinking in language helps humans generalize, extrapolate, adapt to new situations, combine old knowledge in new ways, explore, plan, replan when necessary or beneficial, transfer knowledge via analogies, and the list goes on [2-4]. Despite these benefits, AI agents rarely, if ever, think, at least not in human language. While neural networks have internal vector activations that can be considered thinking, many hypothesize that there are specific benefits to thinking in the discrete, symbolic form of language (e.g. combining ideas in an exponential number of ways) [6, 11, 12], meaning that agents that think in language might learn faster, perform better, and generalize better than non-lingual agents.
In addition to agents being more capable, there are major benefits regarding AI Safety and Interpretability that arise when agents think in our language. If one can watch an agent think during training, one can recognize deficiencies in skills or values that can be improved, or one could decide the agent is not ready to be deployed. During testing, one can constantly scan the thoughts of the agent and intervene when the agent plans to do something undesirable. For example, if an agent thinks "My goal is to take my passenger to the store as fast as possible so I will run through this red light without stopping" one could intervene to stop that behavior ahead of time. Furthermore, watching agents think enhances the steerability of agents. If an agent is confused when solving challenging tasks, one can inject their thoughts into the agent to help it solve the task in a desired way. A final major benefit of agents that think in human language is it makes it easier to train more capable, safer AI agents. One can spot why things are not working, instead of just seeing that they are not working, and that provides ideas for how to debug and or improve AI training.
For all these reasons, adding the ability of AI agents to think in language could produce many significant advantages, and we suggest that the most effective way to achieve this goal is by imitating human thinking. Humans do not acquire thinking skills in isolation; instead, they are learned in part through demonstrations and feedback provided by teachers [2, 13-15]. As such, a promising method is to have agents learn from demonstrations where humans think out loud while acting. This approach is distinct from existing works that leverage pre-trained Large Language Models (LLMs) for planning [16, 17], because such LLMs are not trained on data where humans think out loud while acting. Thought data, such as YouTube videos and transcripts [18, 19], contains millions of hours of people talking out loud while performing tasks, revealing the thinking behind their actions, planning, decisions, and replanning, such as when they play video games [19]. This thought data is greatly valuable and widely available (Section 2), but has not yet been extensively explored, and this work hopes to encourage further research into the use of thought data to teach thinking skills to agents.
Provided we can solve the real, significant challenges of AI Safety and existential risk [20-24], there are tremendous gains to be had by creating more powerful AI or even AGI. In this paper, we propose a novel Imitation Learning framework, Thought Cloning, where agents not only learn to act from human demonstrations, as in Behavioral Cloning [25], but also learn to think from demonstrations where human think out loud while acting. Although we expect Thought Cloning to truly shine when trained on vast online datasets of synchronized human thoughts and actions, this paper validates the concept with synthetic thought data in a challenging domain, BabyAI [26]. Our experimental results illustrate that Thought Cloning outperforms Behavioral Cloning, even when Behavioral Cloning agents have the ability to think (in latent vectors), but have to learn that skill without the supervision of thinking provided by Thought Cloning. We also demonstrate that Thought Cloning generalizes better than Behavioral Cloning in out-of-distribution tasks in both zero-shot and fine-tuning settings. Finally, we provide empirical evidence for the previously discussed advantages of Thought Cloning in terms of Safety and Interpretability, where unsafe behavior can be near perfectly stopped before execution. All told, the results are promising and offer a glimpse of the enormous potential of Thought Cloning to not only make AI smarter, but also safer and more interpretable.
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.
Conventional Imitation Learning methods [27, 28], such as Behavioral Cloning [25], strive to construct a policy that accurately replicates the distribution of behavior in a given dataset of demonstrations. However, our proposed framework, Thought Cloning, diverges from this approach by aiming to teach agents how to also think while acting, utilizing a synchronized dataset of human thinking. The thought dataset, denoted as  $\mathcal{D} = \{D_i\}_{i=1}^N$ , comprises a series of trajectories,  $D_i = (m, \{(o_t, th_t, a_t)\}_{t=1}^T)$ . Each trajectory encompasses a mission,  $m$ , defined in natural language, along with an observation  $o_t$ , an action  $a_t$ , and a corresponding thought  $th_t$  at each timestep,  $t$ . Such datasets are widely available online. For example, by inferring action labels from Youtube videos with VPT [19] and then retrieving the corresponding transcripts, we can obtain a thought dataset that contains both human thinking and action [19, 18]. In such a dataset for Minecraft, a thought like "I need to gather wood to build a shelter before nightfall" might correspond to the player moving towards a tree and collecting wood. It is worth noting that noise is an inevitable aspect of online data. For instance, commentary like "please subscribe to my channel" would be commonly found in YouTube videos. However, we believe this challenge could be effectively mitigated with proper data preprocessing. An example from MineCLIP [18] shows that careful heuristic-based preprocessing of the captions can make data clean enough to train multi-modal models (e.g., a simple heuristic focusing on domain-specific vocabulary could remove most off-topic text). More promisingly, one could use language models to filter out off-topic data, either by removing any off-topic comments or determining that some videos should be excluded because not enough of the commentary is helpful. Additionally, even if noise is present, recent ML history shows that-at scale-such "noise" does not prevent learning the "signal": examples of this occurring on internet-scale data include GPT [29], CLIP [30], and VPT [19]. While not a perfect simulation of online caption noise, as described later, we add noise to the synthetic thought data and Thought Cloning works well despite it (Section 3.1).
In the Thought Cloning training framework, agents learn to produce natural language thoughts at each timestep and subsequently condition their actions based on these generated thoughts. This learning process gives rise to a bi-level architecture (Fig. 1). The architecture comprises an Thought Generator responsible for thought generation, and a Action Generator tasked with executing actions based on the thoughts generated by the Thought Generator. While different choices of what to condition the Thought Generator and Action Generator are possible, in this work, for a particular trajectory of length  $T$  in the thought dataset we minimize:
$$
\min  _ {\theta_ {u}, \theta_ {l}} \sum_ {t = 1} ^ {T} - \alpha \log \pi_ {\theta_ {u}} \left(t h _ {t} | m, \left\{o _ {\tau} \right\} _ {\tau = 1} ^ {t}, \left\{t h _ {\tau} \right\} _ {\tau = 1} ^ {t - 1}\right) - \log \pi_ {\theta_ {l}} \left(a _ {t} | m, \left\{o _ {\tau} \right\} _ {\tau = 1} ^ {t}, t h _ {t}\right) \tag {1}
$$
Here,  $\theta_{u}$  and  $\theta_{l}$  represent the weights for the Thought Generator and Action Generator;  $\alpha$  represents the coefficient for Thought Cloning loss;  $th$ ,  $o$ ,  $a$ , and  $m$  denote thought, observation, action, and mission, as previously described. The first part of the loss is the Thought Cloning loss, where the Thought Generator is conditioned on the history of thoughts, observations, and the mission, to predict the thought. That thought is then compared with the ground truth thought in the dataset. The second part is the action loss, where the Action Generator is conditioned on the current thought, the history
Figure 2: Left: A BabyAI [26] 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.
of observations, and the mission and predicts the action to do, and we then calculate the loss by comparing the predicted action to the ground truth action in the dataset.
For more complex or large-scale scenarios, the Thought Generator can be implemented with pre-trained Vision-Language Models (VLM) either zero-shot or fine-tuned [31], while the Action Generator can be trained from scratch or adapted from existing language-conditioned controllers in the target domain [6, 16]. In this paper, we base both components on the BabyAI 1.1 model architecture [32], which utilizes a memory-augmented architecture—an LSTM [33]–to address the partial observability challenge. The model also employs FiLM [34] for modality fusion, effectively combining visual and text input. The detailed architecture adopted in this paper can be found in Supplementary Material A. While all models in this paper are trained from scratch, we anticipate that the utilization of pre-trained models in complex domains will be beneficial.
3.1 Domain and Synthetic Thought Data
This paper employs BabyAI [26], a simulated partially observable 2D gridworld domain. We focus on the most challenging environment, BossLevel, in BabyAI. An overview of the domain is shown in Fig. 2 (left). Each BabyAI environment consists of a randomly generated room layout, item configuration, and a mission described in natural language, sampled on an environment distribution. Colored items (balls, keys, boxes, doors) and the initial position of the agent are randomly distributed across a  $20 \times 20$  grid world containing nine  $6 \times 6$  rooms. Missions comprise four possible tasks (GoTo, PickUp, OpenDoor, PutNextTo), connected by then/after and and (with or without ordering constraints). GoTo and PickUp require agents to go to or pick up an object; OpenDoor requires agents to open or unlock a door; PutNextTo requires the agent to pick up object A, find object B, and drop A next to B. The mission may implicitly require the agent to open or unlock doors to find the target objects. Relative directional instruction in the mission, e.g., on your right, is based on the agent's initial position. An environment is solved when all tasks in the mission are completed. The agent's observation consists of the  $7 \times 7$  grid cells in front of the agent, except the agent cannot see through walls (Fig. 2 yellow square). This work features the state-based observations provided by BabyAI [26]. Each grid cell in the  $7 \times 7$  observation is represented by three integers: [the item ID, the color ID, and a status code], resulting in a  $7 \times 7 \times 3$  observation matrix. The status code is 1 for closed doors and 2 for locked doors, with 0 for open doors and other items. Occluded grid cells are assigned an item ID of 0. The agent's action space includes [left, right, forward, pickup, drop, toggle door (unlock, open, close)].
The key challenges in BabyAI revolve around partial observability, hard-to-explore mazes, complex missions in natural language, and long-horizon planning. The  $7 \times 7$  observation field is limited
compared to the  $27 \times 27$  maze, and the agent cannot see through walls and closed doors. The maze containing multiple closed rooms is difficult to navigate and explore as the agent needs to find target items across multiple closed (even locked) rooms. The missions are challenging because (1) they are described in natural language and (2) they can consist of multiple tasks, each requiring complicated navigation and actions. Combining all these factors results in a long horizon, with hundreds or even thousands of steps needed to solve a single environment.
One significant advantage of BabyAI is that it provides an Oracle Solver (named BOT in [26]) capable of generating step-by-step solutions for any given environment. This is achieved through hand-coded rules and an internal stack machine to generate plans for solving environments. In our work, we translate the Oracle Solver's internal states into natural language thoughts with pre-defined rules. For example, if the inner logic is to open a red door to explore the room, the translated thought will read, "open red door to explore". This translation process is combined with the generated demonstrations to synthesize the thought dataset with 1 million trajectories. To make the dataset more realistic, noise is added, with a  $1\%$  chance of adding a random noisy segment at each timestep, consisting of a random thought and several random actions, with a random length sampled from 1-6. A trajectory with example noise is shown in Supplementary Material C.
3.2 Experiment Setup
To verify the effectiveness of learning to think, we compare our Thought Cloning (TC) approach to the classic learning algorithm, Behavioral Cloning (BC). BC shares most of its architecture with the Action Generator of TC (Fig. 1), and because it is trained only on action loss, it does not encode thought like the Action Generator of TC. Additionally, since BC has fewer parameters than TC, we introduce an ablation variant called TC w/o Imitating Thought that is trained without the Thought Cloning loss to demonstrate that TC's superiority is not solely due to its larger number of parameters. This variant's architecture is identical to the TC architecture, except for a minor architectural difference where the latent vector from the Thought Generator is directly input to the Thought Generator as thought. This adjustment is necessary because this variant is not trained on the Thought Cloning loss, so we do not have per-word supervision. To train these parameters, we thus need to train them based on how these "thoughts" contribute to the action loss (i.e. they are trained end-to-end to predict actions). If we sampled words (as in Thought Cloning), we could not train these parameters end-to-end because hard sampling of words is non-differentiable, so gradients could not flow backward through this operation. Thus, we make one small change in order to allow the parameters to be trained, which is to pass the logits of the Thought Generator directly into the Action Generator, which is a differentiable operation. We feel this is the closest and fairest control possible to Thought Cloning, allowing virtually the same architecture and the same number of parameters, but not including the Thought Cloning innovation of exploiting human thought data.
Our training setup is based on BabyAI [26, 32]. The training iterates for 8 epochs on the 1 million episode dataset, corresponding to a total of  $7 \times 10^{8}$  training frames. The Thought Cloning loss parameter  $\alpha$  (Eq. 1) is set to 2. During training, we employ teacher-forcing [35], which is adopted when decoding thoughts. It conditions the Action Generator on the ground truth thoughts from the dataset. The teacher-forcing ratio linearly decreases from  $100\%$  to  $0\%$  during the training process. Producing all the main results in the paper took about ten A40 GPUs for one week. More details on training can be found in Supplementary Material A.
In our experiments, the performance of agents is evaluated based on their success rate in held-out test environments. Success for an environment is defined as the completion of all specified tasks in the mission. By controlling random seeds, all test environments are unseen during the training process. All experiment results from Sections 3.3, 3.4 and 3.5 are calculated from five independent runs. The success rate results in Section 3.3 are obtained by testing agents on a set of 512 sampled environments. In Section 3.4, agents are tested in a larger set of 1,024 test environments. During the testing phase, the TC agent has identical observations as the BC agent, i.e. it has no extra information.
3.3 Imitation Learning
In this section, we show the main performance results of training TC, BC, and TC w/o Imitating Thought. The results illustrate that TC learns faster than BC, where BC requires orders of magnitude more time to achieve a performance similar to TC's early-stage results, and TC ultimately outperforms BC at the end of training (Fig. 3). The outperformance of TC compared to BC at  $25\%$ ,  $50\%$ ,  $75\%$ , and
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 3.2.) 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.
$100\%$  of the way through training is statistically significant, as confirmed by the Mann-Whitney U test, with  $p = [0.012, 0.008, 0.021, 0.008] < 0.05$ . These results support our hypothesis that natural language can help the agent learn to explore and plan.
Another comparison is between TC and an ablation variant TC w/o Imitating Thought that shares the same architecture with TC, but without the Thought Cloning loss in training. The results show that TC also substantially outperforms TC w/o Imitating Thought (Fig. 3). Similar to the previous, the results are statistically significant ( $p = [0.008, 0.012, 0.008, 0.008] < 0.05$ ). The results reveal that TC's superior performance is not solely due to a larger number of parameters than BC, and also supports our argument that learning from human thought boosts an agent's ability to think.
An example of a TC agent planning and replanning is shown in Fig 1 (right). After opening the blue door, the agent discovers the target (a purple box) within its observation and thinks about going to it to complete the task. However, the agent realizes that a purple ball is blocking its path. A smart replan emerges here, with the agent inserting a new plan to remove the ball. The agent achieves this subgoal by picking up the ball in its way, finding an empty space, and then dropping it. After completing this new, necessary, intermediate task, the agent resumes its original mission to go to the purple box and successfully solves the environment. From this example, we can see that by thinking like humans in natural language, the agent demonstrates successful planning and replanning abilities. We also see the interpretability benefits, as it is easy to follow along and understand why the agent executes certain actions.
3.4 Generalization to Out-of-Distribution Environments