SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For full transparency, replicability, and to facilitate future research building on our work, we are releasing the source code, model weights, and the dataset used in our experiments<sup>2</sup>. Additionally, we provide key details necessary for the evaluation and replication of our work in this supplementary information. The architectural details of the Thought Cloning models adopted in this paper are shown in Fig. 6. As in [32], all missions and thoughts are encoded with Gated Linear Units (GLUs), with separate encoders employed for the missions and thoughts respectively. After the encoding process, we apply an attention mechanism [56] to dynamically weight the importance of different parts of the text encoding, based on the state history. The observation is encoded with a Convolutional Neural Network (CNN) and a Bag-of-Words [57] encoding approach. In the Thought Generator, a Transformer encoder [56, 58] is adopted to embed the thought history and mission, with the thought history as the query and the mission as the key and value. This Transformer encoder consists of two layers, each with two heads. The Action Generator is identical to the Behavior Cloning Baseline, except with the additional encoding of thoughts. Key architectural parameters, such as memory size and embedding size, are consistent with the baseline in [32], as shown in Table 1.
The pseudocode for Thought Cloning (TC) training framework is shown in Algorithm 1. In the loss function, we follow [32] by including an entropy term for actions. The Adam optimizer [59] is adopted to train TC and TC variant, with a batch size of 180 and a learning rate of  $5e^{-4}$ . Similar to the setting in baseline [32, 26], we train BC with a batch size of 296 and a learning rate  $5e^{-5}$ . The learning rate schedule begins with a warm-up phase of  $5T$  training steps, where  $T = 51200$ , linearly increasing from  $1e^{-4}$  to  $5e^{-4}$  for every  $T$  step, and then decaying by  $50\%$  at  $120T$  training steps, similar to the practices in [26, 60]. The teacher-forcing ratio linearly decreases from  $100\%$  to  $0\%$  from the  $10T$  training step to the end for every  $T$  step. In line 5 of Algorithm 1, the input thought could be the ground truth from the dataset  $(th_t)$  or the generated thought from the Thought Generator  $(\hat{th}_t)$ , depending on with or without teacher forcing. For training efficiency, Backpropagation Through Time was truncated at 20 steps in TC. The mix precision in PyTorch is also adopted during training, which speeds up training without sacrificing much performance [61]. In fine-tuning experiments, due to the increased difficulty of the levels and longer steps requiring more memory, we reduced the batch size from 180 to 40 and trained with an auto-regressive strategy. Detailed hyperparameter settings are shown in Table 1.
Table 1: Hyperparameter Settings
Hyperparameter	Value
Adam β1	0.9
Adam β2	0.99
Adam ε	10-5
Entropy Coefficient	0.01
Image Embedding Dimension	128
Text Embedding Dimension	256
Memory Dimension	2048
B Ablation Study: BC with the same parameter count or data as TC
We designed the control TC w/o Imitating Thought (Section 3.2) to address the concern that TC outperforms BC because TC has more parameters. Although TC w/o Imitating Thought allows the control to have the same number of parameters (and architecture) as Thought Cloning, for completeness, we try to address the concern in another way. Instead of holding the architecture the same, we create a BC control with a more canonical BC architecture, but with the same number of
$^{2}$ https://github.com/ShengranHu/Thought-Cloning
Figure 6: Detailed architecture for Thought Cloning (TC) agent adopted in this paper. At each timestep  $t$ , the inputs to the TC agent include a natural language-defined mission  $m$ , and an observation  $o_t$ . All preceding thoughts  $\{th_{\tau}\}_{\tau=1}^{t-1}$  is embedded with an LSTM in the Thought Generator. The generated thought  $th_t$  from the Thought Generator will be the input to the Action Generator and an action  $a_t$  is predicted by the Action Generator. (Left): The Thought Generator. We employ an LSTM [33] to embed the thought history and a transformer encoder to process both the mission and thought history. The text input is then fused with the visual observation input using FiLM [34]. (Right): The Action Generator is largely similar to the BabyAI agent [26], with the primary difference being the additional embedding of the thought generated by the Thought Generator.
Algorithm 1 Thought Cloning
1: Input: thought dataset  $\mathcal{D} = \{D_i\}_{i=1}^N$ , where each  $D_i = (m, \{(o_t, th_t, a_t)\}_{t=1}^T)$ . Thought Generator  $\pi_{\theta_u}(th|o, m, \{history_th\})$ , Action Generator  $\pi_{\theta_t}(a|o, m, th)$
2: while training do
3: for each  $\bar{D}_i = (m, \{(o_t, th_t, a_t)\}_{t=1}^T)$  in  $\mathcal{D}$  do
4: for each  $(o_t, th_t, a_t)$  in  $D_i$  do
5: Generate thought sentence  $\hat{t} h_{t} = \pi_{\theta_{u}}(\cdot |m,\{o_{\tau}\}_{\tau = 1}^{t},\{th_{\tau}\}_{\tau = 1}^{t - 1})$
6: Predict action probability distribution  $\hat{a}_t = \pi_{\theta_l}(\cdot|m, \{o_\tau\}_{\tau=1}^t, t\hat{h}_t)$
7: Compute the loss:  $\mathcal{L}(\theta_u,\theta_l) = \mathcal{L}_{CE}(a_t,\hat{a}_t) + \alpha \mathcal{L}_{CE}(th_t,\dot{th}_t) - \beta H(\hat{a}_t)$
8: Update the policy network parameters  $\theta_{u},\theta_{l}$  by minimizing  $\mathcal{L}(\theta_u,\theta_l)$
9: end for
0: end for
11: end while
parameters as TC. Results show it does not perform nearly as well as TC (Fig. 7 and Table 2: "Pure BC architecture (and matched num parameters)).
Also, one might argue that TC gets more data (one set of actions and one set of words per episode), and thus that a proper BC control is to give BC twice as much data. A counter is that such a control is unnecessary because noticing and harnessing this additional (often free) and heretofore ignored data stream is a central contribution of this paper, so showing that using this data improves things is the main comparison to be made. However, for completeness, we ran an experiment giving BC twice as much data: results show BC with twice as much data is still far slower to learn and has significantly worse performance at convergence (Fig. 7 and Table 2: "BC w/ 2x data"). Additionally, a prior work Think Before You Act [54] has a similar number of episodes of actions and language data in the same domain, and their model has more parameters than ours. Results show that TC also outperforms that method (Table 2).
Finally, to further confirm that the superiority of TC did not come from both more data and parameters, an ablation study was conducted on a pure BC architecture with the same number of parameters as TC but also with 2x data. The results show that it also underperformed TC (Fig. 7 and Table 2: "Pure BC architecture (and matched num parameters), 2x data"). All of the aforementioned results highlight the effectiveness of the TC framework and the importance of the thought data.
Figure 7: Training progress comparison of Thought Cloning (TC), Behavioral Cloning (BC), and other ablation variants. The BC architecture is identical to the Action Generator of TC and the TC w/o Imitating Thought has the same architecture as TC, but without the TC loss (and has one minor architectural change that is required since it does not have this loss; see response to reviewer gnSR). BC w/ 2x Data trains BC with 2x more episodes worth of actions than TC and BC. Pure BC architecture (and matched num parameters) is to augment BC architecture to have matched parameters with TC. Pure BC architecture (and matched num parameters), 2x data trains BC with both 2x episodes of actions and a BC architecture scaled up to match the number of parameters of TC. See Table 2 in this document for more details on the number of parameters and amount of data for each treatment. The error bars are the  $95\%$  confidence intervals from five runs of experiments. The results show that TC learns dramatically faster and has significantly higher performance at convergence compared with BC and all ablation variants, which indicates the advantages of TC are not solely due to additional training data or model capacity.
C Synthetic Human Thought Dataset
Fig. 8 presents an example trajectory. We translate the inner state of the BabyAI Oracle Solver (called "Bot" in [26]) into natural language thoughts. These thoughts outline the current plan for task completion and also describe the underlying intentions behind these plans, as the same low-level plan can serve different stated high-level purposes. For instance, the plan could be to "open the red door" with the intention of "completing the open mission" or "exploring". The segments with inserted noise are marked in red in Fig. 8.
D Example on Diagnosing Agents by Observing Thoughts
In this section, we provide an example of one time when we were able to diagnose Thought Cloning (TC) agents by observing their thoughts during the development phase of this paper. In the early stages of development, we trained the TC agent with a constant teacher-forcing strategy. We observed that during testing, the agents often got stuck persisting with incorrect thoughts and did not actively explore new ideas. For instance, in the top right example in Fig. 9, after  $t = 53$ , the agent persistently attempted to implement the incorrect thought "pickup blue box to complete putnext mission" until it reached the maximum step limit, without seeking new ideas. This observation led us to realize that, as we exclusively trained the agent with oracle thoughts via a teacher-forcing strategy, the agent had never practiced dealing with incorrect thoughts and consequently had not learned to recover from them by trying alternative ideas. Subsequently before this realization, we had attempted to transition directly to auto-regressive training after the teacher-forcing training stage. However, the agent then started to generate nonsensical thoughts. The trajectory in Fig. 9 (left) shows nonsensical thoughts such as open blue ( $t = 24$ ) and pickup door door ( $t = 75$ ) being generated when a constant auto-regressive strategy is applied. Because of the realization from being able to observe the agent's
mission: put a key next to the blue ball and pick up the green key
t=1-5: Thought "pickup a key to complete PutNext mission", Action [forward, right, ..., forward]
t=6: Thought "pickup blue key to complete PutNext mission", Action [pickup]
t=7-12: Thought "explore unseen area", Action [right, right, ..., right]
t=13-14: Thought "open green door to explore", Action [left, left, ..., toggle]
t=15-22: Thought "explore unseen area", Action [forward, forward, ..., left]
t=23-28: Thought "open red door to explore", Action [forward, forward, ..., toggle]
t=29: Thought "explore unseen area", Action [forward]
t=30-38: Thought "drop blue key to complete PutNext mission", Action [forward, forward, ..., drop]
t=39-44: Thought "explore unseen area", Action [right, right, ..., left]
t=45-48: Thought "open green door to explore", Action [left, left, ..., toggle]
t=49-56: Thought "explore unseen area", Action [forward, forward, ..., left]
t=57-69: Thought "open green door to explore", Action [forward, forward, ..., toggle]
t=70-77: Thought "explore unseen area", Action [forward, forward, ..., right]
t=78-88: Thought "open grey door to explore", Action [forward, forward, ..., toggle]
t=89-95: Thought "explore unseen area", Action [forward, forward, ..., right]
t=96-101: Thought "open purple door to explore", Action [forward, forward, ..., toggle]
t=102-110: Thought "explore unseen area", Action [forward, forward, ..., r
t=111-122: Thought "open red door to explore", Action [forward, forward, ..., toggle]
t=123-130: Thought "explore unseen area", Action [forward, forward, ..., right]
t=131-134: Thought "open grey door to explore", Action [right, right, ..., forward]
t=135-137: Thought "pickup green box to complete PutNext mission", Action [left, left, left]
t=138-204: Thought "open grey door to explore", Action [right, right, ..., toggle]
t=205: Thought "explore unseen area", Action [forward]
t=206-213: Thought "pickup green key to complete pickup mission", Action [forward,...,pickup]
Figure 8: Example trajectories of the synthetic human thought dataset. The inserted noisy segments are highlighted in red.
mission: put the green box next to the purple door after you go to the purple door and put a blue box next to the purple door
Figure 9: Example trajectories of agents trained with different strategies. Constant teacher-forcing training refers to exclusively training with the teacher-forcing strategy. In this scenario, the agent does not learn to recover from incorrect thoughts. Once it adopts an incorrect thought, it continues to follow this thought for thousands of time-steps until it reaches the maximum step count (top right from  $t = 53$  to  $t = 2880$ ). Constant auto-regressive training after teacher-forcing training implies directly transitioning to auto-regressive training following an initial phase of teacher-forcing training. In this case, agents begin to generate nonsensical thoughts, as shown on the left, such as open blue at  $t = 24$  (left) and pickup door door at  $t = 75$  (left). Gradual decay of teacher-forcing rate involves gradually reducing the ratio of teacher-forcing during training. This strategy is adopted in the final version of Thought Cloning. In this setting, the agent might generate some incorrect thoughts as shown at  $t = 53$  (bottom right), but it can recover from these errors to explore new ideas, as evidenced at  $t = 131$  (bottom right).
Table 2: Success rates of TC, BC, ablation variants, and the related work Think Before You Act [54] on BabyAI BossLevel. See the descriptions of ablation variants in the Fig. 7 Caption. The Success Rate is the mean and standard deviation from 5 experiments (except Think Before You Act, which did not include the number of experiments in their paper). The num (nubmer of) parameters of Think Before You Act is estimated from the number of parameters of the Huggingface GPT-2 model [62]. The results show TC outperforms all other methods, illustrating the superiority of TC, and supporting our argument that the advantages of TC are not solely due to additional training data or model capacity.
Method	Num Parameters	Data	Success Rate (%)
Behavioral Cloning	20.6M	1M Episodes Actions	91.2 ± 0.9
Thought Cloning w/o Imitating Thought	82.5M	1M Episodes Actions	65.5 ± 12.4
BC w/ 2x Data	20.6M	2M Episodes Actions	91.4 ± 1.7
Pure BC architecture (and matched num parameters)	83.9M	1M Episodes Actions	91.9 ± 2.1
Pure BC architecture (and matched num parameters), 2x data	83.9M	2M Episodes Actions	92.7 ± 1.0
Think Before You Act [54]	124M (estimated)	1M Episodes Actions + Language	85.2 ± 0.5
Thought Cloning	82.5M	1M Episodes Actions + Language	96.2 ± 0.8
thoughts, we adopted a gradual decay schedule for teacher-forcing rates during training. As shown in Fig. 9 (bottom right), the agent in this setting was able to explore new ideas after failing on an incorrect thought, and it rarely generate nonsensical thoughts. For example, the agent generates an incorrect thought at  $t = 53$ , but it can recover from these errors to explore new ideas, e.g. open grey door to explore. Because we can observe the TC agents thinking out loud, we are able to identify the issue and improve the agent's performance. Without this visibility into the agent's thoughts, simply observing their actions would have made it much harder to pinpoint the underlying problems.