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

Shengran Hu1,2
<EMAIL_ADDRESS>&Jeff Clune1,2,3
<EMAIL_ADDRESS>&
1Department of Computer Science, University of British Columbia
2Vector Institute
3Canada CIFAR AI Chair

###### 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.111The code and dataset are available in
https://github.com/ShengranHu/Thought-Cloning.

##  1 Introduction

Language may be the key to what separates humans from all other animals,
endowing us with an amazing level of general intelligence [1, 2, 3, 4].
Crucially, the benefits of language are not confined to improving our ability
to communicate with others: language also helps us _think_ better [2, 3, 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 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_. Thinking in language helps humans generalize, extrapolate, adapt
to new situations, combine old knowledge in new ways, explore, plan, replan
when necessary or beneficial, and the list goes on [2, 3, 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, 9, 10], 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, 11, 12, 13]. 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 [14,
15], because such LLMs are not trained on data where humans think out loud
_while acting_. Thought data, such as YouTube videos and transcripts [16, 17],
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 [17]. This thought data is
greatly valuable and widely available (Section 1), but has not yet been
extensively explored, and this work hopes to encourage further research into
the utilization of thought data to teach thinking skills to agents.

Provided we can solve the real, significant challenges of AI Safety and
existential risk [18, 19, 20, 21, 22], 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 [23], 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 [24]. 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.

##  2 Proposed Method

Figure 1: Overall framework for Thought Cloning (TC). The TC agent has two
components: the Upper-Level and Lower-level Components. At each timestep, the
TC agent receives an observation, a mission, and a history of thoughts as
inputs. The Upper-Level Component generates thoughts, and the Lower-Level
Component 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 [25, 26], such as Behavioral Cloning
[23], 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 [17] and then retrieving
the corresponding transcripts, we can obtain a thought dataset that contains
both human thinking and action [17, 16]. 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. To
validate Thought Cloning, we construct a synthetic thought dataset to simulate
having internet-scale datasets (see 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 Upper-level
Component responsible for thought generation, and a Lower-level Component
tasked with executing actions based on the thoughts generated by the Upper-
level Component. While different choices of what to condition the Upper-level
and Lower-level Components 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}}(th_{t}|m,\\{o_{\tau}\\}_{\tau=1}^{t},\\{th_{\tau}\\}_{\tau=1}^{t-1})-\log\pi_{\theta_{l}}(a_{t}|m,\\{o_{\tau}\\}_{\tau=1}^{t},th_{t})$
(1)

Here, $\theta_{u}$ and $\theta_{l}$ represent the weights for the Upper-Level
and Lower-level Components; $\alpha$ represents the coefficient for Thought
Cloning loss; $th$, $o$, $a$, and $m$ denote thought, observation, action, and
mission, as previously described.

For more complex or large-scale scenarios, the Upper-level Component can be
implemented with pre-trained Vision-Language Models (VLM) either zero-shot or
fine-tuned [27], while the Lower-level Component can be trained from scratch
or adapted from existing language-conditioned controllers in the target domain
[6, 14]. In this paper, we base both components on the BabyAI 1.1 model
architecture [28], which utilizes a memory-augmented architecture–an LSTM
[29]–to address the partial observability challenge. The model also employs
FiLM [30] 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 Experimental Results

###  3.1 Domain and Synthetic Thought Data

Figure 2:  Left: A BabyAI [24] 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.

This paper employs BabyAI [24], 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, boxs, doors_) and the initial position of the
agent are randomly distributed across a $27\times 27$ grid world containing
nine $3\times 3$ 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 [24]. 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 [24]) 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 B.

###  3.2 Experiment Setup

To verify the effectiveness of _learning to think_ , we compare our Thought
Cloning (TC) approach to the conventional Imitation Learning algorithm,
Behavioral Cloning (BC). BC shares most of its architecture with the Lower-
level Component of TC (Fig. 1), and because it is trained only on action loss,
it does not encode thought like the lower-level component 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. The architecture of the variant is mostly identical to the TC
architecture, except for a minor architectural difference where the latent
vector from the upper level is directly input to the lower level as thought.
This adjustment is required when training occurs without thought supervision
because the discrete sampling to produce a sequence of thoughts (words) is not
differentiable.

Our training setup is based on BabyAI [24, 28], with BC agents identical to
the Imitation Learning baseline from it. The training iterates for 8 epochs on
the 1 million episode dataset, corresponding to a total of 160 training steps
and $7\times 10^{8}$ training frames. The Thought Cloning loss parameter
$\alpha$ (Eq. 1) is set to 2. During training, we employ teacher-forcing [31],
which is adopted when decoding thoughts. It conditions the Lower-level
Component on the ground truth thoughts from the dataset. The teacher-forcing
ratio linearly decreases from 100% to 0% from the 10th training step to the
end. 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 5 are calculated
from five independent runs. The success rate results presented 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

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 Lower-level Component 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 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.

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 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

(a)

(b)

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). (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 Upper-level Component 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.

This section compares the generalization abilities between 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 approximately 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 calculation formula for
Cognitive Difficulty is adapted from the maxStep parameter calculation in
BabyAI environments [24] and is given by ($\\#\text{ of \\{\emph{GoTo, PickUp,
OpenDoor}\\}}+2\times\\#\text{ of \\{\emph{PutNextTo}}\\}+\\#\text{ 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 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. 4(a)),
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. 4(b)). 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

(a)

(b)

Figure 5: (a): A heatmap illustrating the Future Action Declaration Score, a
metric designed to evaluate the interpretability of Thought Cloning agents
(Section 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 results show that Precrime Intervention
effectively eliminates almost all unsafe behaviors.

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 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. 5(a)). Interestingly, TC agents also scored near-
perfectly across all out-of-distribution environments (rest of Fig. 5(a)),
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. 5(b)). 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 phase 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 has 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, which prevented stable training and harmed
performance. 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: we gradually decay the teacher-forcing
rate (fraction) during training, which dramatically improved performance
(Section. 3.3). Supplementary Material C contains more details about this
example.

Lastly, TC enables steerability. 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. We can also take advantage of this
capability to help agents in challenging tasks. The TC agent, when provided
with oracle high-level thoughts, is capable of near-perfect performance across
almost all environments (Fig. 4(a)). Additional empirical evidence of this is
the high fraction of tasks successfully performed by following these oracle
thoughts. This oracle task success rate starts at a median of 96.0 (95%
confident interval: -8.4,+4.0)% for tasks furthest out of distribution
(Behavioral Difficulty $>$ 450 and Cognitive difficulty $=$ 9), and rises to a
near-perfect median of 99.2 (95% confident interval: -5.5,+0.8)% when within
the training distribution. These findings highlight the promising potential of
TC agents in effectively collaborating with humans to accomplish challenging
missions.

##  4 Related Works

###  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
[32, 33, 34, 35, 36]. 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 [37, 38, 39]. The second category of
work involves pre-trained LLMs that generate plans in language for RL systems.
Earlier works [14, 15] 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 [15, 40, 41]. 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 [27], in which a pre-trained Vision-Language Model is adopted as the
planner, allowing it to recognize patterns from observations directly.
However, PALM-E was not trained on synchronized videos of humans thinking out
loud and acting, meaning it does 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.

###  4.2 Learning from Dataset Aligning Action and Language

Several studies have recognized the value of datasets that align action with
language. DIAL [42] 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, (SL)3 [43], 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. [44] collected a dataset from two
human players collaborating on an RTS game. However, the agent in [44] 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 [45].
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 [44, 45] 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.

##  5 Discussion and Conclusion

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, and where we expect it
to truly shine, 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 [17, 16] (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 experimental results presented in Section 3.4, which
suggest that the current bottleneck in agent capability is its high-level
thinking, a skill that could be substantially enhanced by scaling to vast
online data.

Of course, there are also risks associated with such agents, from the same
ones that occur with language models trained on our written thoughts, such as
bias [46, 47, 48] and otherwise emulating undesirable human thoughts, but also
in terms of AI Safety and Existential Risk [18, 19, 21].

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.

## Acknowledgments and Disclosure of Funding

This work was supported by the Vector Institute, a grant from Schmidt Futures,
an NSERC Discovery Grant, and a generous donation from Rafael Cosman. We also
thank Aaron Dharna, Ben Norman, and Jenny Zhang (sorted alphabetically) in our
lab at the University of British Columbia for insightful discussions.

## References

  * Berwick and Chomsky [2016] Robert C Berwick and Noam Chomsky.  _Why Only Us: Language and Evolution_.  MIT press, 2016.
  * Premack [2004] David Premack.  Is language the key to human intelligence?  _Science_ , 303(5656):318–320, 2004.
  * Chomsky et al. [2006] Noam Chomsky et al.  _Language and Mind_.  Cambridge University Press, 2006.
  * Pinker [2003] Steven Pinker.  _The Language Instinct: How the Mind Creates Language_.  Penguin UK, 2003.
  * Shridhar et al. [2021] Mohit Shridhar, Lucas Manuelli, and Dieter Fox.  Cliport: What and where pathways for robotic manipulation.  In _Conference on Robot Learning_ , 2021.
  * Huang et al. [2022a] 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_ , 2022a.
  * Brown et al. [2020] 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.
  * Team et al. [2021] 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.
  * Wei et al. [2022] 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.
  * Marcus [2001] Gary Marcus.  The algebraic mind, 2001.
  * Vygotsky et al. [2011] Lev Vygotsky et al.  _Interaction Between Learning and Development_.  Linköpings universitet, 2011.
  * Bransford et al. [2000] John D Bransford, Ann L Brown, Rodney R Cocking, et al.  _How People Learn_ , volume 11.  Washington, DC: National academy press, 2000.
  * Bada and Olusegun [2015] 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.
  * Ahn et al. [2022] 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, Kuang-Huei Lee, Sergey Levine, Yao Lu, Linda Luu, Carolina Parada, Peter Pastor, Jornell Quiambao, Kanishka Rao, Jarek Rettinghouse, Diego Reyes, Pierre Sermanet, Nicolas Sievers, Clayton Tan, Alexander Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Mengyuan Yan, and Andy Zeng.  Do As I Can and Not As I Say: Grounding language in robotic affordances.  In _arXiv preprint arXiv:2204.01691_ , 2022.
  * Huang et al. [2022b] Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch.  Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.  In _International Conference on Machine Learning_ , pages 9118–9147. PMLR, 2022b.
  * Fan et al. [2022] Linxi Fan, Guanzhi Wang, Yunfan Jiang, Ajay Mandlekar, Yuncong Yang, Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar.  MineDojo: Building open-ended embodied agents with internet-scale knowledge.  In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_ , 2022.  URL https://openreview.net/forum?id=rc8o_j8I8PX.
  * Baker et al. [2022] Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune.  Video pretraining (VPT): Learning to act by watching unlabeled online videos.  _Advances in Neural Information Processing Systems_ , 35:24639–24654, 2022.
  * Ecoffet et al. [2020] Adrien Ecoffet, Jeff Clune, and Joel Lehman.  Open questions in creating safe open-ended AI: Tensions between control and creativity.  In _Conference on Artificial Life_ , pages 27–35. MIT Press, 2020\.
  * Bostrom [2002] N Bostrom.  Existential Risks: analyzing human extinction scenarios and related hazards.  _Journal of Evolution and Technology_ , 9, 2002.
  * Sotala and Yampolskiy [2014] Kaj Sotala and Roman V Yampolskiy.  Responses to catastrophic AGI risk: a survey.  _Physica Scripta_ , 90(1):018001, 2014.
  * Yudkowsky et al. [2008] Eliezer Yudkowsky et al.  Artificial Intelligence as a positive and negative factor in global risk.  _Global catastrophic risks_ , 1(303):184, 2008\.
  * Dietterich and Horvitz [2015] Thomas G Dietterich and Eric J Horvitz.  Rise of concerns about AI: reflections and directions.  _Communications of the ACM_ , 58(10):38–40, 2015\.
  * Bain and Sammut [1995] Michael Bain and Claude Sammut.  A framework for behavioural cloning.  In _Machine Intelligence 15_ , pages 103–129, 1995.
  * Chevalier-Boisvert et al. [2019] Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio.  BabyAI: First steps towards grounded language learning with a human in the loop.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rJeXCo0cYX.
  * Hussein et al. [2017] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne.  Imitation learning: A survey of learning methods.  _ACM Computation Survey_ , 50(2), apr 2017.  ISSN 0360-0300.  doi: 10.1145/3054912.  URL https://doi.org/10.1145/3054912.
  * Pomerleau [1988] Dean A Pomerleau.  Alvinn: An autonomous land vehicle in a neural network.  _Advances in Neural Information Processing Systems_ , 1, 1988.
  * Driess et al. [2023] Danny Driess, Fei Xia, Mehdi S. M. Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, Wenlong Huang, Yevgen Chebotar, Pierre Sermanet, Daniel Duckworth, Sergey Levine, Vincent Vanhoucke, Karol Hausman, Marc Toussaint, Klaus Greff, Andy Zeng, Igor Mordatch, and Pete Florence.  PaLM-E: An embodied multimodal language model.  In _arXiv preprint arXiv:2303.03378_ , 2023.
  * Hui et al. [2020] David Yu-Tung Hui, Maxime Chevalier-Boisvert, Dzmitry Bahdanau, and Yoshua Bengio.  Babyai 1.1.  _arXiv preprint arXiv:2007.12770_ , 2020.
  * Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780, 1997.
  * Perez et al. [2018] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron Courville.  FiLM: Visual reasoning with a general conditioning layer.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 32(1), 2018.  doi: 10.1609/aaai.v32i1.11671.  URL https://ojs.aaai.org/index.php/AAAI/article/view/11671.
  * Williams and Zipser [1989] Ronald J Williams and David Zipser.  A learning algorithm for continually running fully recurrent neural networks.  _Neural computation_ , 1(2):270–280, 1989.
  * Shu et al. [2018] Tianmin Shu, Caiming Xiong, and Richard Socher.  Hierarchical and interpretable skill acquisition in multi-task reinforcement learning.  In _International Conference on Learning Representations_ , 2018.
  * Jiang et al. [2019] Yiding Jiang, Shixiang Shane Gu, Kevin P Murphy, and Chelsea Finn.  Language as an abstraction for hierarchical deep reinforcement learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Blukis et al. [2022] Valts Blukis, Chris Paxton, Dieter Fox, Animesh Garg, and Yoav Artzi.  A persistent spatial semantic representation for high-level natural language instruction execution.  In _Conference on Robot Learning_ , pages 706–717. PMLR, 2022.
  * Zhang and Chai [2021] Yichi Zhang and Joyce Chai.  Hierarchical task learning from language instructions with unified transformers and self-monitoring.  In _Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021_ , pages 4202–4213, Online, August 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.findings-acl.368.  URL https://aclanthology.org/2021.findings-acl.368.
  * Zawalski et al. [2023] Michał Zawalski, Michał Tyrolski, Konrad Czechowski, Tomasz Odrzygóźdź, Damian Stachura, Piotr Piękos, Yuhuai Wu, Łukasz Kuciński, and Piotr Miłoś.  Fast and Precise: Adjusting planning horizon with adaptive subgoal search.  In _International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=7JsGYvjE88d.
  * OpenAI [2023] OpenAI.  GPT-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Bai et al. [2022] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al.  Training a helpful and harmless assistant with reinforcement learning from human feedback.  _arXiv preprint arXiv:2204.05862_ , 2022.
  * Li et al. [2022] Xiang Lorraine Li, Adhiguna Kuncoro, Jordan Hoffmann, Cyprien de Masson d’Autume, Phil Blunsom, and Aida Nematzadeh.  A systematic investigation of commonsense knowledge in large language models.  In _Empirical Methods in Natural Language Processing_ , pages 11838–11855, 2022.
  * Wang et al. [2023] Zihao Wang, Shaofei Cai, Anji Liu, Xiaojian Ma, and Yitao Liang.  Describe, Explain, Plan and Select: Interactive planning with large language models enables open-world multi-task agents.  _arXiv preprint arXiv:2302.01560_ , 2023.
  * Lin et al. [2023] Kevin Lin, Christopher Agia, Toki Migimatsu, Marco Pavone, and Jeannette Bohg.  Text2Motion: From natural language instructions to feasible plans.  _arXiv preprint arXiv:2303.12153_ , 2023.
  * Xiao et al. [2022] Ted Xiao, Harris Chan, Pierre Sermanet, Ayzaan Wahid, Anthony Brohan, Karol Hausman, Sergey Levine, and Jonathan Tompson.  Robotic skill acquistion via instruction augmentation with vision-language models.  In _Workshop on Language and Robotics at CoRL 2022_ , 2022.
  * Sharma et al. [2022] Pratyusha Sharma, Antonio Torralba, and Jacob Andreas.  Skill induction and planning with latent language.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1713–1726, Dublin, Ireland, May 2022. Association for Computational Linguistics.  doi: 10.18653/v1/2022.acl-long.120.  URL https://aclanthology.org/2022.acl-long.120.
  * Hu et al. [2019] Hengyuan Hu, Denis Yarats, Qucheng Gong, Yuandong Tian, and Mike Lewis.  Hierarchical decision making by generating and following natural language instructions.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Mezghani et al. [2023] Lina Mezghani, Piotr Bojanowski, Karteek Alahari, and Sainbayar Sukhbaatar.  Think Before You Act: Unified policy for interleaving language reasoning with actions.  In _Workshop on Reincarnating Reinforcement Learning at ICLR_ , 2023\.
  * Mehrabi et al. [2021] Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan.  A survey on bias and fairness in machine learning.  _ACM Computing Surveys (CSUR)_ , 54(6):1–35, 2021\.
  * Gehman et al. [2020] Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A. Smith.  RealToxicityPrompts: Evaluating neural toxic degeneration in language models.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pages 3356–3369, Online, November 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.findings-emnlp.301.  URL https://aclanthology.org/2020.findings-emnlp.301.
  * Bender et al. [2021] Emily M Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell.  On the dangers of stochastic parrots: Can language models be too big.  In _Proceedings of the 2021 ACM conference on fairness, accountability, and transparency_ , pages 610–623, 2021.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Mikolov et al. [2013] Tomas Mikolov, Kai Chen, Gregory S. Corrado, and Jeffrey Dean.  Efficient estimation of word representations in vector space.  In _International Conference on Learning Representations_ , 2013.
  * Kang et al. [2019] Gi-Cheon Kang, Jaeseo Lim, and Byoung-Tak Zhang.  Dual attention networks for visual reference resolution in visual dialog.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing_ , pages 2024–2033, 2019.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Goyal et al. [2017] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  _arXiv preprint arXiv:1706.02677_ , 2017.
  * Micikevicius et al. [2018] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al.  Mixed precision training.  In _International Conference on Learning Representations_ , 2018.

## Supplementary Material

##  Appendix A Architecture and Training Details

Figure 6: Detailed architecture for Thought Cloning (TC) agent. 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 Upper-level
Component. The generated thought $th_{t}$ from the Upper-level Component will
be the input to the Lower-level Component and an action $a_{t}$ is predicted
by the Lower-level Component. (Left): The Upper-level Component. We employ an
LSTM [29] 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 [30]. (Right): The Lower-level Component
is largely similar to the BabyAI agent [24], with the primary difference being
the additional embedding of the thought generated by the Upper-Level
Component.  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})$, upper level component
$\pi_{\theta_{u}}(th|o,m,\\{history\\_th\\})$, lower level component
$\pi_{\theta_{l}}(a|o,m,th)$

2:while training do

3: for each $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 upper level thought sentence
$\hat{th}_{t}=\pi_{\theta_{u}}(\cdot|m,\\{o_{\tau}\\}_{\tau=1}^{t},\\{th_{\tau}\\}_{\tau=1}^{t-1})$

6: Predict lower level action probability distribution
$\hat{a}_{t}=\pi_{\theta_{l}}(\cdot|m,\\{o_{\tau}\\}_{\tau=1}^{t},\hat{th_{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},\hat{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

10: end for

11:end while

For full transparency, replicability, and to facilitate future research
building on our work, we are releasing both the source code and model weights.
Additionally, we provide key details necessary for the evaluation and
replication of our work in this supplementary information. The architectural
details of Thought Cloning models are shown in Fig. 6. As in [28], 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 [49] 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 [50] encoding approach. In the Upper-level Component,
a Transformer encoder [49, 51] 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 Lower-level Component 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
[28], as shown in Table 1.

The pseudocode for Thought Cloning (TC) training framework is shown in
Algorithm 1. In the loss function, we follow [28] by including an entropy term
for actions. The Adam optimizer [52] 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 [28, 24], 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 5 training steps, linearly increasing from $1e^{-4}$ to $5e^{-4}$,
and then decaying by 50% at 120th training steps, similar to the practices in
[24, 53]. 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 Upper-
level Component ($\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 [54].
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 $\beta_{1}$ | 0.9
Adam $\beta_{2}$ | 0.99
Adam $\epsilon$ | $10^{-5}$
Entropy Coefficient | 0.01
Image Embedding Dimension | 128
Text Embedding Dimension | 256
Memory Dimension | 2048
Figure 7: Example trajectories of the synthetic human thought dataset. The
inserted noisy segments are highlighted in red.

##  Appendix B Synthetic Human Thought Dataset

Fig. 7 presents an example trajectory. We translate the inner state of the
BabyAI Oracle Solver (called “Bot” in [24]) 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. 7.

##  Appendix C Example on Diagnosing Agents by Observing Thoughts

Figure 8: 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).

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. 8, 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. 8 (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 thoughts, we
adopted a gradual decay schedule for teacher-forcing rates during training. As
shown in Fig. 8 (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.


Table 4: Summary of the mock tests we performed to validate the modelling of the non-linear matter power spectrum. Mock tests 1-3 validate the non-linear correction to the matter power spectrum due to dark matter only. The remaining tests validate the baryon feedback mitigation strategies we test in this work; scale cuts (4-7), HM20 (8-14), SP(k) (15-20) and BCEmu (21-29). We test analysing the mocks with various prior and analysis choices for each of the four baryon model approaches. If not otherwise specified, the mocks were created with HM20 as the dark matter-only non-linear correction to the matter power spectrum, and three ’baryonic feedback scenarios’; BAHAMAS, cosmo-OWLS and $A_{\rm{mod}}=0.82$. The ’mock’ column therefore states the model/baryonic feedback scenario used to create the mock and the ’model’ column labels the model/analysis choices used to analyse the mock’. For example, a mock of ’EuclidEmu’ and model of HM20-DM’ denotes a mock created with Euclid Emulator for the non-linear correction to the matter power spectrum, and analysed with HM20. We report the $S_{8}$ constraints and 68% confidence level using the mean-marginal approach $S_{8}$, with $\Delta S_{8}$ quantifying the offset from the true $S_{8}=0.7805$, as a fraction of the $1\sigma$ error. No. | Mock | Model | $S_{8}$ | $\Delta S_{8}$
---|---|---|---|---
1 | HM20-DM | HM20-DM | $0.781^{+0.012}_{0.011}$ | $0.044\sigma$
2 | HM20-DM | HM20-DM $\sum m_{\nu}$: [0.06,6] eV | $0.775^{+0.012}_{0.011}$ | $-0.461\sigma$
3 | EuclidEmu-DM | HM20-DM | $0.785^{+0.012}_{0.011}$ | $0.403\sigma$
4 | EuclidEmu-DM | HM20-DM + Scale cuts | $0.784^{+0.017}_{0.017}$ | $0.239\sigma$
5 | BAHAMAS 8.0 | Scale cuts | $0.771^{+0.015}_{0.018}$ | $-0.566\sigma$
6 | cosmo-OWLS 8.5 | Scale cuts | $0.763^{+0.013}_{0.017}$ | $-1.191\sigma$
7 | $A_{\rm mod}=0.820$ | Scale cuts | $0.762^{+0.016}_{0.018}$ | $-1.114\sigma$
8 | BAHAMAS 8.0 | HM20 BAHAMAS | $0.772^{+0.013}_{0.012}$ | $-0.675\sigma$
9 | cosmo-OWLS 8.5 | HM20 BAHAMAS | $0.765^{+0.013}_{0.014}$ | $-1.182\sigma$
10 | $A_{\rm mod}=0.820$ | HM20 BAHAMAS | $0.765^{+0.012}_{0.013}$ | $-1.270\sigma$
11 | BAHAMAS 8.0 | HM20 wide | $0.792^{+0.021}_{0.023}$ | $0.509\sigma$
12 | BAHAMAS 8.0 | HM20 wide $\sum m_{\nu}$: [0.06,6] eV | $0.788^{+0.017}_{0.023}$ | $0.358\sigma$
13 | cosmo-OWLS 8.5 | HM20 wide | $0.792^{+0.021}_{0.022}$ | $0.537\sigma$
14 | $A_{\rm mod}=0.820$ | HM20 wide | $0.777^{+0.019}_{0.024}$ | $-0.180\sigma$
15 | BAHAMAS 8.0 | SP(k): X-ray | $0.781^{+0.017}_{0.018}$ | $0.017\sigma$
16 | cosmo-OWLS 8.5 | SP(k): X-ray | $0.775^{+0.016}_{0.018}$ | $-0.347\sigma$
17 | $A_{\rm mod}=0.820$ | SP(k): X-ray | $0.761^{+0.012}_{0.012}$ | $-1.643\sigma$
18 | BAHAMAS 8.0 | SP(k) wide | $0.780^{+0.015}_{0.016}$ | $-0.010\sigma$
19 | cosmo-OWLS 8.5 | SP(k) wide | $0.774^{+0.015}_{0.017}$ | $-0.424\sigma$
20 | $A_{\rm mod}=0.820$ | SP(k) wide | $0.773^{+0.015}_{0.016}$ | $-0.479\sigma$
21 | BAHAMAS 8.0 | BCEmu1 | $0.776^{+0.017}_{0.018}$ | $-0.248\sigma$
22 | cosmo-OWLS 8.5 | BCEmu1 | $0.773^{+0.015}_{0.017}$ | $-0.489\sigma$
23 | $A_{\rm mod}=0.820$ | BCEmu1 | $0.759^{+0.016}_{0.019}$ | $-1.239\sigma$
24 | BAHAMAS 8.0 | BCEmu7 | $0.792^{+0.020}_{0.025}$ | $0.510\sigma$
25 | BAHAMAS 8.0 | BCEmu7 $\sum m_{\nu}$: [0.06,6] eV | $0.785^{+0.020}_{0.025}$ | $0.202\sigma$
26 | BAHAMAS 8.0 + EuclidEmu-DM | BCEmu7 | $0.791^{+0.020}_{0.025}$ | $0.466\sigma$
27 | BAHAMAS 8.0 + EuclidEmu-DM | BCEmu7 $\sum m_{\nu}$: [0.06,6] eV | $0.787^{+0.019}_{0.027}$ | $0.287\sigma$
28 | cosmo-OWLS 8.5 | BCEmu7 | $0.791^{+0.022}_{0.026}$ | $0.439\sigma$
29 | $A_{\rm mod}=0.820$ | BCEmu7 | $0.777^{+0.018}_{0.024}$ | $-0.165\sigma$

We validate the choice to model the dark matter-only power spectrum with HM20
using dark matter-only mocks created with HM20 and EE2. Here we summarise the
findings of mock tests 1-4.

  * •

This test analyses a mock with the same model choices used to create it and is
useful for identifying projection effects. We find these to be present for
$\Omega_{\rm m}$ and $\sigma_{8}$, which are under- and over-estimated,
respectively, by $\sim$0.5$\sigma$.

  * •

When modelling dark-matter with HM20 in a HM20 generated mock, allowing the
neutrino mass to vary with the prior $\sum m_{\nu}$: [0.06,6] eV improves the
recovery of $\Omega_{\rm m}$ and $\sigma_{8}$. It however decreases the
accuracy in the recovery of $S_{8}$, resulting in an underestimation of
$0.5\sigma$.

  * •

Modelling the non-linear dark matter-only power spectrum with HM20 in a EE2
generated mock over estimates $S_{8}$ by $\sim 0.4\sigma$, which is reduced
marginally to $\sim 0.2\sigma$ when scale cuts are applied. We note that DES &
KiDS Collaborations (2023) find a smaller bias in the mean of $\sim 0.1\sigma$
when scale cuts are used. In this mock analysis there are two set-up
differences that could explain this: we use a higher value of $\Omega_{\rm m}$
to create the mock and we marginalise over $A_{\rm s}$, instead of $S_{8}$.

Next we validate each of our four baryon models (scale cuts, HM20 with free
$\Theta_{\rm AGN}$, BCEmu and SP(k)) and their respective prior and analysis
variants. We test the modelling choices on three ‘baryonic feedback scenarios’
of increasing ‘extremity’ in terms of their impact on the matter power
spectrum suppression: BAHAMAS 8.0, cosmo-OWLS 8.5 and $A_{\rm{mod}}=0.82$.
Here we summarise the findings of the mock tests 5-29 and Fig. 10. Note that
in the figure, mocks are labelled in the format ’model used to create the
mock: model/analysis choices used to analyse the mock’. For example, ’BAHAMAS:
HM20 BAHAMAS’ denotes a mock created with BAHAMAS-like baryon feedback
scenario and analysed with HM20 with the BAHAMAS $\Theta_{\rm AGN}=7.6-8.0$
prior.

  * •

Optimised scale cuts underestimates $S_{8}$ by $0.5-1.2\sigma$ in the three
feedback scenarios we test, with the tension worsening with more extreme
feedback. As noted in Sec. 4.2, this is as expected given our choice of
feedback scenarios used to build the synthetic data. The OWLS-AGN scenario,
which was used to decide which angular scales of the measurement would be
discarded in the DES Y3 analysis predicts a less extreme feedback scenario
than the three scenarios used here.

  * •

The left panel of Fig. 10 compares prior choices of the HM20 model with
$\Theta_{\rm AGN}$ as a free parameter. We find that HM20 with the BAHAMAS
calibrated prior choice of $\Theta_{\rm AGN}=7.6-8.0$ underestimates $S_{8}$
by $\sim 0.7-1.2\sigma$, with the tension worsening with a more extreme
feedback scenario. Using the wide prior choice of $\Theta_{\rm AGN}=7.3-9.0$
overestimates $S_{8}$ in BAHAMAS 8.0 and cosmo-OWLS 8.5 by $\sim 0.5\sigma$,
but still underestimates $S_{8}$ in an $A_{\rm{mod}}=0.82$-like feedback
scenario by $\sim 0.2\sigma$.

  * •

The right panel of Fig. 10 shows that SP(k) with the X-ray informed prior
demonstrates the best recovery of $S_{8}$ in a BAHAMAS 8.0-like feedback
scenario of all the baryonic feedback modelling choices. In more extreme
feedback scenarios, it however underestimates $S_{8}$ by $\sim 0.3-1.6\sigma$.
Using the wide prior choice improves the recovery of $S_{8}$ to be within
$\sim 0.4\sigma$ for the three feedback scenarios.

  * •

The central panel of Fig. 10 demonstrates that BCEmu7 overestimates $S_{8}$ in
both BAHAMAS 8.0 and cosmo-OWLS 8.5 baryonic feedback scenarios by $\sim
0.4\sigma$, but underestimates $S_{8}$ in an $A_{\rm{mod}}=0.82$-like feedback
scenario by $\sim 0.2\sigma$. BCEmu1 underestimates $S_{8}$ by
$0.3-1.2\sigma$, again with the tension worsening when testing more extreme
feedback scenarios.

  * •

In the mock tests which over-estimate $S_{8}$, freeing the neutrino mass with
the prior $\sum m_{\nu}$: [0.06,6] eV appears to improve recovery of the input
cosmology. The shifts we see in $S_{8}$ are however consistent with the $\sim
0.2\sigma$ shift to lower values we find when testing the dark matter-only
modelling of HM20 with and without free neutrinos, i.e. tests 1 and 2.

Ultimately we find that choosing the restrictive analysis variant of a model
tends to result in an under-estimation of $S_{8}$. This motivates the use of
more flexible modelling choices in cosmic shear analyses, since these tend to
result in improved recovery of $S_{8}$, however this is at the consequence of
a loss in precision.

##  Appendix B Baryon feedback constraints

###  B.1 SP(k)

Figure 12: SP(k) marginalised posteriors when analysing the DES Y3 cosmic shear data with the wide prior on SP(k) (light green) and X-ray informed prior (dark green). We show the 68% and 95% confidence levels of the $S_{8}$ and the three model parameters (upper panel) and the mean total baryon fraction, $f_{\rm b}/\Omega_{\rm b}/\Omega_{\rm m}$, and $1\sigma$ uncertainty as a function of halo mass, $M_{500}$ (lower). The X-ray prior is derived from HSC-XXL $1\sigma$ constraints (black hatched Akino et al., 2022), in this plot scaled to the mean cosmology obtained from the $SPk$: X-ray analysis. Table 5: The constraints on the SP(k) parameters attained by the WL-only SP(k) analyses with both the wide and X-ray informed prior choices. We report the mean and 68% confidence levels.  Parameter | Wide prior | X-ray prior
---|---|---
$\alpha$ | $3.830^{+0.409}_{-0.363}$ | $4.153^{+0.061}_{-0.058}$
$\beta$ | $1.174^{+0.104}_{-0.158}$ | $1.206^{+0.043}_{-0.042}$
$\gamma$ | $0.519^{+0.265}_{-0.157}$ | $0.388^{+0.077}_{-0.076}$

The mean and 68% confidence levels on the SP(k) parameters constrained by the
WL-only SP(k) analyses with the wide and X-ray informed prior choices are
reported in Table 5. The corresponding marginalised posteriors for these
analyses are shown in the upper panel of Fig. 12. Since SP(k) directly maps
between the baryon fraction in groups and clusters and the matter power
spectrum suppression, we also show the constraints on the $f_{\rm
b}/\Omega_{\rm b}/\Omega_{\rm m}$-$M_{500}$ relation predicted when analysing
the DES cosmic shear with the two SP(k) analysis variants in the lower panel
of Fig. 12.

###  B.2 Baryonification

The mean and 68% confidence levels on the baryonification parameters
constrained by the WL-only BCEmu1, BCEmu3 and BCEmu7 analyses, in addition to
the WL + kSZ BCEmu7 joint analysis (for which the average halo mass of the kSZ
sample $M_{\rm h,200}$ is also reported) are given in Table 6. The
corresponding marginalised posteriors for these analyses are shown in Fig. 13.

Table 6: The constraints on the baryonification parameters attained by the WL-only BCEmu1, BCEmu3 and BCEmu7 analyses, as well as the BCEmu7 WL + kSZ BCEmu7 analysis. For the WL + kSZ analysis we also report the constraint on the average halo mass of the kSZ sample ($10^{13}M_{\odot}$). We report the mean and 68% confidence levels.  Param. | BCEmu1 | BCEmu3 | BCEmu7 | WL + kSZ
---|---|---|---|---
$\log_{10}M_{\rm c}$ | $13.42^{+0.60}_{-0.54}$ | $13.13^{+0.79}_{-0.63}$ | $13.06^{+1.01}_{-0.87}$ | $13.22^{+0.42}_{-0.29}$
$\theta_{\rm ej}$ | - | $5.31^{+2.17}_{-1.30}$ | $5.14^{+1.75}_{-1.74}$ | $5.15^{+1.46}_{-1.48}$
$\eta_{\delta}$ | - | $0.22^{+0.09}_{-0.12}$ | $0.23^{+0.10}_{-0.10}$ | $0.22^{+0.10}_{-0.11}$
$\mu$ | - | - | $1.00^{+0.50}_{-0.61}$ | $0.68^{+0.25}_{-0.58}$
$\gamma$ | - | - | $2.57^{+1.00}_{0.71}$ | $2.66^{+0.98}_{0.67}$
$\delta$ | - | - | $6.83^{+2.09}_{-2.43}$ | $7.35^{+2.16}_{-2.10}$
$\eta$ | - | - | $0.22^{+0.10}_{-0.11}$ | $0.21^{+0.08}_{-0.13}$
$M_{\rm h,200}$ | - | - | - | $4.10^{+1.55}_{-1.64}$
Figure 13: The marginalised posteriors for $S_{8}$, the seven baryonification
parameters, and the average halo mass og the kSZ sample $M_{\rm h,200}$
attained by the WL-only BCEmu1 (red solid line), BCEmu3 (light purple dotted
line) and BCEmu7 (pink dashed line) analyses, as well as the WL + kSZ BCEmu7
(dark purple line) analysis. The inner and outer contours show the 68% and 95%
confidence levels respectively. The grey dashed lines show the values the
BCEmu parameters are fixed to in the case of using BCEmu1 and BCEmu3, and the
mean halo mass of the CMASS sample reported by Schaan et al. (2021).

###  B.3 Prior choice for the halo mass of the kSZ sample

Given the significant scatter in the literature, we choose to include an
additional model parameter in the analysis, $M_{\rm h,200}$, corresponding to
the mean $M_{200}$ of the CMASS sample, with a prior range provided in Table
1. For example, the stacked stellar mass–halo mass relation of Sonnenfeld et
al. (2019) of CMASS galaxies derived from HSC galaxy–galaxy lensing
measurements, when combined with the stellar mass distribution of CMASS,
implies a mean $M_{200}\approx 0.5\times 10^{13}{\rm M}_{\odot}$, whereas the
abundance matching methods of Kravtsov et al. (2018), Moster et al. (2018),
and Behroozi et al. (2019) imply mean $M_{200}$ values of $2.7\times
10^{13}{\rm M}_{\odot}$, $4.4\times 10^{13}{\rm M}_{\odot}$, and $6.6\times
10^{13}{\rm M}_{\odot}$, respectively. A HOD-based analysis of the clustering
of BOSS CMASS galaxies by White et al. (2011) found a mean $M_{200}$ of
$3.6\times 10^{13}{\rm M}_{\odot}$. Given this large study-to-study variance,
we choose a flat prior range of $[0.5-7]\times 10^{13}{\rm M}_{\odot}$ and
marginalise over this parameter.

The combined WL + kSZ analysis constrains $M_{\rm
h,200}=4.098^{+1.548}_{-1.639}\times 10^{13}{\rm M}_{\odot}$, indicating we
are prior constrained at the $2\sigma$ level. For comparison to the remainder
of this work which generally quotes halo masses in $M_{500}$, this corresponds
to $M_{h,500}=3.01\times 10^{13}{\rm M}_{\odot}$, assuming a NFW profile and a
concentration-mass relation from Dutton & Macciò (2014).

In this appendix we explore the impact on our cosmology and baryon model
parameter constraints of choosing a wider prior $M_{\rm h,200}:[0.8,30]\times
10^{13}{\rm M}_{\odot}$. We also consider a fixed $M_{\rm h,200}$ analysis
using the mean mass of the CMASS sample as determined by Schaan et al. (2021)
and used in (Amodeo et al., 2021): $M_{\rm h,200}=3\times 10^{13}{\rm
M}_{\odot}$. Fig. 14 shows the marginalised posteriors for $\Omega_{\rm m}$,
$S_{8}$, $\log_{10}(M_{\rm c})$, $\theta_{\rm ej}$ and $M_{\rm h,200}$ in a WL
+ kSZ analysis with the different prior choices on $M_{\rm h,200}$. The halo
mass, $M_{\rm h,200}$, is correlated with $\log_{10}(M_{\rm c})$ and in the
case of a wider prior, both $M_{\rm h,200}$ and $\log_{10}(M_{\rm c})$ prefer
higher values. However, $S_{8}$ is relatively stable to the halo mass prior,
shifting by 0.3$\sigma$.

Figure 14: Upper: The stacked kSZ temperature profile at 98GHz as a function
of angular radius, $\theta$, centred on the group or cluster (bottom) when
varying the mean halo mass of CMASS galaxy sample modelled, $M_{\rm h,200}$,
within the limits of the fiducial prior choice $M_{\rm h,200}:[5\times
10^{12},7\times 10^{13}]M_{\odot}$. The ACT CMASS measurements at 98GHz are
shown as the black data points in the bottom panels and the model profiles are
convolved with the f90 beam profile for comparison. Lower: The marginalised
posteriors for $S_{8}$, $\log_{10}(M_{\rm c})$ and $M_{\rm h,200}$ in a WL +
kSZ analysis with different prior choices on $M_{\rm h,200}$. We show the
prior used in the fiducial analysis $M_{\rm h,200}:[5\times 10^{12},7\times
10^{13}]M_{\odot}$ (purple), as well as a wide prior $M_{\rm h,200}:[8\times
10^{12},3\times 10^{14}]M_{\odot}$ (light blue) and fixed $M_{\rm
h,200}=3\times 10^{13}]M_{\odot}$ (dark blue, also shown as the dashed grey
line). The inner and outer contours show the 68% and 95% confidence levels
respectively.

###  B.4 The relationship between baryon fraction and matter power spectrum
suppression

It has been demonstrated that in hydrodynamical simulations, the mean baryon
fraction in halos of $M\sim 10^{14}\mbox{${\bf M_{\odot}}$}$ can be predictive
for the suppression of the matter power spectrum due to baryonic feedback
effects, robust to a number of feedback prescriptions (van Daalen et al.,
2020). This relationship can be described by an empirical fitting function
relating the mean baryon fraction, $f_{\rm b}/\Omega_{\rm b}/\Omega_{\rm m}$,
measured within $R_{500}$ for halos of mass $M_{500}=10^{14}M_{\odot}$ to the
suppression of the matter power spectrum $P(k)/P_{\mathrm{DMonly}}(k)$ at
$k=1$ $h\mathrm{Mpc}^{-1}$ (equation 5 in their work). The best-fit relation,
shown as the solid grey line in Fig. 15, was fit to the cosmo-OWLS (Le Brun et
al., 2014) and BAHAMAS (McCarthy et al., 2017) simulations and is accurate to
1% for the simulations they test, shown as the grey shaded region in Fig. 15.
In this section, we discuss where constraints of our analysis lie in the
$P(k)/P_{\mathrm{DMonly}}(k)$-$f_{\rm b}/\Omega_{\rm b}/\Omega_{\rm m}$ plane
with respect to the van Daalen et al. (2020) relation. In Fig. 15 we plot the
result for the WL-only and WL + kSZ analyses with BCEmu7 baryon modelling, as
well as the WL-only results with SP(k)’s wide and X-ray informed prior
choices. For comparison, we also show the simulations discussed throughout
this work: FLAMINGO (Schaye et al., 2023); BAHAMAS (McCarthy et al., 2017);
SIMBA (Davé et al., 2019), MillenniumTNG (Pakmor et al., 2023) and FABLE
(Henden et al. 2018, Bigwood et al. in prep.).

This van Daalen et al. (2020) relationship is at the core of the the SP(k)
model (Salcido et al., 2023). Given that the ANTILLES simulations follow this
relationship, the mean baryon fraction is used to calibrate their emulator.
Naturally, both of our WL-only analyses with SP(k) fall on the van Daalen et
al. (2020) relation to within 1%. However, the BCEmu model does not enforce
the van Daalen et al. (2020) relationship between the baryon fraction and
matter power suppression, and allows for greater flexibility in the impact of
feedback on the matter distribution, including scenarios which, according to
the simulations, could be deemed unphysical. Both the WL-only and WL + kSZ
analyses with BCEmu give a mean constraint below the van Daalen et al. (2020)
relation. That is, they allow for a more extreme matter power spectrum
suppression for their predicted mean baryon fraction predicted. In the future,
it will be important to test this relationship with a range observations,
including those probe lower halo masses (see also Pandey et al., 2023).

Figure 15: The relation between the matter power spectrum suppression
$P(k)/P_{\mathrm{DMonly}}(k)$ at $k=1$ $h\mathrm{Mpc}^{-1}$ and the mean
baryon fraction, $f_{\rm b}/\Omega_{\rm b}/\Omega_{\rm m}$, measured within
$R_{500}$ for halos of mass $M_{500}=10^{14}M_{\odot}$. We plot the empirical
best-fit relation of van Daalen et al. (2020) as the grey solid line, with the
1% accuracy on the relation’s ability to predict $P(k)/P_{\mathrm{DMonly}}(k)$
shown as the shaded grey region. We plot the constraints obtained in our
analysis in the $P(k)/P_{\mathrm{DMonly}}(k)$-$f_{\rm b}/\Omega_{\rm
b}/\Omega_{\rm m}$ plane. We show the WL-only (pink starred) and WL + kSZ
(pruple starred) analyses with BCEmu7, and the WL-only results with $SPk$,
showing both the wide (light green starred) and X-ray (dark green starred)
prior choices. We also plot the simulations FLAMINGO (Schaye et al., 2023);
BAHAMAS (McCarthy et al., 2017); SIMBA (Davé et al., 2019), MillenniumTNG
(Pakmor et al., 2023) and FABLE (Henden et al. 2018, Bigwood et al. in prep.)
as the black datapoints.

##  Appendix C Intrinsic alignment (IA) model

In this appendix we consider the impact of the IA model choice on our results.
Throughout this work, we use the NLA IA model, which is a sub-space of the
TATT model. A WL-only analysis using BCEmu7 with TATT constrains
$S_{8}=0.802^{+0.028}_{-0.024}$, lying $\sim 0.5\sigma$ lower than that
constrained using NLA. This shift in $S_{8}$ between the two IA models is
consistent when using BCEmu1 instead of BCEmu7, and similar to that found in
previous work when using HM20 DES & KiDS Collaborations (2023) and scale cuts
(Secco et al., 2022; Amon et al., 2023). As in the literature, both IA model
choices give comparable $\chi^{2}_{\rm red}$ values (see Tab.3). In future
work, it is important to determine the more accurate IA model. Here, we use
this test as validation that our baryon model choice and the results of our
baryon model comparison are independent of the choice of IA model.

##  Appendix D Sampling algorithm choice

In this appendix we compare the cosmological parameter estimates attained
using the multinest and polychord samplers. Throughout the work, we use the
multinest settings $n_{\rm live}=500$, efficiency = 0.3, tolerance = 0.1,
constant efficiency = False, max. iterations = 50,000 for computing
efficiency. For polychord we use $n_{\rm live}=500$, tolerance = 0.01, $n_{\rm
repeats}=60$, fast fraction$=0.1$. Considering the BCEmu7 WL-only analysis we
find that sampling with polychord estimates a mean value of
$S_{8}=0.817^{+0.025}_{0.025}$, consistent with that attained by multinest of
$S_{8}=0.818^{+0.017}_{0.024}$. However, in agreement with the findings of DES
& KiDS Collaborations (2023), we find that the 68% confidence level for
$S_{8}$ attained using multinest is $18\%$ smaller than that estimated with
polychord. Similarly we find a WL + kSZ analysis with polychord returns
$S_{8}=0.821^{+0.020}_{0.023}$, consistent with the
$S_{8}=0.823^{+0.019}_{0.020}$ using multinest but with a $9\%$ larger
confidence region.

## Affiliations

1 Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge
CB3 0HA, U 2 Kavli Institute for Cosmology Cambridge, Madingley Road,
Cambridge, CB3 OHA 3 Department of Astrophysical Sciences, Princeton
University, Peyton Hall, Princeton, NJ 08544, USA 4 Institute for
Computational Science, University of Zurich, Winterthurerstrasse 190, CH-8057
Zurich, Switzerland 5 Astrophysics Research Institute, Liverpool John Moores
University, 146 Brownlow Hill, Liverpool L3 5RF, UK 6 Centre for Energy,
Environmental and Technological Research, Av. Complutense, 40, Moncloa -
Aravaca, 28040 Madrid, Spain 7 Kavli Institute for Particle Astrophysics and
Cosmology, 382 Via Pueblo Mall Stanford, CA 94305-4060, USA 8 SLAC National
Accelerator Laboratory 2575 Sand Hill Road Menlo Park, California 94025, USA 9
Physics Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720,
USA 10 Berkeley Center for Cosmological Physics, Department of Physics,
University of California, Berkeley, CA 94720, USA 11 Department of Astronomy,
Cornell University, Ithaca, NY 14853, USA 12 Department of Physics, University
of Michigan, Ann Arbor, MI 48109, USA 13 Kavli Institute for the Physics and
Mathematics of the Universe (WPI), UTIAS, The University of Tokyo, Kashiwa,
Chiba 277-8583, Japan 14 Department of Physics, Carnegie Mellon University,
Pittsburgh, Pennsylvania 15312, USA 15 NSF AI Planning Institute for Physics
of the Future, Carnegie Mellon University, Pittsburgh, PA 15213, USA 16 Kavli
Institute for Particle Astrophysics & Cosmology, P. O. Box 2450, Stanford
University, Stanford, CA 94305, USA 17 SLAC National Accelerator Laboratory,
Menlo Park, CA 94025, USA 18 Laboratório Interinstitucional de e-Astronomia -
LIneA, Rua Gal. José Cristino 77, Rio de Janeiro, RJ - 20921-400, Brazil 19
Observatório Nacional, Rua Gal. José Cristino 77, Rio de Janeiro, RJ -
20921-400, Brazil 20 Argonne National Laboratory, 9700 South Cass Avenue,
Lemont, IL 60439, USA 21 Institute of Space Sciences (ICE, CSIC), Campus UAB,
Carrer de Can Magrans, s/n, 08193 Barcelona, Spain 22 Department of Astronomy
and Astrophysics, University of Chicago, Chicago, IL 60637, USA 23 Fermi
National Accelerator Laboratory, P. O. Box 500, Batavia, IL 60510, USA 24
Kavli Institute for Cosmological Physics, University of Chicago, Chicago, IL
60637, USA 25 NASA Goddard Space Flight Center, 8800 Greenbelt Rd, Greenbelt,
MD 20771, USA 26 Instituto de Física Gleb Wataghin, Universidade Estadual de
Campinas, 13083-859, Campinas, SP, Brazil 27 Center for Cosmology and Astro-
Particle Physics, The Ohio State University, Columbus, OH 43210, USA 28
Instituto de Astrofisica de Canarias, E-38205 La Laguna, Tenerife, Spain 29
Department of Physics and Astronomy, University of Pennsylvania, Philadelphia,
PA 19104, USA 30 Université Grenoble Alpes, CNRS, LPSC-IN2P3, 38000 Grenoble,
France 31 University Observatory, Faculty of Physics, Ludwig-Maximilians-
Universität, Scheinerstr. 1, 81679 Munich, Germany 32 Jet Propulsion
Laboratory, California Institute of Technology, 4800 Oak Grove Dr., Pasadena,
CA 91109, USA 33 Brookhaven National Laboratory, Bldg 510, Upton, NY 11973,
USA 34 Department of Physics, ETH Zurich, Wolfgang-Pauli-Strasse 16, CH-8093
Zurich, Switzerland 35 Institut de Física d’Altes Energies (IFAE), The
Barcelona Institute of Science and Technology, Campus UAB, 08193 Bellaterra
(Barcelona) Spain 36 Department of Astronomy/Steward Observatory, University
of Arizona, 933 North Cherry Avenue, Tucson, AZ 85721-0065, USA 37 Department
of Physics, University of Arizona, Tucson, AZ 85721, USA 38 School of Physics
and Astronomy, Cardiff University, CF24 3AA, UK 39 Centro de Investigaciones
Energéticas, Medioambientales y Tecnológicas (CIEMAT), Madrid, Spain 40
Institut de Recherche en Astrophysique et Planétologie (IRAP), Université de
Toulouse, CNRS, UPS, CNES, 14 Av. Edouard Belin, 31400 Toulouse, France 41
Department of Physics and Astronomy, University of Waterloo, 200 University
Ave W, Waterloo, ON N2L 3G1, Canada 42 Institute for Astronomy, University of
Edinburgh, Edinburgh EH9 3HJ, UK 43 Department of Physics, Northeastern
University, Boston, MA 02115, USA 44 Lawrence Berkeley National Laboratory, 1
Cyclotron Road, Berkeley, CA 94720, USA 45 Jodrell Bank Center for
Astrophysics, School of Physics and Astronomy, University of Manchester,
Oxford Road, Manchester, M13 9PL, UK 46 Nordita, KTH Royal Institute of
Technology and Stockholm University, Hannes Alfvéns väg 12, SE-10691
Stockholm, Sweden 47 Physics Department, 2320 Chamberlin Hall, University of
Wisconsin-Madison, 1150 University Avenue Madison, WI 53706-1390 48 Department
of Physics, University of Genova and INFN, Via Dodecaneso 33, 16146, Genova,
Italy 49 Center for Astrophysical Surveys, National Center for Supercomputing
Applications, 1205 West Clark St., Urbana, IL 61801, USA 50 Department of
Astronomy, University of Illinois at Urbana-Champaign, 1002 W. Green Street,
Urbana, IL 61801, USA 51 Department of Physics, Duke University Durham, NC
27708, USA 52 Department of Applied Mathematics and Theoretical Physics,
University of Cambridge, Cambridge CB3 0WA, UK 53 Kavli Institute for
Cosmology, University of Cambridge, Madingley Road, Cambridge CB3 0HA, UK 54
Department of Physics, Stanford University, 382 Via Pueblo Mall, Stanford, CA
94305, USA 55 Physics Department, William Jewell College, Liberty, MO, 64068
56 Department of Physics and Astronomy, Stony Brook University, Stony Brook,
NY 11794, USA 57 Department of Astronomy, University of Geneva, ch. d’Écogia
16, CH-1290 Versoix, Switzerland 58 Department of Astronomy, University of
California, Berkeley, 501 Campbell Hall, Berkeley, CA 94720, USA 59 Cerro
Tololo Inter-American Observatory, NSFs National Optical-Infrared Astronomy
Research Laboratory, Casilla 603, La Serena, Chile 60 Institute of Cosmology
and Gravitation, University of Portsmouth, Portsmouth, PO1 3FX, UK 61 CNRS,
UMR 7095, Institut d’Astrophysique de Paris, F-75014, Paris, France 62
Sorbonne Universités, UPMC Univ Paris 06, UMR 7095, Institut d’Astrophysique
de Paris, F-75014, Paris, France 63 Department of Physics & Astronomy,
University College London, Gower Street, London, WC1E 6BT, UK 64 Institut
d’Estudis Espacials de Catalunya (IEEC), 08034 Barcelona, Spain 65 Hamburger
Sternwarte, Universität Hamburg, Gojenbergsweg 112, 21029 Hamburg, Germany 66
Department of Physics, IIT Hyderabad, Kandi, Telangana 502285, India 67
Institute of Theoretical Astrophysics, University of Oslo. P.O. Box 1029
Blindern, NO-0315 Oslo, Norway 68 Instituto de Fisica Teorica UAM/CSIC,
Universidad Autonoma de Madrid, 28049 Madrid, Spain 69 School of Mathematics
and Physics, University of Queensland, Brisbane, QLD 4072, Australia 70 Santa
Cruz Institute for Particle Physics, Santa Cruz, CA 95064, USA 71 Department
of Physics, The Ohio State University, Columbus, OH 43210, USA 72 Center for
Astrophysics $|$ Harvard & Smithsonian, 60 Garden Street, Cambridge, MA 02138,
USA 73 Australian Astronomical Optics, Macquarie University, North Ryde, NSW
2113, Australia 74 Lowell Observatory, 1400 Mars Hill Rd, Flagstaff, AZ 86001,
USA 75 George P. and Cynthia Woods Mitchell Institute for Fundamental Physics
and Astronomy, and Department of Physics and Astronomy, Texas A&M University,
College Station, TX 77843, USA 76 LPSC Grenoble - 53, Avenue des Martyrs 38026
Grenoble, France 77 Institució Catalana de Recerca i Estudis Avançats, E-08010
Barcelona, Spain 78 Perimeter Institute for Theoretical Physics, 31 Caroline
St. North, Waterloo, ON N2L 2Y5, Canada 79 Ruhr University Bochum, Faculty of
Physics and Astronomy, Astronomical Institute, German Centre for Cosmological
Lensing, 44780 Bochum, Germany 80 Department of Physics and Astronomy,
Pevensey Building, University of Sussex, Brighton, BN1 9QH, UK 81 School of
Physics and Astronomy, University of Southampton, Southampton, SO17 1BJ, UK 82
83 Computer Science and Mathematics Division, Oak Ridge National Laboratory,
Oak Ridge, TN 37831 84 Max Planck Institute for Extraterrestrial Physics,
Giessenbachstrasse, 85748 Garching, Germany 85 Universitäts-Sternwarte,
Fakultät für Physik, Ludwig-Maximilians Universität München, Scheinerstr. 1,
81679 München, Germany