<EMAIL_ADDRESS>
# Learning Cognitive Maps from Transformer Representations for Efficient
Planning in Partially Observed Environments

Antoine Dedieu  Google DeepMind  Wolfgang Lehrach  Google DeepMind  Guangyao
Zhou  Google DeepMind  Dileep George  Google DeepMind  Miguel Lázaro-Gredilla
Google DeepMind

###### Abstract

Despite their stellar performance on a wide range of tasks, including in-
context tasks only revealed during inference, vanilla transformers and
variants trained for next-token predictions (a) do not learn an explicit world
model of their environment which can be flexibly queried and (b) cannot be
used for planning or navigation. In this paper, we consider partially observed
environments (POEs), where an agent receives perceptually aliased observations
as it navigates, which makes path planning hard. We introduce a transformer
with (multiple) discrete bottleneck(s), TDB, whose latent codes learn a
compressed representation of the history of observations and actions. After
training a TDB to predict the future observation(s) given the history, we
extract interpretable cognitive maps of the environment from its active
bottleneck(s) indices. These maps are then paired with an external solver to
solve (constrained) path planning problems. First, we show that a TDB trained
on POEs (a) retains the near-perfect predictive performance of a vanilla
transformer or an LSTM while (b) solving shortest path problems exponentially
faster. Second, a TDB extracts interpretable representations from text
datasets, while reaching higher in-context accuracy than vanilla sequence
models. Finally, in new POEs, a TDB (a) reaches near-perfect in-context
accuracy, (b) learns accurate in-context cognitive maps (c) solves in-context
path planning problems.

##  1 Introduction

Large vanilla transformers (Vaswani et al., 2017) trained for next-token
prediction have been successfully applied to a wide range of applications such
as natural language processing (Radford et al., 2018; Chowdhery et al., 2022),
text-conditioned image generation (Ramesh et al., 2021), reinforcement
learning (Chen et al., 2021a), code writing (Chen et al., 2021b). These large
language models (LLMs) also exhibit emergent abilities (Wei et al., 2022),
among which their capacity to in-context learn (ICL), i.e., to adapt to a new
task at inference time given a few examples (Brown et al., 2020). However,
these models suffer from shortcomings that prevent them from being used for
planning and navigation (Valmeekam et al., 2022; Mialon et al., 2023). In
particular, they do not learn an interpretable world model (Momennejad et al.,
2023) which can be flexibly queried.

Herein, we consider a suite of partially observed environments (POEs) Chrisman
(1992) where the agent receives perceptually aliased observations as it
navigates; and cannot deterministically recover its spatial positions from its
observations. Path planning in these POEs is hard: the planner has to
disambiguate aliasing to locate itself, which requires modeling its history of
observations and actions.

Figure 1: An agent is trained on random walks in an aliased room with no
reward and unknown layout. At test time, given a novel random walk (in
fuchsia), it has to find a shortest path (in red) between room positions A and
B. While a vanilla transformer solves this path planning problem with forward
rollouts, which can be exponentially expensive due to aliasing, our
transformer variant pairs its learned cognitive map with an external solver.

Example of a path planning problem in a POE that a vanilla transformer cannot
solve:  We consider the partially observed $\small{15\times 20}$ room in
Fig.1, which only contains four unique observations. The room also has global
aliasing: the $\small{4\times 4}$ patch with a black border appears twice. An
agent executes a series of discrete actions, each action leading to a discrete
observation. The agent does not have access to any reward or to the room
layout. At test time, it wants to find a shortest path (in red) between two
room positions (A and B) it has been to. A transformer trained to predict the
next observation given past observations and actions can only perform forward
rollouts, which, due to aliasing, (a) scale exponentially in the $\ell_{1}$
distance between A and B (b) prevents it from even knowing when it reaches the
target B.

In this paper, we propose the transformer with discrete bottleneck(s), TDB),
which adds a single, or multiple, discrete bottleneck(s) (Van Den Oord et al.,
2017) on top of a transformer to compress the transformer outputs into a
finite number of latent codes. We train TDB with an augmented objective, then
extract a cognitive map from its active bottleneck(s) indices. First, we show
that, on aliased rooms, aliased cubes with non-Euclidean dynamics and visually
rich $3$D environments (Beattie et al., 2016), these maps (a) disambiguate
aliasing (b) nearly recover the ground truth dynamics and (c) can be paired
with external solvers to solve path planning problems—like the one in Fig.1.
Hence, TDB (a) retains the nearly perfect predictive abilities of vanilla
transformers and LSTMs, (b) solves (constrained) path planning problems
exponentially faster. Second, a TDB extracts an interpretable latent graph
from a text dataset (Xie et al., 2021) while achieving higher test accuracies
than vanilla sequence models. Finally, when exploring a new POE at test time,
TDB can (a) in-context predict the next observation given history (b) solve
in-context path planning problems (c) learn accurate in-context cognitive
maps.

The rest of this paper is organized as follows. Sec.2 discusses related work.
Sec.3 details our proposed TDB model. Finally, Sec.4 compares our method with
vanilla transformer and LSTM in a variety of navigation and text experiments.

##  2 Related work

Planning with LLMs:  Despite their success on simple benchmarks including
maths (Cobbe et al., 2021) and logics (Srivastava et al., 2022), growing
evidence suggests that LLMs performance collapse on benchmarks that require
stronger planning skills (Valmeekam et al., 2022). While Pallagani et al.
(2022) finetune LLMs for planning, Guan et al. (2023) extract an approximate
world model with an LLM, then further refine it with an external planner.

Interpretable circuits:  Some works try to reverse-engineer LLMs to extract
interpretable circuits, i.e., internal structures that drive certain model
behaviors (Elhage et al., 2021; Olsson et al., 2022; Nanda et al., 2023).
While circuit analysis can be performed at scale (Lieberum et al., 2023), it
is labor intensive (Conmy et al., 2023), and does not extract explicit
structures to solve planning or navigation tasks.

Decision transformers:  Decision transformers (Chen et al., 2021a) abstract
reinforcement learning (RL) as sequence modeling: they can play Atari games
and stack blocks with a robot arm (Reed et al., 2022). However, they only have
been applied to shortest path problems on small graphs with $20$ nodes.
Decision transformers have also access to a reward, which reveals information
about the environment. They would struggle to solve the shortest path in
Fig.1, where aliasing is high and no external reward is provided.

Learning world models with a transformer:  Lamb et al. (2022) train a
transformer with a multi-step objective on sequences of observations and
actions, and successfully learned a minimal world representation that either
the agent can control or that affects its dynamics. However, the authors (a)
consider fully visible settings and (b) learn representations by discarding
the agent’s history: their approach would not solve the shortest path in
aliased settings as in Fig.1. In another approach, Guo et al. (2022) learn
future latent representations from current latent representations but do not
explicitly extract interpretable world models that can be used for planning.
In a different line of work, Li et al. (2022) show that transformers learn
internal world models, that can be used to control some model’s behavior. In
contrast, here, we are interested in extracting a world model from a
transformer that can be used to solve path planning problems exponentially
faster than vanilla models.

Clone-structured causal graphs (CSCGs):  CSCG (George et al., 2021) is an
interpretable probabilistic model that learns cognitive maps in aliased POEs
and solves path planning tasks. However, CSCGs learn a block-sparse transition
matrix over a large latent space via the expectation-maximization (EM)
algorithm (Dempster et al., 1977). This large transition matrix occupies a
large amount of memory and limits CSCGs’ scalability—see Appendix B.1 for a
detailed discussion. CSCGs would also struggle to solve the ICL experiments in
Sec.4.5 without an external algorithm—see Swaminathan et al. (2023). In
contrast, the transformer variant we introduce in this paper builds this same
transition matrix sparsely, via local counts. Appendix B.2 discusses a CSCG-
inspired variant of our proposed model.

##  3 A path planning-compatible transformer

###  3.1 Problem statement

Problem:  We consider an agent executing a series of discrete actions
$a_{1},\ldots,a_{N-1}$ with $a_{n}\in\\{1,\ldots,N_{\text{actions}}\\}$, e.g.
walking in a room. As a result of each action, the agent receives a
perceptually aliased observation Chrisman (1992), resulting in the stream
$x_{1},\ldots,x_{N}$. The agent does not have access to any reward at any
time. Our goal is to (a) predict the next observation given the history of
past observations and actions and (b) build a world model that disambiguates
aliased observations and that can be called by an external planner to solve
path planning tasks.

Trajectory representation:  We represent a trajectory
$\tau=(x_{1},a_{1},\ldots,a_{N-1},x_{N})$ by alternating observations and
actions. This representation allows us to consider the autoregressive
objective of predicting the next observation given the history of past
observations and actions:

$\mathcal{L}_{\operatorname*{obs}}=\sum_{n=1}^{N-1}\mathcal{L}_{\operatorname*{obs}}(n)=-\sum_{n=1}^{N-1}\log
p(x_{n+1}|x_{1},a_{1},\ldots,x_{n},a_{n})$ (1)

###  3.2 Predicting the next observation with a transformer

We propose to train a vanilla transformer (Vaswani et al., 2017) to minimize
Equation (1). To do so, we first map each categorical observation $x_{n}$
(resp. action $a_{n}$) of $\tau$ to a linear embedding
$E_{\operatorname*{obs}}(x_{n})\in\mathbb{R}^{D}$ (resp.
$E_{\operatorname*{act}}(a_{n})\in\mathbb{R}^{D}$). Second, we feed the
sequence of input embeddings
$z=\left(E_{\operatorname*{obs}}(x_{1}),E_{\operatorname*{act}}(a_{1}),\ldots,E_{\operatorname*{obs}}(x_{N}),E_{\operatorname*{act}}(a_{N})\right)$
to a transformer with causal mask (Radford et al., 2018), resulting in the
output sequence $(T_{1},\ldots,T_{2N})$. We only retain the outputs with even
indices $(T_{2},T_{4},\ldots,T_{2N})$: $T_{2n}\in\mathbb{R}^{D}$ is derived
after applying the successive self-attention layers to the intermediate
representations of $x_{1},a_{1},\ldots,x_{n},a_{n}$. Finally, a linear mapping
applied to $T_{2n}$ returns the predicted logits for the next observation
$x_{n+1}$. We display this transformer architecture in Fig.6, Appendix A.

As we show in Sec.4, this transformer almost perfectly predicts the next
observation on a suite of perceptually aliased POEs. However, it can only
solve the path planning problem in Fig.1 with forward rollouts, which have an
exponential cost in the $\ell_{1}$ distance between the room positions A and
B.

Figure 2: Our proposed transformer with a (single) discrete bottleneck. The
respective linear embeddings of observations and actions go through a causal
transformer. The observation outputs are compressed by the vector quantizer,
then concatenated with the next action embedding in order to predict the next
observation. Finally, a cognitive map of the environment is built from the
active bottleneck indices.

###  3.3 A transformer with a discrete bottleneck

To address these inherent limitations, we add a discrete bottleneck on top of
a vanilla transformer, which compresses all the information needed to minimize
Equation (1).

After applying the causal transformer from Sec.3.2 to the sequence $z$ of
observations and actions embeddings, we now only extract the encodings with
odd indices, which results in the stream $(T_{1},T_{3},\ldots,T_{2N-1})$:
$T_{2n-1}\in\mathbb{R}^{D}$ is derived from
$x_{1},a_{1},\ldots,a_{n-1},x_{n}$, but is not aware of the action $a_{n}$. We
then apply vector quantization (Van Den Oord et al., 2017) and map each vector
$T_{2n-1},\forall n$ to the closest entry in a dictionary of latent codes
$\mathcal{D}=(d_{1},\ldots,d_{K})$, where
$d_{k}\in\mathbb{R}^{D},\leavevmode\nobreak\ \forall k$. That is, we introduce
the operator

$\phi(y)=\operatorname*{argmin}_{k=1,\ldots,K}\|y-d_{k}\|_{2}^{2},\leavevmode\nobreak\
\forall y.$ (2)

The quantized output of $T_{2n-1}$ is $d_{\phi(T_{2n-1})}$. Because $T_{2n-1}$
has not been exposed to $a_{n}$, we derive the next observations logits via
$f\left(d_{\phi(T_{2n-1})}\oplus E_{\operatorname*{act}}(a_{n})\right)$, where
$f$ is a two layers MLP and $\oplus$ represents the concatenation operator. We
refer to our proposed transformer with a discrete bottleneck as TDB, and
visualize it in Figure 2.

###  3.4 Augmenting the objective value

The discrete bottleneck in Sec.3.3 introduces the additional objective term
$\mathcal{L}_{\operatorname*{bot}}=\sum_{n=1}^{N}\mathcal{L}_{\operatorname*{bot}}(n)$
with

$\mathcal{L}_{\operatorname*{bot}}(n)=\|d_{\phi(e)}-\operatorname*{sg}(e)\|_{2}^{2}+\beta\|\operatorname*{sg}\left(d_{\phi(e)}\right)-e\|_{2}^{2},$
(3)

where $e=T_{2n-1}$, $\beta=0.25$ and $\operatorname*{sg}$ the stop-gradient
operator (Van Den Oord et al., 2017). However, we show in Fig.7, Appendix G
that a TDB trained with the simple objective
$\mathcal{L}_{\operatorname*{obs}}+\mathcal{L}_{\operatorname*{bot}}$ fails at
disambiguating observations with identical neighbors. To allow the TDB to
learn richer representations of the observations, we propose to augment the
transformer objective via two solutions.

Solution 1: Predicting the next $S$ observations. At each timestep $n$,
instead of only predicting the next observation—which corresponds to $S=1$—we
predict the next $S$ observations and replace the loss
$\mathcal{L}_{\operatorname*{obs}}(n)$ in Equation (1) with

$\mathcal{L}^{S}_{\operatorname*{obs}}(n)=-\frac{1}{S}\sum_{s=0}^{S-1}\log
p(x_{n+s+1}|x_{1},a_{1},\ldots,x_{n},a_{n},\ldots,a_{n+s})$ (4)

The model does not have access to any observation after $x_{n}$: it only sees
the $S$ actions $a_{n},\ldots,a_{n+S-1}$. The logits of the observation
$x_{n+s+1}$ are predicted via $f^{(s)}\left(d_{\phi(T_{2n-1})}\oplus
E_{\operatorname*{act}}(a_{n})\oplus\ldots\oplus
E_{\operatorname*{act}}(a_{n+s})\right)$, where $f^{(s)}$ is a two-layers MLP.
In Sec.4, we set $S=3$.

Solution 2: Predicting the next encoding. Our second solution is to predict
future latent representations Guo et al. (2022). We refer to the TDB described
so far as a student network, with weights $\Theta$, and introduce a teacher
network (Grill et al., 2020), with the same architecture and whose weights
$\Theta^{\text{teacher}}$ are an exponential moving average of $\Theta$. At
each gradient update, we set
$\Theta^{\text{teacher}}\leftarrow\alpha\Theta^{\text{teacher}}+(1-\alpha)\Theta$,
where $\alpha$ is the decay rate, which we fix to $0.99$.

The student network predicts both (a) the next observation and (b) the next
teacher network’s quantized encoding. This introduces the objective term
$\mathcal{L}_{\operatorname*{enc}}=\sum_{n=1}^{N-1}\mathcal{L}_{\operatorname*{enc}}$
with

$\displaystyle\mathcal{L}_{\operatorname*{enc}}(n)=\left\|\underbrace{\frac{g\left(d_{\phi(T_{2n-1})}\oplus
E_{\operatorname*{act}}(a_{n})\right)}{\|g\left(d_{\phi(T_{2n-1})}\oplus
E_{\operatorname*{act}}(a_{n})\right)\|_{2}}}_{\text{using
}\Theta}-\operatorname*{sg}\left(\underbrace{\frac{d_{\phi(T_{2n+1})}}{\|d_{\phi(T_{2n+1})}\|_{2}}}_{\text{using
}\Theta^{\text{teacher}}}\right)\right\|_{2}^{2}$

where $g$ is a two-layer MLP with same hidden layer as $f$.

###  3.5 Extension to multiple discrete bottlenecks

TDB supports $\small{M>1}$ discrete bottlenecks. Each bottleneck induces its
own (a) dictionary of latent codes
$\mathcal{D}^{(i)}=(d^{(i)}_{1},\ldots,d^{(i)}_{K})$, (b) function
$\phi^{(i)}$ which acts as in Equation (2) and (c) objective
$\mathcal{L}_{\operatorname*{bot}}^{(i)}$. The transformer output $T_{2n-1}$
passes through the discrete bottlenecks in parallel: the $i$th bottleneck
returns the latent code $d^{(i)}_{\phi^{(i)}(T_{2n-1})}$. We get back to the
previous case by defining (a)
$d_{\phi(T_{2n-1})}=\oplus_{i=1}^{M}d^{(i)}_{\phi^{(i)}(T_{2n-1})}$ and (b)
$\mathcal{L}_{\operatorname*{bot}}=\sum_{i=1}^{M}\mathcal{L}_{\operatorname*{bot}}^{(i)}$.
As we show in Sec.4, multiple discrete bottlenecks make training faster and
sometimes better. However, Appendix F introduces a disentanglement metric
which shows that they are highly redundant and do not learn disentangled
representations.

###  3.6 Learning a cognitive map from bottleneck indices

When we use a single bottleneck, TDB maps each observation $x_{n}$ to the
latent code index $s_{n}=\phi(T_{2n-1})\in\\{1,\ldots,K\\}$. We define a count
tensor $C$ which counts the empirical latent transitions:
$C_{ijk}=\sum_{(s_{n},a_{n},s_{n+1})}\mathbf{1}(s_{n}=i,a_{n}=j,s_{n+1}=k)$.
We then threshold $\mathcal{C}$ to build an action-augmented empirical
transition graph $\mathcal{G}$. That is, $\mathcal{G}$ has an edge between the
vertices $i$ and $k$ with the action $j^{*}=\operatorname*{argmax}_{j}C_{ijk}$
iff. $C_{ij^{*}k}\geq t_{\text{ratio}}\leavevmode\nobreak\ c^{*}$, where
$c^{*}=\max_{ijk}C_{ijk}$ and $t_{\text{ratio}}$ is a threshold. This edge
indicates that the action $j^{*}$ leads from the latent node $i$ to the latent
node $k$ a large number of times.

In practice (a) we use multiple discrete bottlenecks (see Sec.3.5) and cluster
the bottleneck indices using the Hamming distance with a threshold
$d_{\text{Hamming}}$ and (b) we map each discarded node of $\mathcal{C}$ to a
retained node in $\mathcal{G}$ so that we can solve all the path planning
problems in Sec.4. We detail (a) and (b) in Appendix C. As a result, each node
in $\mathcal{G}$ is paired with a collection of tuples of bottleneck indices.

$\mathcal{G}$ is a cognitive map of the agent’s environment which (a) is
interpretable, (b) is action-conditioned, hence, models the agent’s dynamics,
and (c) can be paired with an external planner to solve path planning
problems—see Sec.4.

###  3.7 Why does TDB learn an accurate map?

The TDB latent index $\phi(T_{2n-1})$ associated with an observation $x_{n}$
given its history corresponds to a node in the latent graph $\mathcal{G}$ in
Sec.3.6. For $\mathcal{G}$ to correctly model the environment’s dynamics, when
this node is active, the agent must always be at the same ground truth spatial
position.

$d_{\phi(T_{2n-1})}$ does not know the next action $a_{n}$; and compresses all
the information needed to minimize the loss at the $n$th step. When TDB
predicts the next observation, $d_{\phi(T_{2n-1})}$ must then encode the next
observation that each next possible action leads to. As we show in Appendix G,
$d_{\phi(T_{2n-1})}$ may not disambiguate distinct spatial positions with
identical neighbors. When TDB predicts the next $S$ observations (Solution 1,
Sec.3.4), $d_{\phi(T_{2n-1})}$ must now encode all the observations
$x_{n+s},1\leq s\leq S,$ that each sequence of $S$ actions leads to. For a
large $S$, the latent index $\phi(T_{2n-1})$ is encouraged to only be active
at a unique spatial location. An alternative approach is to augment
$d_{\phi(T_{2n-1})}$ so that it compresses all its neighbors’ encoding
(Solution 2).

##  4 Computational results

We assess the performance of our proposed TDB on four experiments: (a)
synthetic $2$D aliased rooms and aliased cubes, (b) simulated $3$D
environments, (c) text datasets and (d) mixtures of $2$D aliased rooms. Each
experiment is run on a $2\times 2$ grid of Tensor Processing Unit v2.

###  4.1 Methods compared

Each experiment compares the following methods:

$\bullet$ The causal transformer described in Sec.3.2, trained with the
autoregressive objective in Equation (1). The architecture considered has $4$
layers, $8$ attention heads, a context length of $400$, an embedding dimension
of $256$, and an MLP hidden layer dimension of $512$. We use relative
positional embedding (Dai et al., 2019) and Gaussian error linear units
activation functions (Hendrycks and Gimpel, 2016).

$\bullet$ Several TDBs, using the same causal transformer. We refer to each
model as TDB$(S,M)$ where (a) $S\in\\{1,3\\}$ is the number of prediction
steps (Solution 1, Sec.3.4), (b) $M\in\\{1,4\\}$ is the number of discrete
bottlenecks (see Sec.3.5)—each one contains $K=1000$ latent codes. We note
TDB$(S,\operatorname*{enc},M)$ when we predict the next encoding (Solution 2,
Sec.3.4). After training, we build a cognitive map as in Sec.3.6 using
$d_{\text{Hamming}}=0.25,t_{\text{ratio}}=0.1$.

$\bullet$ A vanilla LSTM (Hochreiter and Schmidhuber, 1997) with hidden
dimension of 256.

###  4.2 Random walks in 2D aliased rooms

Problem:  We consider a $2$D aliased ground truth (GT) room of size $15\times
20$, containing $\small{O=4}$ unique observations, and a $4\times 4$ patch
repeated twice—similar to Fig.1. An agent walking selects, at each timestep, a
discrete action $a_{n}\in\\{1,2,3,4\\}$ and collects the categorical
observation $x_{n}\in\\{1,\ldots,O\\}$. The actions have unknown semantics:
the agent does not know that $a_{n}=1$ corresponds to moving up. No
assumptions about Euclidean geometry are made either. The training and test
set contains $2048$ random walks of observations and actions of length $400$
each.

Training:  We train each method in Sec.4.1 with Adam (Kingma and Ba, 2014) for
$25$k training iterations, using a learning rate of $0.001$ and a batch size
of $32$. For regularization, we use a dropout of $0.1$.

Solving path planning problems at test time:  For a test sequence
$x^{\text{test}}$ of length $N=400$, let $\operatorname*{pos}_{n}$ be the GT
unknown $2$D spatial position of the observation $x_{n}$. At test time, given
a context $C=50$, the path planning task is to derive a shortest path, i.e. a
sequence of actions, which leads from $\operatorname*{pos}_{C}$ to
$\operatorname*{pos}_{N-C}$ in the GT room. We emphasize that (a) the model
only observes $x_{n}$ and does not know $\operatorname*{pos}_{n}$ (b) the path
$(a_{C},\ldots,a_{N-C-1})$ is a solution, referred to as fallback path. Fig.1
shows this problem for $N=40,C=5$.

For a vanilla sequence model, we autoregressively sample $100$ random rollouts
of length $N-2C$ each and collect all the candidates, i.e., the observations
equal to $x^{\text{test}}_{N-C}$. Because of aliasing, the model cannot know
whether a candidate is at the target position $\operatorname*{pos}_{N-C}$. We
use the tail of the last $C$ observations and actions and only return the
candidates estimated to be at $\operatorname*{pos}_{N-C}$. For a TDB, after
learning the cognitive map $\mathcal{G}$, we map $x_{C}$ and $x_{N-C}$ to the
bottleneck indices $\phi(T_{2C-1})$ and $\phi(T_{2(N-C)-1})$, then to two
nodes of $\mathcal{G}$ using the Hamming distance. We then call the external
solver networkx.shortest_path (Hagberg et al., 2008) to find the shortest path
between these two nodes. See Appendix D for details about both procedures.

#### Metrics:

We evaluate each method for four metrics.

$\bullet$ A prediction metric, TestAccu: next observation accuracy on the test
random walks.

$\bullet$ Two path planning metrics: we consider $200$ shortest paths problems
and report (a) ImpFallback, the percentage of problems for which a model finds
a valid path which improves over the fallback path—a path (i.e. a sequence of
actions) is valid if it correctly leads from $\operatorname*{pos}_{C}$ to
$\operatorname*{pos}_{N-C}$ in the GT room—and (b) RatioSP, when a better
valid path is found, the ratio between its length and an optimal path
length—we derive an optimal path using the GT room.

$\bullet$ Normalized graph edit distance, defined for two graphs
$\mathcal{G}_{1},\mathcal{G}_{2}$ as
$\text{NormGED}(\mathcal{G}_{1},\mathcal{G}_{2})=\text{GED}(\mathcal{G}_{1},\mathcal{G}_{2})\leavevmode\nobreak\
\leavevmode\nobreak\ /\leavevmode\nobreak\ \leavevmode\nobreak\
(\text{GED}(\mathcal{G}_{1},\emptyset)+\text{GED}(\mathcal{G}_{2},\emptyset))$
where GED is the graph edit distance (Sanfeliu and Fu, 1983) and $\emptyset$
is the empty graph. NormGED is lower than $1.0$ and is equal to $0.0$ iff.
$\mathcal{G}_{1}$ and $\mathcal{G}_{2}$ are isomorphic. We use a timeout of
$15$ minutes and empirical positions to accelerate the GED computation—see
Appendix H.

Method |  TestAccu ($\%$) $\uparrow$ |  ImpFallback ($\%$) $\uparrow$ |  RatioSP $\downarrow$ |  NormGED $\downarrow$
---|---|---|---|---
Vanilla transformer |  $99.00\leavevmode\nobreak\ (0.01)$ |  $29.53\leavevmode\nobreak\ (0.75)$ |  $16.82\leavevmode\nobreak\ (0.92)$ |  —
Vanilla LSTM |  $98.85\leavevmode\nobreak\ (0.01)$ |  $31.97\leavevmode\nobreak\ (0.66)$ |  $16.81\leavevmode\nobreak\ (0.83)$ |  —
TDB$(S=1,M=1)$ |  $98.94\leavevmode\nobreak\ (0.05)$ |  $6.61\leavevmode\nobreak\ (0.53)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.532\leavevmode\nobreak\ (0.005)$
TDB$(S=1,\operatorname*{enc},M=1)$ |  $98.74\leavevmode\nobreak\ (0.04)$ |  $99.20\leavevmode\nobreak\ (0.20)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $\mathbf{0.118}\leavevmode\nobreak\ (0.022)$
TDB$(S=3,M=1)$ |  $\mathbf{99.06}\leavevmode\nobreak\ (0.01)$ |  $\mathbf{99.55}\leavevmode\nobreak\ (0.13)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.124\leavevmode\nobreak\ (0.019)$
TDB$(S=3,M=4)$ |  $99.00\leavevmode\nobreak\ (0.01)$ |  $96.59\leavevmode\nobreak\ (0.15)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.186\leavevmode\nobreak\ (0.018)$

Table 1: Results averaged over $10$ aliased $15\times 20$ rooms with $O=4$
observations. See main text for a description of the metrics. Arrows pointing
up (down) indicate that higher (lower) is better. Our TDB with either multi-
step objective or next encoding prediction (a) retains the nearly perfect test
accuracy of vanilla sequence models (b) consistently solves the shortest paths
problems when paired with an external solver—while both transformer or LSTM
catastrophically fail—(c) learns cognitive maps nearly isomorphic to the
ground truth.

Results:  Table 1 averages the metrics over $10$ experiments: each run
considers a different GT room. First, we observe that both vanilla LSTM and
transformer almost perfectly predict the next observation given the history.
However, both sequence models struggle to solve path planning problems with
forward rollouts: they only improve over the fallback path $\sim 30\%$ of the
same. When they do so, they find shorter paths $16$ times longer than an
optimal path1.

Second, all the TDBs retain the predictive performance of vanilla models.
However, TDB$\small{(S=1,M=1)}$ fails at path planning: it only finds better
valid paths $7\%$ of the time as it only solves the simplest problems. Indeed,
the learned cognitive maps merge the observations with identical neighbors,
which introduces unrealistic shortcuts in the latent graphs—see Fig.7,
Appendix G. In contrast, TDB with either (a) multi-step or (b) next encoding
objective learns cognitive maps that recover the GT dynamics and reach low
NormGED—see Fig.3[center left]. They are able to almost consistently find an
optimal shortest path. As a result, ImpFallback is above $99\%$ while RatioSP
is $1.00$.

Multiple discrete bottlenecks accelerate training:  Table 2, Appendix E,
reports the average number of training steps to reach a $98\%$ train accuracy:
for $\small{S=3}$, this number goes from $8300$ for $\small{M=1}$ down to
$4000$ for $\small{M=4}$. In addition, multiple discrete bottlenecks seem more
robust to the environment. Table 6, Appendix I.2, shows that, when the room
contains $\small{O=12}$ unique observations, all the TDBs with $\small{M=4}$
outperform their $\small{M=1}$ counterparts.

Lack of disentanglement in the latent space:  We test whether multiple
discrete bottlenecks learn a disentangled representation by training a
logistic regression to predict each bottleneck index given the other
indices—see Appendix F for details. If the bottlenecks were independent, test
disentanglement accuracy would be $0.20\%$. However, it is $84.99\%(\pm
0.56\%)$, which shows the strong redundancy in the representations learned by
each bottleneck.

The agent can locate itself after a short context:  Appendix I.3 first shows
that, on test random walks, accuracy is initially low then quickly increases:
it is above $99.50\%$ after the $30$th observation—see Fig.9. Second, latent
codes have a bimodal distribution: many codes with low empirical frequencies
model the agent’s early uncertainty and are replaced by high-frequency latent
codes which later express the agent’s confidence—see Fig.10. Finally, the
agent can “teleportate” in the room but such behavior appears less than $1\%$
of the time after the $60$th observation—see Fig.11.

TDB can solve constrained path planning problems:  TDB can inject constraints
on demand into the external solver. To illustrate this, we consider a path
planning problem variant: the new task is to find the shortest path from
$\operatorname*{pos}_{C}$ to $\operatorname*{pos}_{N-C}$ which avoids, when
possible, a randomly picked color. TDB$\small{(S=3,M=1)}$ is still able to
almost consistently solve the problem: ImpFallback is still $99.55\%(\pm
0.13\%)$ while RatioSP is still $1.00(\pm 0.00)$. The average length of the
shortest paths is now $13.13(\pm 0.23)$, as opposed to $10.73(\pm 0.08)$ when
all the colors are allowed.

TDB handles noise during training:  When replacing $10\%$ of the training
observations of a TDB$\small{(S=3,M=1)}$ by random ones, training accuracy
drops to $88.03\%(\pm 0.04\%)$, while TestAccu, ImpFallback and RatioSP are
still $\small{99.08\%(\pm 0.01\%),99.49\%(\pm 0.20\%)}$ and $\small{1.00\%(\pm
0.00\%)}$, showing that training noise does not affect the test prediction and
planning performance.

TDB learn maps for non-Euclidean dynamics:  TDB can be applied to aliased
environments with non-Euclidean dynamics. We generate data from a $3$D aliased
cube with edge size $6$ and $O=12$ unique observations—see Fig.3[bottom left].
Table 5, Appendix I.1 shows that the prediction and planning performance of a
TDB are near perfect. Fig.3[bottom - center left] visualizes the learned graph
with the Kamada-Kawai algorithm (Kamada et al., 1989),

---

---

---

---

Figure 3: [Left] Top: $2$D aliased room of size $\small{15\times 20}$ with
$O=4$ unique observations and $2$ identical $4\times 4$ patches (in fuchsia).
Bottom: Aliased cube of edge size $6$ with $O=12$ and non-Euclidean dynamics.
[Center left] Top: cognitive map learned with a TDB($S=3,M=1$). For
visualization, each latent node is mapped with the observation (resp. is
placed at the $2$D GT spatial position) with higher empirical frequency when
this node is active. Bottom: similar, but we use the Kamada-Kawai algorithm.
[Center] Isometric view of a simulated $3$D environment. The agent navigates
with egocentric actions and collects RGB images. [Center right] Three cluster
centers: the cluster indices serve as categorical observations. [Right]
Cognitive map learned with a TDB($S=3,M=4$): each location is represented by
four nodes in the latent graph, corresponding to the four agent’s heading
directions. Also see Fig.12, Appendix J.

###  4.3 Egocentric views in 3D simulated environments

Problem:  We consider a suite of visually rich $3$D simulated environments
(Beattie et al., 2016) with a similar setting as Guntupalli et al. (2023)—see
Fig.3[center, center right]. At each step, the agent can take any of three
discrete egocentric actions—move $1$m forward, rotate $90\deg$ left, rotate
$90\deg$ right—and it sees a new view, which is a $64\times 64$ RGB image. The
agent’s egocentric views are clustered with a $k$-means quantizer—with
$\small{k=128}$— and the clusters indices are used as categorical
observations. The training and test set sizes are identical to Sec.4.2.

Training:  We train the same models as in Sec.4.2, with the same parameters,
on $10$ synthetic $3$D rooms.

Results:  The full results are reported in Table 7, Appendix J. Despite their
excellent predictive performance, vanilla sequence models only find a path
better than fallback for $26.60\%(\pm 0.60\%)$ of the problems: these “better”
paths are $16.38(\pm 0.87)$ times longer than optimal paths111The transformer
and LSTM performance are almost identical because the same random walks are
used for both models.. Compared with Sec.4.2, training TDBs with a single
bottleneck do not converge. Consequently, these models reach a low test
accuracy and cannot solve the path planning problems. In contrast, for TDBs
with $\small{M=4}$ discrete bottlenecks, averaged train accuracies are at
least $99.70\%$ after $1000$ only training steps (see Table 3, Appendix E).
These models reach nearly perfect (a) test accuracy and (b) path planning
performance. Their latent maps accurately model the POE’s dynamics—see
Fig.3[right] and Fig.12, Appendix J. However, the multiple bottlenecks learn a
highly redundant representation: the test disentanglement accuracy (see
Appendix F) of a TDB$(S=3,M=4)$ is $99.06\%\leavevmode\nobreak\ (\pm 0.06\%)$.

###  4.4 Extension to text via the GINC dataset

Problem: Herein, we extend TDB to text datasets to show that it can extract
interpretable latent graphs on this other modality. We consider the text GINC
dataset, introduced in Xie et al. (2021) to study in-context learning (ICL).
GINC222available at https://github.com/p-lambda/incontext-
learning/tree/main/data. is generated from a uniform mixture of five factorial
HMMs Ghahramani and Jordan (1995)—referred to as a concepts. Each concept has
two factorial chains with $10$ states each. Each state can emit $50$
observations shared across concepts. The training set consists of $1000$
training documents with a total of ${\sim}10$ million tokens: each document
uniformly selects a concept and samples independent sentences from it. The
test set consists of in-context prompts. Each prompt has between $n=0$ and
$n=64$ examples: each example is of length $k\in\\{3,5,8,10\\}$. There are
$2500$ prompts for each setting $(k,n)$. Each prompt uniformly selects a
concept, samples $n-1$ examples $x^{(1)}_{:k},\ldots,x^{(n-1)}_{:k}$ of length
$k$, and one example $x^{(n)}_{:k-1}$ of length $k-1$. The in-context task is
to infer the most likely last token of the last example, i.e.,
$\operatorname*{argmax}_{x_{k-1}^{(n)}}p\left(x_{k-1}^{(n)}|x^{(1)}_{:k},\ldots,x^{(n-1)}_{:k},x^{(n)}_{:k-1}\right)$.

Since the vocabulary is shared among the concepts, observations in GINC are
aliased like in natural language. Solving the task requires the model to
disambiguate the aliased observations and correctly infer the latent concepts.

Training: We train the same models as above using the same parameters, except
the batch size which we set to $24$.

Figure 4: [Left] For all context lengths $k$, TDB($S=1,M=4$) achieves higher
in-context accuracies than an LSTM and a vanilla transformer on the GINC test
dataset (Xie et al., 2021)—while TDB($S=5,M=1$) is the best model for large
contexts. [Right] The learned latent graph of TDB($S=5,M=1$) exhibits five
clusters, each corresponding to a color-coded concept.

Results:  Fig. 4[left] reports the in-context accuracy—defined as the average
ratio of correct predictions—for each pair $(k,n)$ of the GINC test set.
$\texttt{TDB}(S=1,M=1)$ is outperformed by vanilla sequence models. However,
$\texttt{TDB}(S=1,M=4)$ outperforms both transformer and LSTM for all context
lengths $k$. For larger contexts $k\in\\{8,10\\}$, the best model is
$\texttt{TDB}(S=5,M=1)$—it is however slower to converge (see Table 4,
Appendix E). Finally, the highest in-context accuracies for TDB are around
$95\%$ and are higher than what was reported in Xie et al. (2021) for
transformers. The numerical values are in Table 8, Appendix K.

Fig. 4[right] displays the latent graph learned by TDB$(S=5,M=1)$. We use
$t_{\text{ratio}}=0.001$ and the Kamada-Kawai algorithm. Each latent node has
a color corresponding to the GT concept with higher empirical frequency when
this node is active. The learned latent graph has an interpretable structure:
TDB learns five latent subgraphs corresponding to the five concepts in the
GINC dataset.

---

---

Figure 5: [Left] In novel $2$D aliased test rooms, TDB$(S=3,M=4)$ perfectly
(a) predicts the next observation (b) solves in-context path planning
problems. These in-context capacities emerge when the number of training rooms
increases. A vanilla transformer only solves the prediction problem (a)—which
an LSTM struggles to do. [Right] Two latent graphs in-context learned by the
TDB on two new test rooms. By design (a) a $3\times 3$ fuchsia-coded patch is
repeated twice and (b) the room partitions induced by the colors are the same.

###  4.5 In-context learning experiments

Problem:  Our last experiment explores whether vanilla models and TDB can, in
a new test POE, (a) in-context predict the next observation (b) solve in-
context path planning problems (c) learn in-context latent graphs. Similar to
Sec.4.2, we first generate a $2$D aliased room of size $6\times 8$ with $O=4$
unique colors, and with global aliasing: a $3\times 3$ patch is repeated
twice. We define the room partition as the partition
$\mathcal{S}_{1},\ldots,\mathcal{S}_{O}$ of the GT room spatial positions such
that all the room spatial positions in a set $\mathcal{S}_{i}$ have the same
observation. We then generate many aliased rooms of the same size by
preserving the room partition: each room picks $4$ colors out of $30$ without
replacement, and assigns all the positions in $\mathcal{S}_{i}$ to the same
color. See Appendix L.1.

We train each model for a varying number of training rooms:
$200,500,1\text{k},2\text{k},5\text{k},10\text{k}$ 333That is, we train on at
most $1.52\%$ of all the possible rooms.. The training set contains $8092$
sequences: each sequence picks a training room at random, then generates a
random walk of length $400$ in it. The test set only contains four random
walks in each new test room. We define in-context path planning problems as
before, using a context $C=100$. Given a new test room, TDB derives the test
bottleneck indices, builds an in-context cognitive map, and uses it to solve
the in-context path planning problems.

Training in base:  We map a sequence of observations
$x=(x_{1},x_{2},\ldots,x_{N})$ in a room—where $x_{n}$ is one of the four room
observations—to a room-agnostic sequence
$\tilde{x}=(\tilde{x}_{1},\ldots,\tilde{x}_{N})$ where (a)
$\tilde{x}_{n}\in\\{1,\ldots,O\\}$ (b) $\tilde{x}_{n}=k$ iff. $\tilde{x}_{n}$
is the $k$th lowest observation seen between indices $1$ and
$n$444$\tilde{x}_{n}=1$ may map to different $x_{n}$ through the sequence.. We
train each model to take $x$ as input and to predict $\tilde{x}$. Targets are
always in the base $\\{1,\ldots,O\\}$, which encourages the models to share
structure across rooms—e.g. TDB can reuse the same latent codes across rooms.
As before, we train for $25$k Adam iterations, using a learning rate of
$0.001$, a batch size of $32$ and a dropout of $0.1$.

Results:  Fig. 5[left] averages the results over $10$ runs—the numerical
values are in Table 9, Appendix L.2. For NormGED, we use a timeout of $20s$.
As the number of training rooms increases, all the models display a phase
transition where both prediction and path planning performance improve. When
the number of training rooms is larger than $5000$, in-context learning
emerges for the transformer and TDB: both achieve almost perfect in-context
accuracy on the new test rooms. In contrast, LSTM reaches lower in-context
accuracy. In addition, in-context path planning emerges in the TDB: it can
nearly perfectly solve the shortest paths problems, which both vanilla models
struggle with. The in-context latent graphs reach low NormGED and are nearly
isomorphic to the GT, and correctly model the room’s dynamics—see Fig.
5[right]. TDB$(S=3,M=4)$ trained on $10$k rooms has a disentanglement accuracy
(see Appendix F) of $94.02\%(\pm 0.27\%)$: again, multiple discrete
bottlenecks do not learn a disentangled representation. Finally, Appendix L.4
studies how spatial exposure to base targets in the training data drives in-
context performance.

##  5 Conclusion

We propose TDB, a transformer variant that addresses the shortcomings of
vanilla transformers for path planning. TDB (a) introduces discrete
bottlenecks which compress the information necessary to predict the next
observation given history then (b) builds interpretable cognitive maps from
the active bottlenecks indices. On perceptually aliased POEs, TDB (a) retains
the near-perfect prediction of transformers, (b) calls an external solver on
its latent graph to solve path planning problems exponentially faster, (c)
learns interpretable structure from text datasets, (d) exhibits emergent in-
context prediction and path planning abilities.

Our approach has two main limitations. First, TDB only accepts categorical
inputs. Second, though multiple discrete bottlenecks accelerate training, they
do not learn a disentangled latent space. To address these points, we want to
modify the TDB architecture (a) to accept high-dimensional continuous
observations (images) (b) so that different bottlenecks learn non-redundant
representations. Furthermore, we would like to extend this work into a
framework to build planning-compatible world models in rich environments. To
do so, we want to learn disentangled latent dynamics in factored Markov
decision processes by (a) extracting knowledge with transformers, (b) using
multiple discrete bottlenecks to compress this knowledge and to generate
latent nodes—or local latent graphs—and (c) learning factorized transition
matrices over these latent nodes.

## Acknowledgments

We thank Kevin Murphy, Daan Wierstra and Théophane Weber for useful
discussions during the preparation of this manuscript.

## References

  * Beattie et al. (2016) Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al.  Deepmind lab.  _arXiv preprint arXiv:1612.03801_ , 2016.
  * 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.
  * Carbonneau et al. (2020) Marc-André Carbonneau, Julian Zaidi, Jonathan Boilard, and Ghyslain Gagnon.  Measuring disentanglement: A review of metrics.  _arXiv preprint arXiv:2012.09276_ , 2020.
  * Chen et al. (2021a) Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  _Advances in neural information processing systems_ , 34:15084–15097, 2021a.
  * Chen et al. (2021b) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_ , 2021b.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Chrisman (1992) Lonnie Chrisman.  Reinforcement learning with perceptual aliasing: The perceptual distinctions approach.  In _AAAI_ , volume 1992, pages 183–188. Citeseer, 1992.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.  Training verifiers to solve math word problems.  _arXiv preprint arXiv:2110.14168_ , 2021.
  * Conmy et al. (2023) Arthur Conmy, Augustine N Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso.  Towards automated circuit discovery for mechanistic interpretability.  _arXiv preprint arXiv:2304.14997_ , 2023.
  * Dai et al. (2019) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov.  Transformer-xl: Attentive language models beyond a fixed-length context.  _arXiv preprint arXiv:1901.02860_ , 2019.
  * Dempster et al. (1977) Arthur P Dempster, Nan M Laird, and Donald B Rubin.  Maximum likelihood from incomplete data via the em algorithm.  _Journal of the royal statistical society: series B (methodological)_ , 39(1):1–22, 1977.
  * Elhage et al. (2021) Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al.  A mathematical framework for transformer circuits.  _Transformer Circuits Thread_ , 1, 2021.
  * George et al. (2021) Dileep George, Rajeev V Rikhye, Nishad Gothoskar, J Swaroop Guntupalli, Antoine Dedieu, and Miguel Lázaro-Gredilla.  Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps.  _Nature communications_ , 12(1):2392, 2021.
  * Ghahramani and Jordan (1995) Zoubin Ghahramani and Michael Jordan.  Factorial hidden markov models.  _Advances in Neural Information Processing Systems_ , 8, 1995.
  * Grill et al. (2020) Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent-a new approach to self-supervised learning.  _Advances in neural information processing systems_ , 33:21271–21284, 2020.
  * Guan et al. (2023) Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati.  Leveraging pre-trained large language models to construct and utilize world models for model-based task planning.  _arXiv preprint arXiv:2305.14909_ , 2023.
  * Guntupalli et al. (2023) J Swaroop Guntupalli, Rajkumar Vasudeva Raju, Shrinu Kushagra, Carter Wendelken, Danny Sawyer, Ishan Deshpande, Guangyao Zhou, Miguel Lázaro-Gredilla, and Dileep George.  Graph schemas as abstractions for transfer learning, inference, and planning.  _arXiv preprint arXiv:2302.07350_ , 2023.
  * Guo et al. (2022) Zhaohan Guo, Shantanu Thakoor, Miruna Pîslar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al.  Byol-explore: Exploration by bootstrapped prediction.  _Advances in neural information processing systems_ , 35:31855–31870, 2022.
  * Hagberg et al. (2008) Aric Hagberg, Pieter Swart, and Daniel S Chult.  Exploring network structure, dynamics, and function using networkx.  Technical report, Los Alamos National Lab.(LANL), Los Alamos, NM (United States), 2008.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Kamada et al. (1989) Tomihisa Kamada, Satoru Kawai, et al.  An algorithm for drawing general undirected graphs.  _Information processing letters_ , 31(1):7–15, 1989.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Lamb et al. (2022) Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Rajiv Didolkar, Dipendra Misra, Dylan J Foster, Lekan P Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford.  Guaranteed discovery of control-endogenous latent states with multi-step inverse models.  _Transactions on Machine Learning Research_ , 2022.
  * Li et al. (2022) Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg.  Emergent world representations: Exploring a sequence model trained on a synthetic task.  _arXiv preprint arXiv:2210.13382_ , 2022.
  * Lieberum et al. (2023) Tom Lieberum, Matthew Rahtz, János Kramár, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik.  Does circuit analysis interpretability scale? evidence from multiple choice capabilities in chinchilla.  _arXiv preprint arXiv:2307.09458_ , 2023.
  * Mialon et al. (2023) Grégoire Mialon, Clémentine Fourrier, Craig Swift, Thomas Wolf, Yann LeCun, and Thomas Scialom.  Gaia: a benchmark for general ai assistants.  _arXiv preprint arXiv:2311.12983_ , 2023.
  * Momennejad et al. (2023) Ida Momennejad, Hosein Hasanbeig, Felipe Vieira, Hiteshi Sharma, Robert Osazuwa Ness, Nebojsa Jojic, Hamid Palangi, and Jonathan Larson.  Evaluating cognitive maps and planning in large language models with cogeval.  _arXiv preprint arXiv:2309.15129_ , 2023.
  * Nanda et al. (2023) Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt.  Progress measures for grokking via mechanistic interpretability.  _arXiv preprint arXiv:2301.05217_ , 2023.
  * Olsson et al. (2022) Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al.  In-context learning and induction heads.  _arXiv preprint arXiv:2209.11895_ , 2022.
  * Pallagani et al. (2022) Vishal Pallagani, Bharath Muppasani, Keerthiram Murugesan, Francesca Rossi, Lior Horesh, Biplav Srivastava, Francesco Fabiano, and Andrea Loreggia.  Plansformer: Generating symbolic plans using transformers.  _arXiv preprint arXiv:2212.08681_ , 2022.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al.  Improving language understanding by generative pre-training.  2018\.
  * Ramesh et al. (2021) Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In _International Conference on Machine Learning_ , pages 8821–8831. PMLR, 2021.
  * Reed et al. (2022) Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al.  A generalist agent.  _arXiv preprint arXiv:2205.06175_ , 2022.
  * Sanfeliu and Fu (1983) Alberto Sanfeliu and King-Sun Fu.  A distance measure between attributed relational graphs for pattern recognition.  _IEEE transactions on systems, man, and cybernetics_ , (3):353–362, 1983.
  * Srivastava et al. (2022) Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso, et al.  Beyond the imitation game: Quantifying and extrapolating the capabilities of language models.  _arXiv preprint arXiv:2206.04615_ , 2022.
  * Swaminathan et al. (2023) Sivaramakrishnan Swaminathan, Antoine Dedieu, Rajkumar Vasudeva Raju, Murray Shanahan, Miguel Lazaro-Gredilla, and Dileep George.  Schema-learning and rebinding as mechanisms of in-context learning and emergence.  _arXiv preprint arXiv:2307.01201_ , 2023.
  * Valmeekam et al. (2022) Karthik Valmeekam, Alberto Olmo, Sarath Sreedharan, and Subbarao Kambhampati.  Large language models still can’t plan (a benchmark for llms on planning and reasoning about change).  _arXiv preprint arXiv:2206.10498_ , 2022.
  * Van Den Oord et al. (2017) Aaron Van Den Oord, Oriol Vinyals, et al.  Neural discrete representation learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wei et al. (2022) Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al.  Emergent abilities of large language models.  _arXiv preprint arXiv:2206.07682_ , 2022.
  * Xie et al. (2021) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.  An explanation of in-context learning as implicit bayesian inference.  _arXiv preprint arXiv:2111.02080_ , 2021.

##  Appendix A Vanilla transformer architecture

Fig.6 presents the vanilla transformer detailed in Sec.3.2. In our numerical
experiments, Sec.4, we train this model to minimize the autoregressive
objective Equation (1).

Figure 6: A vanilla transformer with causal mask takes as inputs the
respective linear embeddings of the observations and actions. A linear layer
on top of the action representations returns the next observation logits.

##  Appendix B On the connection between TDB and CSCG

###  B.1 Computational and space complexity

A clone-structured causal graph (CSCG) (George et al., 2021) is a hidden
Markov model variant used to model sequences of categorical observations. The
latent space of a CSCG is partitioned, such that all the latent variables in
one set of the partition deterministically emit the same observation. These
latent variables are the clones of the observation. Because of its latent
structure, CSCG models a transition between two observations as a transition
between their respective clones—as opposed to a transition between all the
latent variables.

We compare below the complexity between a vanilla CSCG and a vanilla self-
attention layer when running inference on a sequence of length $N$.

Computation complexity:  For a vanilla CSCG with $M$ clones per observation,
the computational cost for computing the forward and backward messages is
$O(M^{2}N)$. In contrast, for a vanilla self-attention layer with a
$D$-dimensional embedding, the forward pass cost is $O(N^{2}D)$. In this
paper, we use $N=400$, $D=256$. CSCG is then more expensive when the number of
clones per observation satisfies $M\gg\sqrt{ND}=320$.

Space complexity:  The memory cost of storing a vanilla CSCG with $M$ clones
per observation and $E$ different observations is $O(M^{2}E^{2})$. This large
dense memory matrix can grow quickly as the number of observations increases,
and prohibit vanilla CSCGs from scaling to very large latent space. In
contrast, a vanilla self-attention layer with a $D$-dimensional embedding
requires a $O(N^{2}+ND)$ memory. The quadratic dependence on the sequence
length is problematic for large corpora, but not for the problem we study
here, where $N=400$.

###  B.2 A transformer with clone-structured discrete bottleneck

We propose herein a variant of the TDB architecture inspired by the CSCG
model, which injects a “clone-structured” inductive bias into its discrete
bottlenecks. We consider the dictionary of latent codes introduced in Sec.3.3:
$\mathcal{D}=(d_{1},\ldots,d_{K})$, where
$d_{k}\in\mathbb{R}^{D},\leavevmode\nobreak\ \forall k$. We now allocate a
subset of the latent codes for each observation $x$ in the vocabulary $V$.
That is, we first define a partition of the indices
$\\{\mathcal{C}(x)\\}_{x\in V}$ such that:

$\bigcup_{x\in V}\mathcal{C}(x)=\\{1,\ldots,K\\}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ ;\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\
\mathcal{C}(x)\cap\mathcal{C}(y)=\emptyset,\leavevmode\nobreak\ \forall x\neq
y\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
;\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\mathcal{C}(x)\neq\emptyset,\leavevmode\nobreak\ \forall x.$

Second, we generalize the operator $\phi$ in Equation (2) so that it now
depends on both (a) the transformer output and (b) the observation itself.
That is, we define:

$\tilde{\phi}(y,x)=\operatorname*{argmin}_{k\in\mathcal{C}(x)}\|y-d_{k}\|_{2}^{2},\leavevmode\nobreak\
\leavevmode\nobreak\ \forall y\in\mathbb{R}^{d},\leavevmode\nobreak\ \forall
x\in V$ (5)

For the observation $x_{n}$ which leads to the transformer output $T_{2n-1}$,
the discrete quantizer now returns the latent code of $\mathcal{D}$ with index
$\tilde{\phi}(T_{2n-1},x_{n})$. That is, each observation can now only
activate a subset of the latent codes. We refer to this new quantizer module
as clone-structured discrete bottleneck.

Computational gain:  This architecture change decreases the computational cost
of the discrete bottleneck by a linear term $O(E)$—which can be of interest
when the number of observations $E$ is large. However, the overall
computational cost is still dominated by the quadratic cost of computing the
attention matrices in the causal transformer.

Performance gain:  Overall, we did not find any major benefit from replacing
the discrete bottleneck of a TDB with its clone-structure variants. Our most
promising result was that we were able to solve the $3$D synthetic
environments in Sec.4.3 using a single clone-structured discrete
bottleneck—where a TDB needs more than one discrete bottleneck. However, a
transformer with clone-structured bottlenecks would struggle to solve the in-
context learning experiments in Sec.4.5. Indeed, a new test room with
observations $5,6,7,8$ would, by definition, use different latent codes than a
training room with observations $1,2,3,4$, which prevents the model from
learning to share structure across rooms. This convinced us to not use the
clone-structured discrete bottleneck in our experiments.

##  Appendix C Clustering steps to build the transition graph

In this section, we detail the clustering steps that we apply to learn a
cognitive map from the bottleneck indices in Sec.3.6.

#### Step 1. Cluster the bottlenecks indices with Hamming distance:

This first clustering step only applies in the case where the TDB has multiple
$M>1$ discrete bottlenecks. In this case, TDB maps each observation $x_{n}$ to
a tuple of latent code indices:

$s_{n}=\left(\phi^{(1)}(T_{2n-1}),\ldots,\phi^{(M)}(T_{2n-1})\right)\in\\{1,\ldots,K\\}^{M}.$

Before building the cognitive map, we first collect in a set $\mathcal{J}$ all
the unique tuples of latent indices appearing in the training set. For two
tuples $s=(s^{(1)},\ldots,s^{(M)})\in\mathcal{J}$ and
$\ell=(\ell^{(1)},\ldots,\ell^{(M)})\in\mathcal{J}$, we compute their Hamming
distance:

$H(s,\ell)=\frac{1}{M}\sum_{j=1}^{M}\mathbf{1}(s^{(j)}\neq\ell^{(j)})\in[0,1].$

We introduce a distance threshold $d_{\text{Hamming}}$—which we set to $0.25$
in Sec.4—and merge any pair of tuples such that $H(s,\ell)\leq
d_{\text{Hamming}}$. This means that we will not distinguish $\ell$ and $s$
when we build the count tensor $C$—see Sec.3.6. Consequently, these two tuples
will be mapped to the same node on the cognitive map $\mathcal{G}$—and any
node in $\mathcal{G}$ will be associated with a collection of tuples of
bottleneck indices.

Intuitively, this means that when the entries of two tuples of bottleneck
indices are “almost all the same”, then we assume that their associated
observations are at the same spatial position.

#### Step 2. Merge identical nodes:

Second, we build the count tensor $C$ (after this first clustering step), and
threshold it with $t_{\text{ratio}}$ to build an action-augmented transition
graph $\mathcal{G}$—as detailed in Sec. 3.6. We only retain nodes that have at
least two incoming and outgoing neighbors. Finally, our second clustering step
merges every pair of nodes in $\mathcal{G}$ that are connected to the same set
of neighbors.

#### Step 3. Map each discarded node to its closest retained node:

Thresholding the count tensor $\mathcal{C}$ may discard some (tuple of)
bottleneck indices with low empirical counts. In particular, as we show in
Fig.10, Appendix I.3, many bottleneck indices that are active early in the
sequences are discarded. This may be problematic for solving path planning
problems, as it means that the corresponding observations cannot be mapped to
nodes in $\mathcal{G}$. Indeed, when TDB activates a discarded (tuple of)
bottleneck indices, it cannot locate itself in the learned cognitive map
$\mathcal{G}$.

To remedy this, we divide the set of unique (tuples of) latent indices
$\mathcal{J}$, between the elements that are discarded—i.e. not mapped to a
node in $\mathcal{C}$—and retained—i.e. mapped to a node in $\mathcal{C}$. Our
last clustering step maps any discarded (tuple of) latent indices
$s_{\text{discarded}}$ to its closest retained latent index for the $\ell_{1}$
distance, defined as:

$\hat{s}_{\text{retained}}\leavevmode\nobreak\ \leavevmode\nobreak\
\in\leavevmode\nobreak\ \leavevmode\nobreak\
\operatorname*{argmin}_{s_{\text{retained}}}\sum_{a=1}^{N_{\text{actions}}}\sum_{\ell=1}^{K}\left|p(\ell\leavevmode\nobreak\
|\leavevmode\nobreak\ s_{\text{retained}},a)-p(\ell\leavevmode\nobreak\
|\leavevmode\nobreak\ s_{\text{discarded}},a)\right|.$

After this step, $s_{\text{discarded}}$ and $\hat{s}_{\text{retained}}$ are
mapped to the same node in $\mathcal{C}$, and all the entries in $\mathcal{J}$
are retained.

##  Appendix D Solving path planning problem at test time on partially
observed environments

We detail herein how a vanilla sequence model and a TDB can solve the test
shortest path problem presented in Sec.4.2.

#### Defining the shortest path problem:

We consider a test sequence of observations $x=(x_{1},\ldots,x_{N})$ and
actions $a=(a_{1},\ldots,a_{N})$ in a POE—the environment can either be the
$2$D aliased rooms of Sec.4.2 and Sec.4.5, the aliased cube of Sec.4.2, or the
egocentric $3$D synthetic environments of Sec.4.3. We also denote
$\operatorname*{pos}=(\operatorname*{pos}_{1},\ldots,\operatorname*{pos}_{N})$
the unobserved sequence of room spatial positions associated with each
observation. Note that, because of aliasing, we cannot deterministically infer
$\operatorname*{pos}_{n}$ from $x_{n}$.

For the $2$D aliased rooms (Sec.4.2), the spatial position of an observation
is its $2$D coordinates. For the egocentric $3$D synthetic environments
(Sec.4.3), the spatial position of an observation is a $3$D vector: the first
two entries correspond to the $2$D coordinates of the observation, while the
last entry corresponds to the agent’s heading direction, which is one of
$\\{0^{\circ},90^{\circ},180^{\circ},270^{\circ}\\}$

Let $C$ be a context length—we set $C=50$ or $C=100$ in practice. The path
planning problem we consider in Sec.4 consists of finding the shortest path,
i.e., the shortest sequence of actions, which leads from the room position
$\operatorname*{pos}_{C}$ (associated with the observation $x_{C}$) to the
room position $\operatorname*{pos}_{N-C}$ (associated with the observation
$x_{N-C}$). Naturally, the sequence $(a_{C},\ldots,a_{N-C-1})$ is a path of
length $N-2C$ from $\operatorname*{pos}_{C}$ to $\operatorname*{pos}_{N-C}$,
which we refer to as the fallback path. In addition, for evaluating a model’s
performance, we use an optimal shortest path, which solves the shortest path
problem in the ground truth room—note that the agent does not know the ground
truth layout. While the optimal shortest path is not unique, all the optimal
shortest paths have the same length.

Fig.1 in the main text illustrates this path planning problem on a $2$D
aliased room—we use $N=40,C=5$ for the sake of visualization.

We detail below our proposed procedure for solving this shortest path problem
with a TDB; and with a vanilla transformer or a vanilla LSTM. As mentioned in
the main text, a vanilla sequence model finds this challenging because (a) it
can only perform forward rollouts—without being able to collapse redundant
visitations of the same spatial position—and (b) it cannot tell whether it has
reached its destination. In contrast, a TDB solves this problem by querying
its learned latent graph.

We are interested in measuring (a) whether each model can derive a path that
improves over the fallback path and (b) when a model derives such a better
path, how much worse it is compared to the optimal shortest path.

#### Deriving shortest paths with TDB:

When TDB is evaluated on the test sequences $x$ and $a$, it maps each
observation $x_{n}$ to either (a) a single latent code index
$s_{n}=\phi(T_{2n-1})\in\\{1,\ldots,K\\}$ when we have a single discrete
bottleneck or (b) a tuple of latent codes indices
$s_{n}=\left(\phi^{(1)}(T_{2n-1}),\ldots,\phi^{(M)}(T_{2n-1})\right)$ when we
have multiple discrete bottlenecks.

As a reminder, each node in the learned cognitive map $\mathcal{G}$ is
associated with a collection of (tuples of) bottleneck indices—see Appendix C.
However, the (tuples of) bottleneck indices $s_{C}$ and $s_{N-C}$ may not be
associated with a node in $\mathcal{G}$. Consequently, we find two tuples of
bottleneck indices $\tilde{s}_{C}$ and $\tilde{s}_{N-C}$ (a) with lowest
Hamming distance from $s_{C}$ and $s_{N-C}$ and (b) which are associated with
two nodes $n_{C}$ and $n_{N-C}$ in $\mathcal{G}$ . Finally, we call the
external networkx.shortest_path solver to find a shortest path between the
nodes $n_{C}$ and $n_{N-C}$.

#### Deriving shortest paths from rollouts with vanilla sequence models:

A naive option that finds the shortest path between $\operatorname*{pos}_{C}$
and $\operatorname*{pos}_{N-C}$ with a vanilla sequence model consists in
exploring all the possible sequences of observations and actions. For a number
of evaluations exponential in the length of the optimal shortest path, this
approach is guaranteed to find an optimal shortest path.

We propose herein an alternative approach, which (a) is more computationally
efficient, (b) is not guaranteed to find the optimal shortest path, and (c) in
practice, often improves over the fallback path. Beforehand, let us define the
subsequences of observations context $x_{\text{context}}=(x_{1},\ldots,x_{C})$
and actions context $a_{\text{context}}=(a_{1},\ldots,a_{C-1})$. We also
define the subsequences of observations tail
$x_{\text{tail}}=(x_{N-C+1},\ldots,x_{N})$ and actions tail
$a_{\text{tail}}=(a_{N-C},\ldots,a_{N-1})$.

First, we generate $N_{\text{random walks}}$ random sequences of actions, each
of length $N-2C$: we denote $a^{(i)}=(a^{(i)}_{C},\ldots,a^{(i)}_{N-C-1})$ the
$i$th random sequence—we initialize the indices at $C$. In Sec.4.3, we set
$N_{\text{random walks}}=10$.

Second, we augment each sequence of actions (of length $N-2C$) into two
sequences of observations and actions, of respective lengths $N-2C$ and
$N-2C-1$, as follows. The first $C$ observations and $C-1$ actions are the
shared context $x_{\text{context}}$ and $a_{\text{context}}$. The next $N-2C$
actions correspond to $a^{(i)}$. The corresponding sequence of observations
corresponds to the sequence model’s autoregressive predictions. That is, for
$n=C,\ldots,N-C-1$, we iteratively define $\hat{x}^{(i)}_{n+1}$ as:

$\hat{x}^{(i)}_{n+1}\in\operatorname*{argmax}_{x}p\left(x\leavevmode\nobreak\
\leavevmode\nobreak\ \bigg{|}\leavevmode\nobreak\ \leavevmode\nobreak\
\underbrace{x_{1},\leavevmode\nobreak\
a_{1},\ldots,a_{C-1},\leavevmode\nobreak\ x_{C}}_{\text{shared
context}},\leavevmode\nobreak\ a^{(i)}_{C},\leavevmode\nobreak\
\hat{x}^{(i)}_{C+1},\leavevmode\nobreak\
a^{(i)}_{C+1},\ldots,\hat{x}^{(i)}_{n},\leavevmode\nobreak\ a_{n}\right)$

For each generated random walk, we look for all the observations equal to
$x_{N-C}$, and define the set of candidates:

$\mathcal{C}=\left\\{(i,n):n\geq C\text{ and
}\hat{x}^{(i)}_{n}=x_{N-C}\right\\}.$

#### Estimating whether the target is reached via tail evaluation:

As we mentioned in the main text, because of aliasing, the vanilla sequence
model cannot know whether it has reached the target position when it observes
a candidate. That is, it does not know whether the spatial position associated
with the candidate is equal to $\operatorname*{pos}_{N-C}$.

We could return all the candidates and evaluate them using the ground truth
map. Instead, we propose herein to estimate whether a candidate $(i,n)$ is at
the target position $\operatorname*{pos}_{N-C}$ by evaluating it—with its
history—on the tail subsequences of observations $x_{\text{tail}}$ and actions
$a_{\text{tail}}$. That is, for a candidate $(i,n)$ we define, for
$k=1,\ldots,C$:

$\hat{x}^{(i)}_{\text{tail},k}\in\operatorname*{argmax}_{x}p\left(x\leavevmode\nobreak\
\leavevmode\nobreak\ \bigg{|}\leavevmode\nobreak\ \leavevmode\nobreak\
\underbrace{x_{1},\leavevmode\nobreak\
a_{1},\ldots,a_{C-1},\leavevmode\nobreak\ x_{C}}_{\text{shared
context}},\leavevmode\nobreak\ \underbrace{a^{(i)}_{C},\leavevmode\nobreak\
\hat{x}^{(i)}_{C+1},\leavevmode\nobreak\
a^{(i)}_{C+1},\ldots,\hat{x}^{(i)}_{n}=x_{N-C}}_{\text{autoregressive random
walk}},\leavevmode\nobreak\ a_{N-C},x_{N-C+1},\ldots,a_{N-C+k-1}\right)$

We estimate that a candidate is at the target position
$\operatorname*{pos}_{N-C}$ when
$(\hat{x}^{(i)}_{\text{tail},1},\ldots,\hat{x}^{(i)}_{\text{tail},C})=x_{\text{tail}}$.
Note that, because of aliasing, even when the tail evaluation succeeds, the
candidate may not be at the target position.

#### Evaluating valid paths:

For the TDB, we return a single shortest path proposal, which is estimated by
the external solver. For the transformer and the LSTM, we return all the path
proposals that are estimated to be at the target position.

For each model, we first evaluate whether the proposed paths are valid. That
is, we test whether the proposed sequence of actions correctly leads from
$\operatorname*{pos}_{C}$ to $\operatorname*{pos}_{N-C}$ in the ground truth
environment. For a vanilla sequence model, when multiple returned paths are
valid, we only retain the shortest one. We then compute two metrics (a) a
binary indicator of whether the method has found a valid path and (b) if (a)
is correct, the ratio between the length of the valid found path and the
optimal path length. Finally, ImpFallback averages (a) while RatioSP averages
(b) over $200$ shortest paths problems.

#### Tail evaluation improves the quality of the paths returned:

For a transformer trained on $2$D aliased rooms in Sec.4.2, if we return all
the candidates, then the ratio of valid paths among the paths returned is
$0.78\%(\pm 0.03\%)$. Indeed, most of the candidates are not at the target
spatial position. When we estimate whether we are at the target position via
tail evaluation, this ratio goes up to $38.76\%(\pm 1.17\%)$—prediction errors
prevent it from being at $100\%$. Similarly, on the $3$D synthetic
environments of Sec.4.3, a transformer without tail evaluation returns
$1.75\%(\pm 0.20\%)$ of valid paths while a transformer with tail evaluation
returns $25.78\%(\pm 1.17\%)$. Finally, on the ICL experiments of Sec.4.5, for
a transformer trained on $10k$ rooms, this same ratio goes from $2.58\%(\pm
0.05\%)$ without tail evaluation up to $51.10\%(\pm 1.81\%)$ with it. Overall,
tail evaluation drastically improves the quality of the paths returned by
sequence models.

##  Appendix E Multiple discrete bottlenecks accelerate training

Tables 2, 3 and 4 below illustrate that multiple discrete bottlenecks
accelerate training.

2D aliased rooms:  For each TDB trained on the $2$D aliased rooms and reported
in Sec.4.2, we compute, for every $1000$ training steps, its training accuracy
on the entire training set. We define a new metric, TrainingAbove98, which
reports the first evaluation when this training accuracy is higher than
$98\%$. Table 2 averages this metric over the $10$ experiments run in Sec.4.2,
for various TDB models.

Method |  TrainingAbove98 $\downarrow$
---|---
TDB$(S=1,M=1)$ |  $9400\leavevmode\nobreak\ (290)$
TDB$(S=1,M=4)$ |  $4300\leavevmode\nobreak\ (145)$
TDB$(S=1,\operatorname*{enc},M=1)$ |  $9300\leavevmode\nobreak\ (425)$
TDB$(S=1,\operatorname*{enc},M=4)$ |  $4500\leavevmode\nobreak\ (158)$
TDB$(S=3,M=1)$ |  $8300\leavevmode\nobreak\ (284)$
TDB$(S=3,M=4)$ |  $4000\leavevmode\nobreak\ (0)$

Table 2: On $2$D aliased rooms, training a TDB with $M=4$ discrete bottlenecks
converges faster than training its counterpart with a single discrete
bottleneck.

$3$D synthetic environments:  As discussed in Sec.4.2, we need $M=4$ discrete
bottlenecks for the TDB training to converge in $3$D synthetic environments.
Table 3 illustrates this by reporting TrainingAfter1k—the training accuracy
after $1000$ training steps—for the different models reported in Table 7,
Appendix J.

Method |  TrainingAfter1k $(\%)$ $\uparrow$
---|---
TDB$(S=1,M=1)$ |  $41.40\leavevmode\nobreak\ (0.09)$
TDB$(S=1,M=4)$ |  $99.72\leavevmode\nobreak\ (0.03)$
TDB$(S=1,\operatorname*{enc},M=1)$ |  $40.81\leavevmode\nobreak\ (0.93)$
TDB$(S=1,\operatorname*{enc},M=4)$ |  $99.71\leavevmode\nobreak\ (0.04)$
TDB$(S=3,M=1)$ |  $28.05\leavevmode\nobreak\ (1.88)$
TDB$(S=3,M=4)$ |  $99.63\leavevmode\nobreak\ (0.06)$

Table 3: On the $3$D synthetic environments, training a TDB with $M=4$
discrete bottlenecks converges faster than training its counterpart with a
single discrete bottleneck.

GINC text dataset:  On the GINC dataset, we also evaluate the training
accuracy of each model reported in Sec.4.5 at every $1000$ training steps.
Here, TrainingAbove98 reports the first iteration when the training accuracy
becomes higher than $98\%$ of the highest value it reaches during training.

Method |  TrainingAfter1k $(\%)$ $\uparrow$
---|---
TDB$(S=1,M=1)$ |  $16000$
TDB$(S=1,M=4)$ |  $6000$
TDB$(S=3,M=1)$ |  $13000$
TDB$(S=3,M=4)$ |  $5000$
TDB$(S=5,M=1)$ |  $11000$
TDB$(S=5,M=4)$ |  $5000$

Table 4: On the GINC dataset, training a TDB with $M=4$ discrete bottlenecks
converges faster than training its counterpart with a single discrete
bottleneck.

##  Appendix F Do multiple discrete bottlenecks learn a disentangled
representation?

#### Predicting a discrete bottleneck index given the other indices:

We propose herein a metric to test whether multiple discrete bottlenecks learn
a disentangled representation of the observations.

If the discrete bottlenecks learn of a disentangled representation, then (a)
they have to be independent (b) each one has to explain a distinct factor of
the data (Carbonneau et al., 2020). Given a TDB with $M$ discrete bottlenecks,
we propose to estimate (a), that is, whether the bottlenecks are independent.
To do so, we train a logistic regression estimator to predict each bottleneck
index given the other $M-1$ bottleneck indices.

Given a sequence of observations $(x_{1},\ldots,x_{N})$, a TDB with $M$
discrete bottlenecks maps each observation $x_{n}$ to the latent codes
$\left(s_{n}^{(1)},\ldots,s_{n}^{(M)}\right)$ where
$s_{n}^{(i)}=\phi^{(i)}(T_{2n-1})\in\\{1,\ldots,K\\}$—see Sec.3.5.

We denote $e^{(i)}_{n}=\left(\mathbf{1}(s_{n}^{(i)}=k)\right)_{1\leq k\leq K}$
the one-hot encoding of $s_{n}^{(i)}$.

Let us now assume that we want to predict the $M$th bottleneck index given the
first $M-1$ bottleneck indices.

We note $z=\left(e^{(1)}_{n},\ldots,e^{(M-1)}_{n}\right)\in\\{0,1\\}^{(M-1)K}$
and $y=e^{(M)}_{n}\in\\{0,1\\}^{K}$, where we drop the dependency over $n$. We
also introduce a matrix $W\in\mathbb{R}^{K\times(M-1)K}$. The logistic
regression objective at timestep $n$ is:

$\mathcal{L}^{(M)}_{\text{disentanglement}}(n)=-\sum_{i=1}^{K}y_{i}\leavevmode\nobreak\
\log\left(\frac{\exp(W_{i}^{T}z)}{\sum_{j=1}^{K}\exp(W_{j}^{T}z)}\right)$

We also define
$\mathcal{L}^{(M)}_{\text{disentanglement}}=\sum_{n=1}^{N}\mathcal{L}^{(M)}_{\text{disentanglement}}(n)$.
Note that we do not use an intercept term or a regularization term.

#### Training:

Given a training and test set from Sec.4, each consisting of $N_{\text{seq}}$
test sequences of length $N$, we first compute all the active bottleneck
indices, which gives us two tensors, $\mathcal{T}_{\text{train}}$ and
$\mathcal{T}_{\text{test}}$, each of dimensions $N_{\text{seq}}\times N\times
M$.

We train $M$ logistic regression estimators on $\mathcal{T}_{\text{train}}$:
the $i$th estimator is trained to minimize
$\mathcal{L}^{(i)}_{\text{disentanglement}}$, that is, to predict the $i$th
bottleneck index given the other $M-1$. For training, we use Adam (Kingma and
Ba, 2014) for $5,000$ iterations, with a learning rate of $0.001$ and a batch
size of $32$.

#### Disentanglement metric:

Our proposed disentanglement metric is the averaged test accuracy of the $M$
logistic regression estimators on the test tensor $\mathcal{T}_{\text{test}}$.

#### Results:

In Sec.4, we use $K=1000$ latent codes for each discrete bottleneck: if the
bottleneck representations are independent, the test disentanglement accuracy
would be $0.20\%$. However, all the test disentanglement accuracies reported
in the main paper, are above $85\%$, which shows that multiple discrete
bottlenecks learn highly redundant representations.

##  Appendix G TDB with single-step objective merges nodes with identical
neighbors

In Sec.3.7, we make the argument that a TDB only trained to predict the next
observation is not encouraged to disambiguate observations with identical
neighbors in the ground truth (GT) room and, as a result, may merge their
representations. We use this same argument in Sec.3.4 to justify the need for
augmenting the loss via either (a) a multi-step objective or (b) next encoding
prediction.

Fig.7 demonstrates the argument. We consider a $2$D aliased GT room from Sec.
4.2 with $O=12$ unique observations—the numerical results are presented in
Table 6, Appendix I.2. Similar to Fig.1, the GT room contains a $4\times 4$
patch with black borders repeated twice. Each $4\times 4$ patch contains an
inner $2\times 2$ patch with fuchsia borders. For each node inside this inner
patch, all four actions in the GT room (move up, down, left, and right) lead
to the same neighbors, regardless of whether the node in the top-left $2\times
2$ patch or the bottom-right one.

Fig.7[right] shows the latent graph learned by a TDB$(S=1,M=4)$ with single
step prediction and four discrete bottlenecks. Because this model is only
trained to predict the next observation given the history, it learns the same
representation for each pair of nodes with same location inside the inner
$2\times 2$ patch. This results in the presence of four merged nodes in the
latent graph—colored in fuchsia. When a merged node is active, the agent can
(a) perfectly predict the next observation but (b) it cannot know which one of
the two possible $2\times 2$ patches it is in. In addition, the merged nodes
introduce unrealistic shortcuts in the learned latent graph: nodes that are
far apart in the GT room are close in the latent graph. Because the learned
cognitive map does not accurately model the ground truth dynamics, the
external solver will propose shortest paths that are invalid in the GT room.
As a consequence, as indicated in Table 6, Appendix I.2, TDB$(S=1,M=4)$ can
only solve $59.37\%(\pm 3.93\%)$ of the path planning problems.

Figure 7: [Left] A 2D aliased room with $O=12$ unique observations: for each
observation inside the $2\times 2$ patches with fuchsia boundaries, all four
actions (move up, down, left, right) lead to the same neighbors, regardless of
whether the node is in the top-left $2\times 2$ patch or the bottom-right one.
[Right] The latent graph learned by a TDB with single step prediction fails at
learning the ground truth dynamics. TDB learns the same representation for the
two elements of each pair. This results in four merged nodes—colored in
fuchsia—which are active at two spatial locations. While these merged nodes do
not affect test prediction performance, they introduce unrealistic shortcuts
in the latent graph, which fool the external planner.

For aliased rooms with 4 colors: When $O=4$, a higher frequency of
observations have identical neighbors. The learned latent graph merges these
observations, which introduces a larger number of unrealistic shortcuts. This
explains the poor performance of TDB($S=1,M=1$) in Table 1, which only
improves $6.61\%(\pm 0.53\%)$ of the time over the fallback path. Because of
this inherent model failure, TDB($S=1,M=4$) similarly only improves over the
fallback path $6.21\%(\pm 0.62\%)$ of the time. In fact, TDB is only able to
solve the easiest path planning problems, for which the optimal shortest path
consists of only a small number of actions. Indeed, the average length of the
optimal shortest paths solved by TDB($S=1,M=4$) is $2.15\leavevmode\nobreak\
(\pm 0.07)$ while the average length of the optimal shortest paths over all
the path planning problems is $10.80\leavevmode\nobreak\ (\pm 0.08)$.

For 3D synthetic environments: Fig.8[middle] shows that, for a $3$D
environment from Section 4.3, the cognitive map learned with a TDB$(S=1,M=4)$
merges three pairs of nodes. These merged nodes introduce unrealistic
shortcuts on the learned latent graph: as a result, TDB$(S=1,M=4)$ fails at
path planning and only improves $65.30\%(\pm 4.24\pm\%)$ of the time over the
fallback path.

Fig.8[right] shows the latent graph learned when the objective is augmented to
predict the next three steps, and highlights the three pairs of nodes merged
in the middle image. Note that each node is mapped to the observation with
higher empirical frequency when this node is active. As before, we see that
the merged nodes have identical neighbors555Because there are $128$ clusters,
the reader cannot precisely infer the observation from the colors: we looked
at the graph indices to make sure that the neighbors of the merged nodes are
identical..

Figure 8: [Left] Top-down view of a $3$D simulated rooms used in Sec.4.3.
[Middle] Latent graph learned with a TDB$(S=1,M=4)$ with single-step
objective: each pair of observations with identical neighbors are represented
with the same node, which introduces unrealistic shortcuts. [Right] Color-
coded latent graph learned with a multi-step objective.

##  Appendix H Accelerating the graph edit distance

The graph edit distance (GED) (Sanfeliu and Fu, 1983) between two graphs
$\mathcal{G}_{1}$ and $\mathcal{G}_{2}$ is a graph similarity measure defined
as minimum cost of edit path—i.e. sequence of node and edge edit operations—to
transform the graph $\mathcal{G}_{1}$ into a graph isomorphic to
$\mathcal{G}_{2}$.

Computing the exact GED is NP-complete.
We therefore rely on the networkx.optimize_graph_edit_distance (Hagberg et
al., 2008) function to return a sequence of approximations. This method
accepts as argument a function to indicate which nodes (resp. edges) of
$\mathcal{G}_{1}$ and $\mathcal{G}_{2}$ should be considered equal during
matching.

To accelerate GED, we say that a node $n_{1}$ of $\mathcal{G}_{1}$ and a node
$n_{2}$ of $\mathcal{G}_{2}$ have to be considered equal — and we note
$n_{1}\approx n_{2}$ if the ground truth spatial positions with higher
empirical frequency when $n_{1}$ (resp. $n_{2}$) is active are the same.
Similarly, we say that an edge of $\mathcal{G}_{1}$ connecting the nodes
$(n^{s}_{1},n^{t}_{1})$ and an edge of $\mathcal{G}_{2}$ connecting the nodes
$(n^{s}_{2},n^{t}_{2})$ are considered to be equal if either (a)
$n_{1}^{s}\approx n_{2}^{s}$ and $n_{1}^{t}\approx n_{2}^{t}$ or (b)
$n_{1}^{s}\approx n_{2}^{t}$ and $n_{1}^{t}\approx n_{2}^{s}$.

We use a timeout of $900$s to compute GED in all the experiments but the ICL
one, for which we use $20$s (due to the large number of normalized graph edit
distances being computed).

##  Appendix I Additional materials for aliased cube and for $2$D aliased room
cubes

###  I.1 Table of results for non-Euclidean aliased cube

Table 6 reports the prediction and path planning metrics for an aliased cube
with edge size $6$, similar to Fig. 3[bottom-left]. All the settings and
parameters are the same as in Section 4.2. We average the results over $10$
experiments: each experiment considers a different cube.

Method |  TestAccu ($\%$) $\uparrow$ |  ImpFallback ($\%$) $\uparrow$ |  RatioSP $\downarrow$ |  NormGED $\downarrow$
---|---|---|---|---
Vanilla transformer |  $\mathbf{99.76}\leavevmode\nobreak\ (0.00)$ |  $43.82\leavevmode\nobreak\ (1.17)$ |  $22.74\leavevmode\nobreak\ (0.84)$ |  —
Vanilla LSTM |  $\mathbf{99.76}\leavevmode\nobreak\ (0.00)$ |  $43.82\leavevmode\nobreak\ (1.17)$ |  $22.74\leavevmode\nobreak\ (0.84)$ |  —
TDB$(S=1,\operatorname*{enc},M=1)$ |  $99.36\leavevmode\nobreak\ (0.06)$ |  $99.07\leavevmode\nobreak\ (0.88)$ |  $\mathbf{1.02}\leavevmode\nobreak\ (0.00)$ |  $\mathbf{0.334}\leavevmode\nobreak\ (0.016)$
TDB$(S=3,M=1)$ |  $99.76\leavevmode\nobreak\ (0.00)$ |  $\mathbf{99.05}\leavevmode\nobreak\ (0.23)$ |  $1.03\leavevmode\nobreak\ (0.00)$ |  $0.304\leavevmode\nobreak\ (0.011)$

Table 5: Results averaged over $10$ aliased cubes of edge size $6$. Our TDB
with either multi-step objective or next encoding prediction (a) retains the
nearly perfect test accuracy of vanilla sequence models (b) consistently
solves the shortest paths problems when paired with an external solver—while
both transformer or LSTM catastrophically fail—(c) learns cognitive maps
almost isomorphic to the ground truth.

###  I.2 Table of results for 12 unique observations

Table 6 reports the prediction and path planning metrics for an aliased room
with $O=12$ unique observations and global aliasing—as a $4\times 4$ patch is
repeated twice. All the settings and parameters are the same as in Section
4.2.

First, during training, a TDB with a single discrete bottleneck either (a) is
stuck in a local optimum or (b) converges after a long time. As a consequence,
contrary to Table 6 which considers $O=4$, all the TDBs with a single
bottleneck are here outperformed by their counterparts with four bottlenecks.
TDBs with a single bottleneck also learn worse cognitive maps and solve a
lower frequency of path planning problems.

Second, regardless of the number of bottleneck used, a TDB with single-step
objective cannot disambiguate the global aliasing and introduces unrealistic
shortcuts—see Appendix G—which result in a drop in planning metrics.

In contrast, TDBs with $M=4$ bottlenecks and either multi-step or next
encoding objective (a) achieves nearly perfect prediction accuracy, (b) are
able to almost consistently find the optimal shortest path, and (c) learn
cognitive maps that recover the GT dynamics.

Method |  TestAccu ($\%$) $\uparrow$ |  ImpFallback ($\%$) $\uparrow$ |  RatioSP $\downarrow$ |  NormGED $\downarrow$
---|---|---|---|---
Vanilla transformer |  $\mathbf{99.75}\leavevmode\nobreak\ (0.00)$ |  $29.89\leavevmode\nobreak\ (0.94)$ |  $17.45\leavevmode\nobreak\ (0.57)$ |  —
Vanilla LSTM |  $\mathbf{99.75}\leavevmode\nobreak\ (0.00)$ |  $29.89\leavevmode\nobreak\ (0.94)$ |  $17.45\leavevmode\nobreak\ (0.57)$ |  —
TDB$(S=1,M=1)$ |  $89.46\leavevmode\nobreak\ (2.48)$ |  $17.31\leavevmode\nobreak\ (4.52)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.265\leavevmode\nobreak\ (0.053)$
TDB$(S=1,M=4)$ |  $\mathbf{99.75}\leavevmode\nobreak\ (0.00)$ |  $59.37\leavevmode\nobreak\ (3.93)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $\mathbf{0.033}\leavevmode\nobreak\ (0.002)$
TDB$(S=1,\operatorname*{enc},M=1)$ |  $92.57\leavevmode\nobreak\ (3.04)$ |  $58.21\leavevmode\nobreak\ (10.98)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.145\leavevmode\nobreak\ (0.038)$
TDB$(S=1,\operatorname*{enc},M=4)$ |  $99.71\leavevmode\nobreak\ (0.00)$ |  $96.58\leavevmode\nobreak\ (2.13)$ |  $1.01\leavevmode\nobreak\ (0.00)$ |  $0.157\leavevmode\nobreak\ (0.016)$
TDB$(S=3,M=1)$ |  $99.33\leavevmode\nobreak\ (0.01)$ |  $97.09\leavevmode\nobreak\ (2.06)$ |  $\mathbf{1.00}\leavevmode\nobreak\ (0.00)$ |  $0.157\leavevmode\nobreak\ (0.010)$
TDB$(S=3,M=4)$ |  $99.73\leavevmode\nobreak\ (0.01)$ |  $\mathbf{99.75}\leavevmode\nobreak\ (0.14)$ |  $1.01\leavevmode\nobreak\ (0.00)$ |  $0.160\leavevmode\nobreak\ (0.017)$

Table 6:  Prediction and path planning metrics averaged over $10$ aliased
$15\times 20$ rooms with $O=12$ observations. Training a TDB single bottleneck
($M=1$) is slow and sometimes does not converge. A TDB with single-step
objective cannot disambiguate global aliasing, fails at recovering the ground
truth dynamics. TDB with multiple discrete bottlenecks ($M=4$) outperform
their single bottleneck counterpart ($M=1$) for both prediction and path
planning performance.

###  I.3 The agent can locate itself after a short context

#### Test accuracy by timestep:

Fig.9 shows the average prediction accuracy as a function of the observation
index for our best TDB$(S=3,M=1)$, reported in Sec.4.2 for the $2$D aliased
room with $O=4$ observations.

Because of aliasing, at the start of a test random walk, the model cannot know
its location in the $2$D room: the prediction accuracy for the first timesteps
is low. The agent then quickly locates itself and test accuracy rapidly
increases: it is above $99.50\%$ after only $30$ observations.

Figure 9: After a short context of $30$ observations, TDB$(S=3,M=4)$ predicts
the next observations $99.50\%$ of the time.

Some latent codes model early uncertainty while others represent late
confidence:  Fig.10 displays, for each latent code active on the test data of
a TDB$(S=3,M=1)$ (a) the average timestep at which this latent code is
active—on the horizontal axis—and (b) its number of appearances on the entire
test data—on the vertical axis. In addition, each bottleneck is colored in
green (resp. red) to indicate that it is retained (resp. discarded) when we
threshold the empirical count tensor $C$ to build the cognitive map—as
detailed in Sec.3.6.

Figure 10: A first cluster of latent indices, in red, has low appearance
counts and models the early agent’s uncertainty in the test sequences. A
second cluster, in green, has high appearance counts and represents the
agent’s confidence when it knows its location.

The figure displays two clusters, which highlight the two roles played by the
latent codes. The first cluster represents a large number of codes with low
appearance counts, which appear early in the test sequences. These codes model
the agent’s uncertainty when it does not know its location.

The latent codes in the second cluster have a high number of appearances and
appear later in the test sequences. They represent the agent’s high confidence
when it knows its locations. In addition, because we use a high count
threshold $t_{\text{ratio}}=0.1$ (see Sec.3.6) to build the cognitive map,
only the transitions between the latent codes of this second cluster are used
to build the graph. As a result, a vast majority of the nodes of the second
(resp. first) cluster are green (resp. red). Consequently, our third
clustering step, detailed in Appendix C, maps the discarded (red) latent
indices to their closest retained (green) latent index.

The agent can “teleportate”:  We finally study how TDB$(S=3,M=1)$ moves in the
learned latent graph. In particular, we say that the agent “teleportates” when
the node of the transition graph where it estimates its location at the $n$th
timestep is not a neighbor of its estimated position at the $n-1$th timestep.
That is, teleportation means that the shortest path between two consecutive
nodes is larger than one.

Fig.11 looks at the average number of teleportation as a function of the
observation index. Around the $20$th observation, the agent still teleportates
$5\%$ of the time. As the agent becomes confident about its location, the
teleportation rate drops below $1\%$ after $100$ observations and stays around
$0.75\%$.

Interestingly, Fig.9[right] shows a subsequence of $10$ observations
$(x_{10},\ldots,x_{19})$ on a test random walks, with perfect next observation
accuracy. Nonetheless, the agent still “teleportates” five times out of ten.

Here, teleportation is an artifact of the clustering procedures that we use
when we build the cognitive map—see Appendix. C. For instance, our third
clustering step may incorrectly map some discarded bottleneck indices to their
closest bottleneck index. As a consequence, when the agent activates a
discarded bottleneck index, it may be mapped to an invalid room location. If
we use a lower threshold $t_{\text{ratio}}$, then teleportations should
disappear but the learned latent graph would not accurately model the
environment’s dynamics.

Figure 11: [Left] The agent frequently teleportates, i.e., moves between nodes
that are not neighbors early in the test sequences. Teleportation rate drops
below $5\%$ after $20$ observations, and below $1\%$ after $100$ observations.
[Right] Despite perfectly predicting the next observation between indices $10$
and $20$, the agent also teleportates five steps out of ten. Here,
teleportation is explained by incorrect clustering when we build the cognitive
map.

##  Appendix J Additional materials for simulated 3D environment

#### Table of results:

We first present the table of results for the simulated $3$D environment
experiments in Sec.4.3.

Method |  TestAccu ($\%$) $\uparrow$ |  ImpFallback ($\%$) $\uparrow$ |  RatioSP $\downarrow$ |  NormGED $\downarrow$
---|---|---|---|---
Vanilla transformer |  $\mathbf{99.80}\leavevmode\nobreak\ (0.00)$ |  $26.60\leavevmode\nobreak\ (0.58)$ |  $16.38\leavevmode\nobreak\ (0.87)$ |  —
Vanilla LSTM |  $99.79\leavevmode\nobreak\ (0.00)$ |  $26.85\leavevmode\nobreak\ (0.58)$ |  $16.38\leavevmode\nobreak\ (0.87)$ |  —
TDB$(S=1,\operatorname*{enc},M=1)$ |  $72.88\leavevmode\nobreak\ (2.86)$ |  $2.85\leavevmode\nobreak\ (1.33)$ |  $1.32\leavevmode\nobreak\ (0.21)$ |  $0.733\leavevmode\nobreak\ (0.082)$
TDB$(S=3,M=1)$ |  $58.51\leavevmode\nobreak\ (1.43)$ |  $1.35\leavevmode\nobreak\ (0.27)$ |  $1.35\leavevmode\nobreak\ (0.13)$ |  $0.946\leavevmode\nobreak\ (0.007)$
TDB$(S=1,\operatorname*{enc},M=4)$ |  $99.79\leavevmode\nobreak\ (0.01)$ |  $97.55\leavevmode\nobreak\ (0.42)$ |  $1.03\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.005}\leavevmode\nobreak\ (0.001)$
TDB$(S=3,M=4)$ |  $99.79\leavevmode\nobreak\ (0.00)$ |  $\mathbf{99.95}\leavevmode\nobreak\ (0.05)$ |  $\mathbf{1.02}\leavevmode\nobreak\ (0.00)$ |  $0.017\leavevmode\nobreak\ (0.004)$

Table 7: Metrics averaged over $10$ simulated $3$D environments. Training a
TDB with a single discrete bottleneck does not converge. In contrast, a TDB
with multiple discrete bottlenecks can almost perfectly (a) predict the next
observation given history, (b) solve shortest paths problems and (c) learn
cognitive maps nearly isomorphic to the ground truth map.

#### Latent graphs visualizations:

Second, Fig.12 presents the isometric views of four $3$D simulated rooms, as
well as the respective cognitive maps learned with a TDB($S=3,M=4$).

Figure 12: [Top] Isometric views of four $3$D simulated rooms used in Sec.4.3.
[Bottom] Corresponding cognitive maps learned with a TDB($S=3,M=4$): each
location in the room is represented by four nodes in the latent graph,
corresponding to the four possible heading directions of the agent.

##  Appendix K Additional materials for GINC dataset

Table 8 presents numerical results associated with Fig.4[left].

Context length |  Num. of examples |  LSTM |  Transformer |  TDB$(S=3,M=4)$ |  TDB$(S=5,M=1)$
---|---|---|---|---|---
$3$ |  $0$ |  $0.426\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.502}\leavevmode\nobreak\ (0.01)$ |  $0.426\leavevmode\nobreak\ (0.01)$ |  $0.284\leavevmode\nobreak\ (0.009)$
|  $1$ |  $0.369\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.488}\leavevmode\nobreak\ (0.01)$ |  $0.464\leavevmode\nobreak\ (0.01)$ |  $0.369\leavevmode\nobreak\ (0.01)$
|  $2$ |  $0.387\leavevmode\nobreak\ (0.01)$ |  $0.538\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.554}\leavevmode\nobreak\ (0.01)$ |  $0.443\leavevmode\nobreak\ (0.01)$
|  $4$ |  $0.394\leavevmode\nobreak\ (0.01)$ |  $0.553\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.588}\leavevmode\nobreak\ (0.01)$ |  $0.526\leavevmode\nobreak\ (0.01)$
|  $8$ |  $0.386\leavevmode\nobreak\ (0.01)$ |  $0.595\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.618}\leavevmode\nobreak\ (0.01)$ |  $0.549\leavevmode\nobreak\ (0.01)$
|  $16$ |  $0.390\leavevmode\nobreak\ (0.01)$ |  $0.631\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.655}\leavevmode\nobreak\ (0.01)$ |  $0.574\leavevmode\nobreak\ (0.01)$
|  $32$ |  $0.380\leavevmode\nobreak\ (0.01)$ |  $0.633\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.670}\leavevmode\nobreak\ (0.009)$ |  $0.618\leavevmode\nobreak\ (0.01)$
|  $64$ |  $0.377\leavevmode\nobreak\ (0.01)$ |  $0.651\leavevmode\nobreak\ (0.01)$ |  $\mathbf{0.676}\leavevmode\nobreak\ (0.009)$ |  $0.646\leavevmode\nobreak\ (0.01)$
$5$ |  $0$ |  $0.692\leavevmode\nobreak\ (0.009)$ |  $\mathbf{0.805}\leavevmode\nobreak\ (0.008)$ |  $0.777\leavevmode\nobreak\ (0.008)$ |  $0.656\leavevmode\nobreak\ (0.01)$
|  $1$ |  $0.742\leavevmode\nobreak\ (0.009)$ |  $0.822\leavevmode\nobreak\ (0.008)$ |  $\mathbf{0.853}\leavevmode\nobreak\ (0.007)$ |  $0.768\leavevmode\nobreak\ (0.008)$
|  $2$ |  $0.754\leavevmode\nobreak\ (0.009)$ |  $0.831\leavevmode\nobreak\ (0.008)$ |  $\mathbf{0.857}\leavevmode\nobreak\ (0.007)$ |  $0.802\leavevmode\nobreak\ (0.008)$
|  $4$ |  $0.749\leavevmode\nobreak\ (0.009)$ |  $0.841\leavevmode\nobreak\ (0.007)$ |  $\mathbf{0.905}\leavevmode\nobreak\ (0.006)$ |  $0.872\leavevmode\nobreak\ (0.007)$
|  $8$ |  $0.761\leavevmode\nobreak\ (0.009)$ |  $0.847\leavevmode\nobreak\ (0.007)$ |  $\mathbf{0.901}\leavevmode\nobreak\ (0.006)$ |  $0.882\leavevmode\nobreak\ (0.006)$
|  $16$ |  $0.756\leavevmode\nobreak\ (0.009)$ |  $0.871\leavevmode\nobreak\ (0.007)$ |  $\mathbf{0.922}\leavevmode\nobreak\ (0.005)$ |  $0.908\leavevmode\nobreak\ (0.006)$
|  $32$ |  $0.767\leavevmode\nobreak\ (0.008)$ |  $0.871\leavevmode\nobreak\ (0.007)$ |  $\mathbf{0.926}\leavevmode\nobreak\ (0.005)$ |  $0.915\leavevmode\nobreak\ (0.006)$
|  $64$ |  $0.758\leavevmode\nobreak\ (0.009)$ |  $0.88\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.924}\leavevmode\nobreak\ (0.005)$ |  $0.917\leavevmode\nobreak\ (0.006)$
$8$ |  $0$ |  $0.849\leavevmode\nobreak\ (0.007)$ |  $0.881\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.895}\leavevmode\nobreak\ (0.006)$ |  $0.816\leavevmode\nobreak\ (0.008)$
|  $1$ |  $0.869\leavevmode\nobreak\ (0.007)$ |  $0.888\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.930}\leavevmode\nobreak\ (0.005)$ |  $0.899\leavevmode\nobreak\ (0.006)$
|  $2$ |  $0.869\leavevmode\nobreak\ (0.007)$ |  $0.886\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.930}\leavevmode\nobreak\ (0.005)$ |  $0.917\leavevmode\nobreak\ (0.006)$
|  $4$ |  $0.877\leavevmode\nobreak\ (0.007)$ |  $0.897\leavevmode\nobreak\ (0.006)$ |  $0.933\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.936}\leavevmode\nobreak\ (0.005)$
|  $8$ |  $0.873\leavevmode\nobreak\ (0.007)$ |  $0.897\leavevmode\nobreak\ (0.006)$ |  $0.935\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.940}\leavevmode\nobreak\ (0.005)$
|  $16$ |  $0.875\leavevmode\nobreak\ (0.007)$ |  $0.916\leavevmode\nobreak\ (0.006)$ |  $0.942\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.950}\leavevmode\nobreak\ (0.004)$
|  $32$ |  $0.862\leavevmode\nobreak\ (0.007)$ |  $0.909\leavevmode\nobreak\ (0.006)$ |  $0.948\leavevmode\nobreak\ (0.004)$ |  $\mathbf{0.956}\leavevmode\nobreak\ (0.004)$
|  $64$ |  $0.883\leavevmode\nobreak\ (0.006)$ |  $0.920\leavevmode\nobreak\ (0.005)$ |  $0.952\leavevmode\nobreak\ (0.004)$ |  $\mathbf{0.955}\leavevmode\nobreak\ (0.004)$
$10$ |  $0$ |  $0.878\leavevmode\nobreak\ (0.007)$ |  $0.888\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.915}\leavevmode\nobreak\ (0.006)$ |  $0.858\leavevmode\nobreak\ (0.007)$
|  $1$ |  $0.889\leavevmode\nobreak\ (0.006)$ |  $0.898\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.930}\leavevmode\nobreak\ (0.005)$ |  $0.926\leavevmode\nobreak\ (0.005)$
|  $2$ |  $0.878\leavevmode\nobreak\ (0.007)$ |  $0.901\leavevmode\nobreak\ (0.006)$ |  $\mathbf{0.938}\leavevmode\nobreak\ (0.005)$ |  $0.935\leavevmode\nobreak\ (0.005)$
|  $4$ |  $0.902\leavevmode\nobreak\ (0.006)$ |  $0.904\leavevmode\nobreak\ (0.006)$ |  $0.936\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.945}\leavevmode\nobreak\ (0.005)$
|  $8$ |  $0.881\leavevmode\nobreak\ (0.006)$ |  $0.906\leavevmode\nobreak\ (0.006)$ |  $0.947\leavevmode\nobreak\ (0.004)$ |  $\mathbf{0.956}\leavevmode\nobreak\ (0.004)$
|  $16$ |  $0.882\leavevmode\nobreak\ (0.006)$ |  $0.908\leavevmode\nobreak\ (0.006)$ |  $0.934\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.947}\leavevmode\nobreak\ (0.004)$
|  $32$ |  $0.889\leavevmode\nobreak\ (0.006)$ |  $0.918\leavevmode\nobreak\ (0.005)$ |  $0.946\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.960}\leavevmode\nobreak\ (0.004)$
|  $64$ |  $0.89\leavevmode\nobreak\ (0.006)$ |  $0.911\leavevmode\nobreak\ (0.006)$ |  $0.941\leavevmode\nobreak\ (0.005)$ |  $\mathbf{0.956}\leavevmode\nobreak\ (0.004)$

Table 8: In-context accuracy for vanilla sequence models and TDBs, for each
each pair $(k,n)$ of context length $k$ and number of examples $n$ of the GINC
test set. Our proposed TDBs consistently reach the highest in-context
accuracies.

##  Appendix L Additional materials for in-context learning

###  L.1 Defining the in-context learning problem

#### Preserving the room partition:

We consider a first $2$D aliased room of size $6\times 8$ with $O=4$
observations, $c_{1},\ldots c_{O}$ such that $c_{i}\in\\{1,\ldots,30\\}$. We
represent this room by a matrix $R\in\\{c_{1},\ldots c_{O}\\}^{6\times 8}$,
where, for $x\in\\{1,\ldots,6\\}$ and $y\in\\{1,\ldots,8\\}$, $R_{xy}$ is the
room observation at the spatial position $(x,y)$.

We define the room partition induced by the room colors as the partition
$\mathcal{S}_{1},\ldots,\mathcal{S}_{O}$ of the $2$D room spatial positions
such that, for each element $\mathcal{S}_{i}$ of the partition, all the
spatial positions in $\mathcal{S}_{i}$ have the same observation.

Formally, the room partition is defined as

$\bigcup_{i=1}^{O}\mathcal{S}_{i}=\\{(x,y):\leavevmode\nobreak\ 1\leq x\leq
6,\leavevmode\nobreak\ 1\leq y\leq 8\\}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \text{;}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\
\mathcal{S}_{i}\cap\mathcal{S}_{j}=\emptyset,\leavevmode\nobreak\ \forall
i\neq j\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\text{;}\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\mathcal{S}_{i}\neq\emptyset,\leavevmode\nobreak\ \forall i,$

and satisfies

$\forall i\leq O,\leavevmode\nobreak\
\forall(x,y)\in\mathcal{S}_{i},\leavevmode\nobreak\ R_{xy}=o_{i}.$

For each injective mapping $\phi:\\{1,\ldots,O\\}\to\\{1,\ldots,30\\}$, we can
build a new room $\tilde{R}$ with (a) observations $\phi(1),\ldots,\phi(O)$
and which (b) preserves the room partition, by assigning all the $2$D spatial
positions in $\mathcal{S}_{i}$ to the observation $\phi(i)$. That is

$\forall i\leq O,\leavevmode\nobreak\
\forall(x,y)\in\mathcal{S}_{i},\leavevmode\nobreak\ \tilde{R}_{xy}=\phi(i).$

The training sets in Sec.4.5 contain at most $10$k such training rooms, each
one preserving the room partition. Note that, for a given partition, the
number of possible rooms which preserves the room partition is the number of
injective mapping from $\\{1,\ldots,O\\}$ to $\\{1,\ldots,30\\}$, which is
$\prod_{i=0}^{O-1}30-i$.

#### Learning in base:

As discussed in the main text, given a sequence of observations
$x=(x_{1},\ldots,x_{N})$ in a room defined as above, each model is trained to
predict targets in the room-agnostic base, $\tilde{x}_{n}\in\\{1,\ldots,O\\}$,
such that $\tilde{x}_{n}=k$ iff. $\tilde{x}_{n}$ is the $k$th lowest
observation seen between indices $1$ and $n$.

In particular, let us define the permutation $i_{1},\ldots,i_{O}$ of
$1,\ldots,O$ such that $\phi(i_{1})\leq\ldots\leq\phi(i_{O})$. Let $n^{*}$ be
such a large enough integer such that all the $O$ different room observations
have been seen between $x_{1}$ and $x_{n^{*}}$. Then, for $n\geq n^{*}$, the
target $\tilde{x}_{n}=k$ corresponds to the $k$th highest color of the room,
which is equal to $\phi(i_{k})$.

###  L.2 Table of results

Table 9 reports the numerical results displayed in Fig. 5[left].

Metric |  Num. training rooms |  LSTM |  Transformer |  TDB$(S=3,M=4)$ |  TDB$(S=1,\operatorname*{enc},M=4)$
---|---|---|---|---|---
TestAccu ($\%$) $\uparrow$ |  $200$ |  $\mathbf{0.508}\leavevmode\nobreak\ (0.009)$ |  $0.426\leavevmode\nobreak\ (0.005)$ |  $0.38\leavevmode\nobreak\ (0.004)$ |  $0.379\leavevmode\nobreak\ (0.002)$
|  $500$ |  $\mathbf{0.741}\leavevmode\nobreak\ (0.039)$ |  $0.595\leavevmode\nobreak\ (0.006)$ |  $0.461\leavevmode\nobreak\ (0.009)$ |  $0.481\leavevmode\nobreak\ (0.007)$
|  $1000$ |  $\mathbf{0.899}\leavevmode\nobreak\ (0.006)$ |  $0.864\leavevmode\nobreak\ (0.016)$ |  $0.568\leavevmode\nobreak\ (0.015)$ |  $0.58\leavevmode\nobreak\ (0.03)$
|  $2000$ |  $0.919\leavevmode\nobreak\ (0.009)$ |  $\mathbf{0.977}\leavevmode\nobreak\ (0.003)$ |  $0.868\leavevmode\nobreak\ (0.03)$ |  $0.8\leavevmode\nobreak\ (0.056)$
|  $5000$ |  $0.848\leavevmode\nobreak\ (0.045)$ |  $\mathbf{0.985}\leavevmode\nobreak\ (0.001)$ |  $\mathbf{0.974}\leavevmode\nobreak\ (0.001)$ |  $0.784\leavevmode\nobreak\ (0.091)$
|  $10000$ |  $0.924\leavevmode\nobreak\ (0.007)$ |  $\mathbf{0.985}\leavevmode\nobreak\ (0.001)$ |  $0.977\leavevmode\nobreak\ (0.001)$ |  $0.912\leavevmode\nobreak\ (0.059)$
ImpFallback ($\%$) $\uparrow$ |  $200$ |  $0.005\leavevmode\nobreak\ (0.001)$ |  $0.002\leavevmode\nobreak\ (0.001)$ |  $\mathbf{0.038}\leavevmode\nobreak\ (0.005)$ |  $0.032\leavevmode\nobreak\ (0.004)$
|  $500$ |  $0.148\leavevmode\nobreak\ (0.04)$ |  $0.007\leavevmode\nobreak\ (0.002)$ |  $0.05\leavevmode\nobreak\ (0.004)$ |  $\mathbf{0.064}\leavevmode\nobreak\ (0.004)$
|  $1000$ |  $0.382\leavevmode\nobreak\ (0.054)$ |  $0.189\leavevmode\nobreak\ (0.032)$ |  $0.115\leavevmode\nobreak\ (0.014)$ |  $\mathbf{0.124}\leavevmode\nobreak\ (0.016)$
|  $2000$ |  $0.447\leavevmode\nobreak\ (0.059)$ |  $0.725\leavevmode\nobreak\ (0.024)$ |  $\mathbf{0.690}\leavevmode\nobreak\ (0.067)$ |  $0.495\leavevmode\nobreak\ (0.092)$
|  $5000$ |  $0.312\leavevmode\nobreak\ (0.086)$ |  $0.809\leavevmode\nobreak\ (0.009)$ |  $\mathbf{0.964}\leavevmode\nobreak\ (0.011)$ |  $0.605\leavevmode\nobreak\ (0.132)$
|  $10000$ |  $0.305\leavevmode\nobreak\ (0.072)$ |  $0.771\leavevmode\nobreak\ (0.011)$ |  $\mathbf{0.970}\leavevmode\nobreak\ (0.009)$ |  $0.782\leavevmode\nobreak\ (0.084)$
RatioSP $\downarrow$ |  $200$ |  $51.79\leavevmode\nobreak\ (20.83)$ |  $34.79\leavevmode\nobreak\ (11.07)$ |  $1.46\leavevmode\nobreak\ (0.09)$ |  $\mathbf{1.33}\leavevmode\nobreak\ (0.06)$
|  $500$ |  $29.73\leavevmode\nobreak\ (3.67)$ |  $45.26\leavevmode\nobreak\ (23.87)$ |  $\mathbf{1.17}\leavevmode\nobreak\ (0.05)$ |  $1.27\leavevmode\nobreak\ (0.06)$
|  $1000$ |  $25.52\leavevmode\nobreak\ (1.19)$ |  $26.14\leavevmode\nobreak\ (1.34)$ |  $\mathbf{1.18}\leavevmode\nobreak\ (0.04)$ |  $1.20\leavevmode\nobreak\ (0.04)$
|  $2000$ |  $27.01\leavevmode\nobreak\ (0.89)$ |  $25.98\leavevmode\nobreak\ (0.83)$ |  $\mathbf{1.07}\leavevmode\nobreak\ (0.03)$ |  $1.11\leavevmode\nobreak\ (0.04)$
|  $5000$ |  $24.79\leavevmode\nobreak\ (0.96)$ |  $24.53\leavevmode\nobreak\ (0.63)$ |  $\mathbf{1.01}\leavevmode\nobreak\ (0.0)$ |  $1.06\leavevmode\nobreak\ (0.05)$
|  $10000$ |  $24.27\leavevmode\nobreak\ (1.73)$ |  $25.13\leavevmode\nobreak\ (0.45)$ |  $\mathbf{1.01}\leavevmode\nobreak\ (0.0)$ |  $\mathbf{1.01}\leavevmode\nobreak\ (0.0)$
NormGED $\downarrow$ |  $200$ |  — |  — |  $\mathbf{0.707}\leavevmode\nobreak\ (0.005)$ |  $0.723\leavevmode\nobreak\ (0.006)$
|  $500$ |  — |  — |  $\mathbf{0.658}\leavevmode\nobreak\ (0.006)$ |  $0.660\leavevmode\nobreak\ (0.008)$
|  $1000$ |  — |  — |  $\mathbf{0.579}\leavevmode\nobreak\ (0.018)$ |  $0.602\leavevmode\nobreak\ (0.033)$
|  $2000$ |  — |  — |  $\mathbf{0.164}\leavevmode\nobreak\ (0.042)$ |  $0.29\leavevmode\nobreak\ (0.073)$
|  $5000$ |  — |  — |  $\mathbf{0.019}\leavevmode\nobreak\ (0.003)$ |  $0.294\leavevmode\nobreak\ (0.121)$
|  $10000$ |  — |  — |  $\mathbf{0.017}\leavevmode\nobreak\ (0.002)$ |  $0.118\leavevmode\nobreak\ (0.079)$

Table 9: In-context metrics for vanilla sequence models and TDBs on the new
test rooms. For a large number of training rooms, our best TDB$(S=3,M=4)$ (a)
reaches nearly perfect in-context accuracy, close to the best performing
transformer, (b) almost perfectly solves the in-context path planning
problems, and (c) learns in-context latent graphs nearly isomorphic to the
ground truth.

First, we observe that for each model, in-context capacities emerge when we
increase the number of rooms.

Second, while vanilla LSTM performs best for a small number of training rooms,
it cannot reach near-perfect in-context accuracies when the number of training
rooms increases—which a vanilla transformer can do. We believe that attention-
like mechanisms are important for in-context capacities to fully emerge in
novel test rooms.

Third, as observed in other experiments, vanilla sequence models cannot solve
path planning problems. Indeed, for a large number of training rooms, a
transformer improves at best $77\%$ of the times over the fallback path.
However, when it does so, it finds paths that are on average $25$ times longer
than the optimal paths. In contrast, our best TDB$(S=3,M=4)$ (a) almost
matches the nearly perfect in-context predictive performance of a vanilla
transformer, (b) almost perfectly solves in-context path planning problems and
(c) learns in-context latent graphs nearly isomorphic to the ground truth.

Finally, our TDB$(S=3,\operatorname*{enc},M=4)$—which also predicts the next
latent encoding—performs worse than the other models. This is explained by its
large variances across the different experiments. In fact, we observe that for
two out of ten experiments, its training accuracy remains very low—below
$40\%$—as the model struggles to learn to predict the next observation. For
the remaining eight runs, prediction and path planning performance are near
optimal and the model competes with our best TDB$(S=3,M=4)$. We could try
increasing the number of bottlenecks and training a
TDB$(S=3,\operatorname*{enc},M=16)$ to mitigate this issue.

###  L.3 In-context accuracy by timestep

Fig.13 shows the average in-context accuracy on the new test rooms, as a
function of the observation index, for our best TDB$(S=3,M=4)$ trained on $5$k
rooms. Interestingly, in-context accuracy drops after a few iterations.
Indeed, as the model sees new observations in the new test rooms, the mapping
between the base labels $\tilde{x}_{n}$ and the observations $x_{n}$ may
change—which may confuse the model. After a few tenths of iterations, TDB is
able to locate itself in the new test room. After $70$ iterations, in-context
accuracy becomes higher than $99\%$.

Figure 13: After $70$ timesteps, in-context accuracy for a TDB$(S=3,M=4)$
trained on $5$k training rooms becomes higher than $99\%$.

###  L.4 In-context learning is driven by spatial exposure to base targets in
the training data

Restricting training rooms:  For this experiment, we build training rooms such
that the $k$th highest color (almost) never appears at any of the spatial
positions indicated by the set $\mathcal{S}_{k}$ of the room partition. With
the notations of Appendix L.1, each injective mapping $\phi$ used to generate
a training room has to satisfy the rule

$\mathbf{\mathcal{R}}=\\{\forall k\leq O:\leavevmode\nobreak\ i_{k}\neq k\\}.$

For instance, the mapping $\phi(1)=5,\leavevmode\nobreak\
\phi(2)=13,\leavevmode\nobreak\ \phi(3)=20,\leavevmode\nobreak\ \phi(4)=8$
cannot be used to generate a training room. Indeed, when we rank the entries
we get $i_{1}=1,\leavevmode\nobreak\ i_{2}=4,\leavevmode\nobreak\
i_{3}=2,\leavevmode\nobreak\ i_{4}=3$ and $i_{1}=1$ violates
$\mathbf{\mathcal{R}}$. Similarly $\phi(1)=20,\leavevmode\nobreak\
\phi(2)=2,\leavevmode\nobreak\ \phi(3)=15,\leavevmode\nobreak\ \phi(4)=11$
violates $\mathbf{\mathcal{R}}$ as $i_{3}=3$.

Given a training sequence $x$, let $n^{*}$ be such a large enough integer such
that all the $O$ different room observations have been seen between $x_{1}$
and $x_{n^{*}}$. For $n\geq n^{*}$, let $k\leq O$ be such that the spatial
position $\operatorname*{pos}_{n}$ of the agent at timestep $n$ belongs to the
set $\mathcal{S}_{k}$ of the room partition. Then, by definition, the base
target $\tilde{x}_{n}$ satisfies $\tilde{x}_{n}\neq k$. Note that, by
definition, when $n\leq n^{*}$ and we have not yet observed all the $O$
different room observations, we may have $\tilde{x}_{n}=k$ when
$\operatorname*{pos}_{n}\in\mathcal{S}_{k}$.

Consequently, for each room spatial position, there is a position-specific
base target that the agent will never see (when $n\geq n^{*}$) at this
position on the training set.

We generate $5$k such training rooms and train a TDB$(S=3,M=4)$ as before.

In-context accuracies:  We consider four families of test rooms:
(A): Test rooms where $\mathbf{\mathcal{R}}$ is respected.
(B): Test rooms where $\mathbf{\mathcal{R}}$ is violated once.
(C): Test rooms where $\mathbf{\mathcal{R}}$ is violated twice.
(D): Test rooms where $\mathbf{\mathcal{R}}$ is violated four times. These are
the rooms satisfying $\phi(0)\leq\phi(1)\leq\phi(2)\leq\phi(3)$.

On test rooms (A), in-context accuracy is $98.05\%(\pm 0.07\%)$.
On test rooms (B), in-context accuracy is $56.97\%(\pm 0.82\%)$.
On test rooms (C), in-context accuracy is $50.25\%(\pm 0.51\%)$.
On test rooms (D), in-context accuracy is $46.20\%(\pm 0.73\%)$.

This drastic drop in in-context accuracy demonstrates that—in addition to the
number of training rooms—in-context capacities are also driven by spatial
exposure to base targets during training.