<EMAIL_ADDRESS>
# Graph schemas as abstractions for transfer learning, inference, and planning

J. Swaroop Guntupalli  Google DeepMind  Rajkumar Vasudeva Raju  Google
DeepMind  Shrinu Kushagra  Google DeepMind  Carter Wendelken  Google DeepMind
Danny Sawyer  Google DeepMind  Ishan Deshpande  Google Research  Guangyao Zhou
Google DeepMind  Miguel Lázaro-Gredilla  Google DeepMind  Dileep George
Google DeepMind

###### Abstract

Transferring latent structure from one environment or problem to another is a
mechanism by which humans and animals generalize with very little data.
Inspired by cognitive and neurobiological insights, we propose graph schemas
as a mechanism of abstraction for transfer learning. Graph schemas start with
latent graph learning where perceptually aliased observations are
disambiguated in the latent space using contextual information. Latent graph
learning is also emerging as a new computational model of the hippocampus to
explain map learning and transitive inference. Our insight is that a latent
graph can be treated as a flexible template — a schema — that models concepts
and behaviors, with slots that bind groups of latent nodes to the specific
observations or groundings. By treating learned latent graphs (schemas) as
prior knowledge, new environments can be quickly learned as compositions of
schemas and their newly learned bindings. We evaluate graph schemas on two
previously published challenging tasks: the memory & planning game and one-
shot StreetLearn, which are designed to test rapid task solving in novel
environments. Graph schemas can be learned in far fewer episodes than previous
baselines, and can model and plan in a few steps in novel variations of these
tasks. We also demonstrate learning, matching, and reusing graph schemas in
more challenging 2D and 3D environments with extensive perceptual aliasing and
size variations, and show how different schemas can be composed to model
larger and more complex environments. To summarize, our main contribution is a
unified system, inspired and grounded in cognitive science, that facilitates
rapid transfer learning of new environments using schemas via map-induction
and composition that handles perceptual aliasing.

##  1 Introduction

Discovering and using the right abstractions in new situations affords
efficient transfer learning as well as quick inference and planning. Humans
excel at this ability, argued to be a key factor behind intelligence and a
fundamental limitation in current AI systems (Shanahan and Mitchell, 2022).
Common reusable structured representations of concepts or
behaviors—schemas—have been proposed as a powerful way to encode abstractions
(Mitchell, 2021; Tenenbaum et al., 2011). Having a computational model with
the ability to discover and reuse previously-learned schemas to behave and
plan in novel situations will be essential for AI.

Experimental evidence suggests that several animals have this ability
(Farzanfar et al., 2023). Rats and mice tend to learn new environments faster
if they can reuse past schemas (Tse et al., 2007; Zhou et al., 2021), and
macaque hippocampus cells encode spatial schemas (Baraduc et al., 2019).
Neural circuits in the hippocampus and prefrontal cortex (PFC) are implicated
in schema learning, recognition, update, and maintenance and these processes
are considered as the foundation of memory consolidation (Preston and
Eichenbaum, 2013; Gilboa and Marlatte, 2017; Samborska et al., 2022). New
experiences are learned in a single trial if it fits an existing schema. An
updated complementary learning systems theory based on this evidence was
proposed in Kumaran et al. (2016), but there is so far no explicit
demonstration of such rapid learning with schema reuse as far as we know.

Structured relational representations have been proposed as a common mechanism
in the hippocampus that reconciles spatial and non-spatial tasks and memory
into the original cognitive mapping view (Eichenbaum and Cohen, 2014;
Stachenfeld et al., 2017). Several recent studies model cognitive maps as
higher order latent graph structures and show generalization to disparate
functions implicated in the hippocampus (George et al., 2021; Raju et al.,
2022; Sharma et al., 2021; Whittington et al., 2020, 2021). We take one such
model of cognitive maps, the clone-structured cognitive graph (CSCG)(George et
al., 2021), and extend it to provide a concrete computational model of
abstractions using graph schemas. Desiderata for our model is a unified system
that facilitates learning new environments by using schemas, handles
perceptual aliasing (Whitehead and Ballard, 1991), and generalizes via map
induction (Sharma et al., 2021) and schema composition.

Our setup is an agent navigating in an environment modeled as a directed
graph. The agent observes the emissions at its current node and traverses to
new nodes via edges with labeled actions. Multiple nodes may emit the same
observations (i.e., they are aliased (Lajoie et al., 2019; Whitehead and
Ballard, 1991)), so the agent cannot observe the state directly. This can be
considered as a discrete time partially observable Markov decision process
(POMDP). As the agent navigates a new environment, our goal is to learn the
underlying latent graph (i.e., map induction Sharma et al. (2021)) and to do
so rapidly by reusing previously learnt graph topologies or graph schemas. We
choose the CSCG model to construct graph schemas as it has been shown to learn
higher-order graphs in highly aliased settings using a smooth, probabilistic,
parameterization of the graph learning problem (George et al., 2021). We
extend this model to describe how learned graphs can be reused as schemas for
transfer learning, quick inference, and planning for behavior in new
situations by rapidly learning observation bindings and discovering the best
schema online.

##  2 Related Work

Research on rapid transfer to new tasks in novel environments has focused on
different aspects, from exploration to modeling and planning. Some of the
recent work has been predominantly done in a reinforcement learning (RL)
framework, and different RL approaches focus on one or more of these aspects.

Meta-RL Model-free meta-RL approaches focus on exploration policy
generalization to new tasks and environmental variations without explicit
model learning. Some show generalization to new tasks in a known environment,
but not to new environments nor do they handle aliasing (Rakelly et al., 2019;
Wang et al., 2016). Recurrent model-free RL has shown some generalization to
environmental variations and in the POMDP settings (Ni et al., 2021). These
variations are of parameters for generating the environment, and the goal is
to be able learn an exploration policy that generalizes to these parametric
variations. These methods do not re-use explicit knowledge of past
environments to model new environments nor do they handle composition of known
environments (Packer et al., 2018). More recent works added episodic memory
with attention heads to selectively attend and reuse stored memories to
rapidly adapt to tasks in new environments in few-shot settings (e.g. episodic
planning network)(Lampinen et al., 2021; Ritter et al., 2020). These models
match optimal planning only after training for billions of steps. Further,
none of these approaches handle aliasing, nor explicitly build models of the
environment and plan over them. Gupta et al. (2017) show that using explicit
model building (mapping) and navigation via planning in spatial environments
can handle partial observability and outperform methods without this ability,
but this study does not reuse those models to rapidly learn novel
environments.

Model based RL Model based RL works learn an explicit model of the environment
(Gregor et al., 2019), some even under a POMDP setup (Igl et al., 2018), and
can transfer to new tasks within that environment, but cannot transfer to
novel environments. Some studies use a version of sub-goal or landmark
discovery and reuse to generalize to novel tasks (Eysenbach et al., 2019; Kipf
et al., 2019; Zhang et al., 2021). However, these approaches do not address
partial observability, and more importantly the generalization is restricted
to the same environments. Semi-parametric topological memory is a method that
extends landmark based navigation to generalize to novel environments (Savinov
et al., 2018). This method trains a network to estimate connectivity between
observations at different locations and uses this to build a graph of
observations in a novel environment. This method does generalize to novel
environments, but its performance degrades significantly with increasing
aliasing. Further, a human-generated exploration path was used for
constructing the graph.

Our work on schema matching and reuse is related to finding correspondence
between graphs in different contexts. The first neural network approach to
structure mapping was proposed in Crouse et al. (2021). But this approach is
restricted to the matching problem and has no mechanism to resolve or learn
new structures or to plan with partially matching schemas. Another line of
related work focused on solving simplified relational tasks inspired by
Raven’s Progressive Matrices (Kerg et al., 2022; Webb et al., 2021). The main
idea is to separate abstract relations from sensory observations during
training, and learn the observation mapping to solve new tasks with the same
relations but novel mapping, but in a deterministic and simplified setting.

The main contribution of this paper is a unified system that (i) facilitates
rapid transfer learning of new environments using schemas via map-induction
and composition, (ii) handles perceptual aliasing, (iii) builds explicit
latent graph models of environments directly from actions and observations,
(iv) supports planning and inference, and (v) and is inspired and grounded in
cognitive and neurosciences. Ours is the first model to combine all these
aspects in a single system.

Figure 1: Overview of using graph schemas in simulated 3D environments. A. An
example 3D environment. B. The agent navigates the environment with egocentric
actions and gets observations as RGB images. The images are passed through a
vector quantizer to obtain cluster indices, which are used as observations for
training a CSCG. C. Graphical model of continuous extension of CSCG with
action conditioning. D. Given a sequence of actions and observations, a CSCG
learns the transition graph of the environment with the corresponding emission
binding. We use simplified versions of the transition graphs in this
schematic. See Fig. S1 for real graphs. E. Schemas are extracted from learned
models with transitions and emissions by unbinding the specific emissions
(represented here as floor colors in the environment, and node colors in the
corresponding graph), but retaining the clone structure (shown as node
patterns). F. In a novel environment with size, color, and structure
variations, the agent navigates and finds the schema that best explains its
observations by learning new emission bindings.

##  3 Methods

###  3.1 Problem Setup

Consider an agent navigating in a directed graph $G$. When the agent visits a
node in the graph, the node emits an observation. However, multiple nodes may
emit the same observation (i.e., they are aliased), so the observation is not
enough to disambiguate where in the graph the agent is located. Additionally,
actions do not have deterministic results—executing the same action at the
same node might result in the agent navigating to different nodes. The
outgoing edges from a node are labeled with the action that would traverse
them, and with the probability of traversing them under that action. The sum
of the probabilities of all outgoing edges from a node with the same action
label sum up to 1. We use the graph $G$ to model the agent’s environment.

When the agent performs a sequence of actions $a_{1},\ldots,a_{N}$ (with
discrete $a_{n}\in\\{1,\ldots N_{\rm actions}\\}$), it receives a sequence of
observations $x_{1},\ldots,x_{N}$ (with each $x_{n}\in\mathbb{R}^{d}$ or
$\\{1,\ldots,N_{\rm obs}\\}$ for continuous and discrete observations,
respectively). The goal of learning is to recover the topology of the
environment $G$ from sequences of actions and observations. The goal of
transfer is to reuse the topology of a previously learned graph to model a new
environment by relabeling the nodes with new observations. In both cases, the
learned graphs can be exploited for tasks such as goal-directed navigation.

###  3.2 Model

Clone-structured cognitive graphs (CSCGs) were introduced in Dedieu et al.
(2019); George et al. (2021) to recover (an approximation of) the graph $G$
from sequences of action-observation pairs. To do this, they use categorical
hidden variables $z_{1},\ldots,z_{N}$ to model the node of the graph that the
agent is at in each time step. With this, it is possible to formulate a
graphical model for the sequence of observations given the actions. Here, we
use the conditional version of their model and extend it to continuous
observations:

$p(\boldsymbol{x}|\boldsymbol{a})=\sum_{\boldsymbol{z}}\sum_{\boldsymbol{y}}P(z_{1})\prod_{n=2}^{N}P(z_{n}|z_{n-1},a_{n-1})\prod_{n=1}^{N}P(y_{n}|z_{n})p(x_{n}|y_{n}),$
(1)

with $p(x_{n}|y_{n})=\mathcal{N}(x_{n}|\mu_{y_{n}},\sigma^{2}I)$. We use the
following shorthand for a sequence of actions:
$\boldsymbol{a}\equiv\\{a_{1}\ldots,a_{N}\\}$
($\boldsymbol{x},\boldsymbol{y},\boldsymbol{z}$ are similarly defined). The
transitions are fully parameterized through an action-conditional transition
tensor $T$ with elements $T_{ijk}=P(z_{n}=k|z_{n-1}=j,a_{n-1}=i)$.

To extend CSCG to continuous observations, we introduced a new variable
$y_{n}$ between the hidden state $z_{n}$ and the observation $x_{n}$. In this
formulation, the observation model is parameterized as an isotropic Gaussian
with variance $\sigma^{2}$ and mean $\mu_{y_{n}}$, which is the centroid
associated to a discrete emission $y_{n}$. The emission model is parameterized
by an emission matrix $E$ with elements $E_{ij}=P(y_{n}=j|z_{n}=i)$. In a
CSCG, by design, multiple hidden states are forced to share the same emission:
if states $i$ and $j$ are _clones_ of the same emission, then
$p(y|z_{i})=p(y|z_{j})$. Also, the emissions are deterministic: each row in
$E$ has a single entry set to 1, with the remaining elements set to zero.
Cloned rows (corresponding to the same emission) will be identical. The idea
is that different latent states can produce similar visible observations, and
the latent state can only be disambiguated through temporal context. We define
$C(x_{n})$ as the set of clones of the emission $x_{n}$.

In a simple case, if we have a finite number of potential observations, we can
allocate a centroid to each of them and set $\sigma^{2}\to 0$. When the
observations are discrete, the observation model becomes deterministic and we
can set $y=x$. This exactly recovers the discrete CSCG from George et al.
(2021). To learn the discrete CSCG, the clone structure is decided a priori by
choosing how many clones to assign to each observation (e.g., a constant
number). This fully defines $E$. Then, learning a graph $G$ that describes a
new environment only involves learning $T$, i.e. maximizing
$P(\boldsymbol{x}|\boldsymbol{a};T,E)$ w.r.t. $T$. This can be done via
expectation-maximization (EM), see Suppl. §A.1. The procedure is analogous to
HMM learning, but computationally more efficient by leveraging the sparsity
pattern of the emission matrix.

For continuous observations, see Suppl. §A.3 for a description of how to learn
the centroids $\mu_{y}$ and allocate hidden states to them (i.e., assigning
the clone structure).

###  3.3 Schemas for transfer learning

We call a learned model with fixed 3-tuple $S_{g}\equiv\\{T,C,E\\}$ a grounded
schema. Transfer learning between different environments that share the same
graph topology but with different emissions associated to its nodes is
possible by reusing the previously learned transitions $T$. We can maximize
$p(\boldsymbol{x}|\boldsymbol{a};T,\boldsymbol{\mu})$ w.r.t.
$\boldsymbol{\mu}$ (or $p(\boldsymbol{x}|\boldsymbol{a};T,E)$ w.r.t. $E$ in
the discrete case), while keeping $T$ fixed to its known value from the
previous environment (see Suppl. §A.2 and §A.4 for discrete and continuous
observations respectively). We call this $T$ without an associated clone
structure $C$ and the emission matrix $E$ as an ungrounded schema or simply a
schema.

Further, if we know that the new environment also preserves the clone
structure $C$ (when two nodes were clones in the original environment, they
are also clones in the new environment), then we can further restrict the
observation model to respect this constraint during learning and reuse the
tuple $S_{u}\equiv\\{T,C\\}$ and learn $\boldsymbol{\mu}$ or $E$. We call this
tuple $S_{u}$ an ungrounded schema with clone structure (Fig. 1E).

For example, in the room navigation setting, a schema models the agent’s
location and head-directions in a room and how actions move the agent, as well
as the knowledge that the floor or door can look the same at multiple
locations in the room. Using EM, we show how schemas allow for rapid model
learning in new environments that have matching topologies and emission
structures with fast binding (Fig. 1F). Inference can be performed with the
matched schemas to actively plan and pursue goals. We can also detect
transitions to another known schema, or to unknown territory by comparing
likelihoods of observations under different schemas.

Schemas can also be used as building blocks to rapidly learn novel
environments that are composed of matching topologies. This comprises learning
transitions and emissions but reusing known schemas where they fit (See Suppl.
§B.5 and Suppl. Alg.2).

##  4 Results

We show results from two sets of experiments: (i) standard benchmarks for
evaluating rapid adaptation, where the environments do not have perceptual
aliasing, and (ii) more challenging setup of environments with extensive
perceptual aliasing.

###  4.1 Rapid adaptation and task solving in novel environments

We first evaluate our model on two benchmarks proposed by Ritter et al. (2020)
to evaluate rapid adaptation and task solving in new environments: Memory &
Planning Game (MPG) and One-Shot StreetLearn.

####  4.1.1 Memory & Planning Game

In the MPG, the agent can navigate on a $4\times 4$ grid, observing symbols,
and the task is to collect reward at a specified goal location in that grid
(Fig. 2A). All grid positions have unique symbols and the symbol-position
mapping is randomized after each episode, which lasts for 100 steps. See
Suppl. §B.1 for details. This setup lets us evaluate our model on ground truth
graph recovery and schema reuse, as the structure is maintained across
episodes. The agent needs to explore to collect the observations and bind them
rapidly to the schema to maximize rewards with optimal navigation.

Figure 2: Performance on MPG and One-shot StreetLearn benchmarks. Example
episodes of MPG (A) and One-Shot StreetLearn (D) along with their respective
learned latent graphs. Average rewards in novel test environments (B, E). CSCG
performance reaches its peak by 9 episodes in MPG and stays the same
thereafter. EPN performance is shown on evaluation at different stages of
training. (C, F) Number of steps to get the reward in subsequent tasks in an
episode. The CSCG agent explores first without optimizing for reward
collection, so our first few tasks take longer to complete but we reach
optimal planning after that. Values for optimal exploration and planning are
also plotted as dotted lines for comparison. First few tasks in One-Shot
StreetLearn take variable numbers of steps as we collect rewards incidentally.
EPN baselines and optimal values for both environments are re-plotted using
the data from Ritter et al. (2020). Error bars are 95% CI of the SEM.

CSCG schema learns the graph structure in few episodes. The CSCG schema agent
first explores the grid randomly collecting observations for a few episodes.
After each episode, we learn a CSCG model that best explains the experiences
across all episodes so far observed. We reuse the same schema ($T$) across all
episodes and learn a new binding (emission matrix $E$) per episode. It takes
only 9 episodes (900 steps) to learn a perfect schema of this $4\times 4$ grid
environment. In subsequent episodes, we rapidly learn the schema bindings and
do planning to maximize the rewards. We employ two different hard coded
exploration policies: random navigation actions and an optimal set of actions
to cover the learned transition graph. Average reward per episode $\pm$
standard error of the mean (SEM) after learning the schema is: $17.3\pm 0.57$
for random and $26.4\pm 0.17$ for optimal exploration policy, which is
comparable to Episodic planning network (EPN) (Ritter et al., 2020). In
contrast, EPN takes more than 10 million episodes ($>10^{9}$ steps) of
training to reach its optimal performance (Fig. 2B). Planning in our model is
optimal in the number of steps to the reward, on par with EPN and the oracle
(Fig. 2C). CSCG performance remains the same since the first reward is
collected after exploration and the plans are optimal thereafter. Note that
the number of steps to finish the first task is longer in our case ($18\pm
0.09$ steps) than EPN, but the average reward in an episode is comparable.

####  4.1.2 One-Shot StreetLearn

One-Shot StreetLearn is a challenging variant of the StreetLearn task
(Mirowski et al., 2019) with varying connectivity structure across episodes to
evaluate rapid task solving ability (Fig. 2D). In each episode, the agent is
placed in a new neighborhood of a city and the task is to navigate to a goal,
specified by the goal street view image, and collect the reward. After
collecting a reward, the agent is re-spawned in a new location and a new goal
is specified. Unlike the MPG, the transition graph changes every episode. We
evaluate our model’s ability to rapidly learn in an episode and to navigate
optimally to the goals to maximize the rewards. Note that there is no schema
reuse in this setting: we learn a new model for every episode. This showcases
the ability to learn rapidly within a few steps without any prior training and
plan efficiently with the learned model.

CSCG matches optimal planning in One-Shot StreetLearn. For the CSCG agent, we
follow an explore and exploit strategy with a hard coded exploration policy.
During exploration, the agent navigates every action from every observation it
encounters while collecting the rewards as it encounters the goals, and uses
this experience to learn a CSCG. This is a guided exploration to cover every
possible edge in the transition graph. After exploration, the agent plans with
the CSCG and collects rewards (See Suppl. §B.2 for details). Average reward
$\pm$ SEM over 100 episodes is $21.7\pm 3.7$, which is lower than EPN (28.7)
as our exploration strategy is not optimal (Fig. 2E). Since we do not consider
optimal exploration in this work, we compare the planning performance on the
learned model after exploration. Post exploration, our agent takes on average
$4.8\pm 0.03$ steps to reach the goal, which matches the optimal value (Ritter
et al., 2020) (Fig. 2F). Note that we do not transfer any learning across
episodes in this setting since the graph changes every episode. In cities with
re-usable graph structures such as the grid layout in Manhattan, CSCG schemas
benefit from the reuse. We evaluate this schema reuse in detail in much harder
settings in the following experiments on navigating in rooms with extensively
aliased observations.

Figure 3: Rapid learning using schemas in novel 3D environments. A. As the
agent walks in a test room, we learn the emission matrix and measure the
negative log likelihood (NLL) of observation sequences under different
schemas. The schema with the least NLL (conversely the highest likelihood) is
considered the best match. The left panel shows the isometric view of the the
training rooms used for learning the schemas. The middle and right panels show
schema matching for two example test rooms. The test rooms differ in color and
lighting conditions from the training rooms. B. Schema stitching in a novel 3D
environment (left) composed of four out of the six rooms. Transition graphs
show the prior model, with potential connections between all six schemas, and
the learned models after 5,000 and 10,000 steps. (Right) NLL of test data for
models learned at different walk lengths, for both the model stitched from
known schemas and the model learned from scratch. Error bars are SEM, computed
across 5 repetitions.

###  4.2 Schema matching and transfer learning in highly aliased environments

In the next set of experiments, we evaluate schema matching and transfer
learning in novel environments that vary in observations and in some cases
size from training environments. For these experiments, we use a more
challenging setting than the benchmarks above, with larger 2D and 3D simulated
environments and extensive aliasing (Beattie et al., 2016). Note that even
though we use spatial navigation as our setup, we do not use any assumptions
about space (Euclidean or otherwise) and model it as a graph navigation
problem. For both 2D and 3D environments, we use a simple model of the agent’s
observations that does not rely on 3D perception, so our results are agnostic
to the specific 3D spatial setup.

We start with a set of environments of different shapes and topologies with
extensive aliasing. Analogous to a large empty arena, observations in the
interior of these environments are perceptually aliased Lajoie et al. (2019);
Whitehead and Ballard (1991), see Fig. 1D. Aliasing occurs by construction in
the 2D case, and as a result of clustering in the 3D case. We first learn the
schemas in training rooms (Fig. 1D), as described in section 3.3 . We evaluate
schema matching on test rooms with similar layouts but with novel observations
and size variations. In a novel test room, the agent takes a random walk and
evaluates the likelihood of the observation sequences given the actions
executed under different learned schemas (Fig. 1E). Note that this requires
first learning the new emission matrix from the data collected during this
random walk and computing the likelihood for each schema. The schema with the
best likelihood is considered the matching schema. We evaluate the likelihoods
of different schemas at multiple intervals during the random walks. See Suppl.
§B.3 for details and parameters used in the following experiments.

CSCG schemas rapidly learn matching bindings in novel environments. To
demonstrate this, we used 3D environments (Beattie et al., 2016) with 6
different layouts. The agent can navigate with 3 discrete egocentric actions
(move forward, turn left, turn right) and the observations are RGB images
corresponding to the agent’s view (Fig. 1A, B). The observation space in this
setup is large and complex and demonstrates the model’s applicability to such
use cases. We follow the procedure for continuous observations described in
$\S$3.2 and Suppl. §A.3 to learn models of the training environments (see Fig.
S1 for the learned graphs after training using random walks). We evaluated
rapid learning on test environments with the same layout but with different
colored walls, floor, and environment lighting, which corresponds to entirely
new RGB observations, for an agent navigating in those rooms using the
procedure described in Suppl. §A.4. Fig. 3A shows successful rapid matching of
correct schemas in test rooms as evaluated by negative log likelihoods (See
Suppl. Fig. S2 for full results). The correct schema was identified for all
six test environments, usually within 1,000 steps in the environment and in
all cases within 2,000 steps, compared to about 50,000 steps used to learn
without schema reuse.

To test the transfer to size variations, we used 2D rooms of five different
layouts (cylinder, rectangle, square with hole, torus, and U-shape) each in 3
different sizes (small, medium, large) (Fig. S3). We learn schemas for the
medium-sized versions of these rooms using a random walk of $50,000$ steps. In
the test rooms with novel observation mappings, the agent takes a random walk
while we learn the new emissions and evaluate the likelihoods of these
observations under all schemas every 5 steps. See Appendix B.3.2 for details.
Fig. S3 shows the negative log likelihoods of all test room under all schemas.
By reusing clone structure, we are able to correctly match the schema in all
cases by 95 steps demonstrating a rapid matching and adaptation to novel
environments with size and observation variations (Appendix Fig. S4 shows
results without using clone structure). We reproduced these results using ten
digit exemplars of binarized MNIST dataset as our room layouts as they provide
an interesting variety of shapes and topologies that are not designed by us
(see Suppl. §B.3.3 for results).

Schema matching also works in environments composed of multiple schemas. We
demonstrate this in novel test environments that are composed of pairs of
MNIST digit rooms, and show online schema matching of individual digits. See
Suppl. §B.5 for details and results.

###  4.3 Rapid learning of novel environments with compositionality of graph
schemas

Schemas can also be used to efficiently learn and navigate in larger
environments composed of known schemas. We do this by matching schemas and
learning the transition structure between them. Fig. 3B shows an example 3D
simulated environment, composed of four smaller 3D rooms. The agent walks in
this environment and learns the composed model (both joint transitions and
emissions) in far fewer steps than needed to learn without using schemas. Fig.
3B also shows the prior model with all potential connections and the learned
model at two different walk lengths. Model quality is measured by the negative
log likelihood over $10,000$-step test walks, for models trained after walks
of different lengths with and without using schemas. Using schemas, we were
able to learn a perfect model of the environment in less than $10,000$ steps,
whereas learning from scratch was significantly worse even after $30,000$
steps. See Suppl. §B.5 for the learning algorithm and experimental details.
This ability to learn by composing and reusing previously learned schemas
enables rapid adaptation to novel environments which only gets better with
more experience.

Figure 4: Schemas enable rapid and robust planning in new environments. A.
Given partial experience in a novel environment (here, a differently colored
variation of the four-room environment from Fig. 3), the joint transition
model and the agent’s location within it can be identified and used to rapidly
navigate around obstacles and find the shortest path to the goal (here, to
return to the starting position). B. (Left) An example test room and its
matching schema. Note the mismatch in the size and the off-center hole in the
test room. (Right). Navigation performance as measured by the distance from
the true goal and the number of plans/re-plans the agent makes, as a function
of size and structure variation between test room and the matching schema.
Error bars are 95% CI of the SEM.

###  4.4 Rapid planning and navigation in novel environments

Rapid schema matching and binding enables planning in novel environments with
limited experience. We first demonstrate this capacity in a novel variation of
the four-room 3D environment that was introduced in Fig. 4A. In this
demonstration, an agent first walks in the test environment, and is then
tasked with finding the shortest route back to its starting location. We used
a manually specified initial walk in order to cover a good portion of the
environment in a minimal number of steps. In this case, the schema – the
composite schema that was learned previously for the four-room combination
(Fig 3) – is known a priori. The agent uses the observations from the initial
walk to learn new emission bindings and decode its current and goal positions
within the model. The agent then plans the shortest path to the goal using
this model (Fig. 4B). Note that the planned route passes through a yet
unvisited part of the environment.

While executing a plan, if the agent estimates that it has not reached the
goal after accounting for the new observations, it is possible that there is a
schema mismatch or the estimated emission matrix is inaccurate. When this
occurs, the agent can initiate re-planning after updating the model using new
experiences gained while executing the plan. This process is iterated until
the agent believes that it has reached its goal after decoding observations
from the initial random walk and all subsequent re-planning steps. In an
experiment designed to test this capacity, we systematically evaluate this
robustness to schema mismatch in 2D environments in terms of Manhattan
distance from the goal location, and the number of re-plans required. Fig. 4B
shows results for one example schema. We can successfully navigate even with
size and structure variations and the performance degrades gracefully as the
variation between the schema and the test room increases. The number of
planning attempts required to reach the goal also increases smoothly. The
diagonal smoothing term that adds self-transition probabilities (see Suppl.
§B.6) is critical for generalization to size variations. Without this
smoothing, the agent never reaches the goal in larger size variations of the
test rooms as shown in Fig 4B (See Suppl. §B.6 for results in another room
layout).

##  5 Discussion and Future Work

Learning abstractions that can rapidly bind to observations from environments
that share the same underlying structure is the hypothesized mechanism for
transfer learning in humans and animals (Kumaran et al., 2016; Tse et al.,
2007; Zhou et al., 2021). We have proposed a concrete computational model of
abstractions and rapid binding using graph schemas that learn higher-order
structures from aliased observation sequences, and uses a slot-binding
mechanism for transferring those schemas to rapidly learn models of new
environments. CSCG schemas learned graph structure in far fewer episodes than
a deep RL agent and matched optimal planning in MPG and One-Shot StreetLearn
tasks. In highly aliased environments, CSCG schemas found matching schema
bindings in novel rooms of different sizes. In composed rooms, we were able to
match the correct schemas as the agent moved across rooms corresponding to
different schemas. We showed successful planning to goals in rooms with shape
and size variations from the matching schema by re-planning and updating the
model while walking to the goal. More importantly, known schemas can be
composed to rapidly learn novel environments and new larger schemas. This
ability bootstraps on itself and only gets better with more experience. There
are many clear directions for potential future works to build on top of our
current work, and we list some of them below.

Schema learning from experiences. We learned schemas independently and
explicitly in this work, but in the real world, it might not be feasible to
have access to differentiated experiences belonging to distinct schemas.
Learning reusable schemas from a continuous stream of experiences (Farzanfar
et al., 2023) could be an interesting future work.

Schemas vs memories. We discard previously learned emissions from past
experiences and learn a new binding. However, in some cases, previously
learned emissions are directly applicable and therefore keeping those in
addition could enable even faster zero-shot adaptation when there is a match.
This can be thought as keeping specific memories versus using abstract
schemas.

Schema maintenance. Our schemas in this work are fixed. However, it is
possible to update the schemas with new experiences. In fact, children
initially tend to perceive and remember experiences that fit in their existing
schemas and develop the flexibility later (Piaget and Cook, 1952). Similarly,
we could update schemas based on new experiences, and even make the schemas
themselves flexible to encapsulate related abstractions but still constrained
by rules to allow consistent inference.

Active exploration. We used either random or known optimal exploration
policies to learn and bind schemas. But the schemas provide action-conditioned
beliefs on future observations, and by choosing actions that could optimally
disambiguate different schemas and seek to learn connections between them, we
could potentially do much better than random exploration. Similarly, instead
of random exploration to learn a new environment and schema, we could direct
the exploration policy by composing known schemas (Sharma et al., 2021), and
even actively learn them while exploring.

## References

  * Baraduc et al. (2019) P. Baraduc, J.-R. Duhamel, and S. Wirth.  Schema cells in the macaque hippocampus.  _Science_ , 363(6427):635–639, Feb. 2019.
  * Beattie et al. (2016) C. Beattie, J. Z. Leibo, D. Teplyashin, T. Ward, M. Wainwright, H. Küttler, A. Lefrancq, S. Green, V. Valdés, A. Sadik, J. Schrittwieser, K. Anderson, S. York, M. Cant, A. Cain, A. Bolton, S. Gaffney, H. King, D. Hassabis, S. Legg, and S. Petersen.  Deepmind lab, 2016.  URL https://arxiv.org/abs/1612.03801.
  * Cormack and Horspool (1987) G. V. Cormack and R. N. S. Horspool.  Data compression using dynamic Markov modelling.  _The Computer Journal_ , 30(6):541–550, 1987\.
  * Crouse et al. (2021) M. Crouse, C. Nakos, I. Abdelaziz, and K. Forbus.  Neural analogical matching.  _AAAI_ , 35(1):809–817, May 2021.
  * Dedieu et al. (2019) A. Dedieu, N. Gothoskar, S. Swingle, W. Lehrach, M. Lázaro-Gredilla, and D. George.  Learning higher-order sequential structure with cloned hmms, 2019.  URL https://arxiv.org/abs/1905.00507.
  * Eichenbaum and Cohen (2014) H. Eichenbaum and N. J. Cohen.  Can we reconcile the declarative memory and spatial navigation views on hippocampal function?  _Neuron_ , 83(4):764–770, Aug. 2014.
  * Eysenbach et al. (2019) B. Eysenbach, R. R. Salakhutdinov, and S. Levine.  Search on the replay buffer: Bridging planning and reinforcement learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Farzanfar et al. (2023) D. Farzanfar, H. J. Spiers, M. Moscovitch, and R. S. Rosenbaum.  From cognitive maps to spatial schemas.  _Nature Reviews Neuroscience_ , 24(2):63–79, 2023\.
  * George et al. (2021) D. George, R. V. Rikhye, N. Gothoskar, J. S. Guntupalli, A. Dedieu, and M. Lázaro-Gredilla.  Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps.  _Nat. Commun._ , 12(1):2392, Apr. 2021.
  * Gilboa and Marlatte (2017) A. Gilboa and H. Marlatte.  Neurobiology of schemas and Schema-Mediated memory.  _Trends Cogn. Sci._ , 21(8):618–631, Aug. 2017\.
  * Gregor et al. (2019) K. Gregor, D. Jimenez Rezende, F. Besse, Y. Wu, H. Merzic, and A. van den Oord.  Shaping belief states with generative environment models for rl.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Gupta et al. (2017) S. Gupta, J. Davidson, S. Levine, R. Sukthankar, and J. Malik.  Cognitive mapping and planning for visual navigation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 2616–2625, 2017.
  * Igl et al. (2018) M. Igl, L. Zintgraf, T. A. Le, F. Wood, and S. Whiteson.  Deep variational reinforcement learning for pomdps.  In _International Conference on Machine Learning_ , pages 2117–2126. PMLR, 2018.
  * Kerg et al. (2022) G. Kerg, S. Mittal, D. Rolnick, Y. Bengio, B. A. Richards, and G. Lajoie.  Inductive biases for relational tasks.  In _ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality_ , 2022.
  * Kipf et al. (2019) T. Kipf, Y. Li, H. Dai, V. Zambaldi, A. Sanchez-Gonzalez, E. Grefenstette, P. Kohli, and P. Battaglia.  Compile: Compositional imitation learning and execution.  In _International Conference on Machine Learning_ , pages 3418–3428. PMLR, 2019.
  * Kumaran et al. (2016) D. Kumaran, D. Hassabis, and J. L. McClelland.  What learning systems do intelligent agents need? complementary learning systems theory updated.  _Trends Cogn. Sci._ , 20(7):512–534, July 2016\.
  * Lajoie et al. (2019) P.-Y. Lajoie, S. Hu, G. Beltrame, and L. Carlone.  Modeling perceptual aliasing in slam via discrete–continuous graphical models.  _IEEE Robotics and Automation Letters_ , 4(2):1232–1239, 2019.  10.1109/LRA.2019.2894852.
  * Lampinen et al. (2021) A. Lampinen, S. Chan, A. Banino, and F. Hill.  Towards mental time travel: a hierarchical memory for reinforcement learning agents.  _Advances in Neural Information Processing Systems_ , 34:28182–28195, 2021.
  * Mirowski et al. (2019) P. Mirowski, A. Banki-Horvath, K. Anderson, D. Teplyashin, K. M. Hermann, M. Malinowski, M. K. Grimes, K. Simonyan, K. Kavukcuoglu, A. Zisserman, et al.  The streetlearn environment and dataset.  _arXiv preprint arXiv:1903.01292_ , 2019.
  * Mitchell (2021) M. Mitchell.  Abstraction and analogy-making in artificial intelligence.  _Ann. N. Y. Acad. Sci._ , 1505(1):79–101, Dec. 2021.
  * Ni et al. (2021) T. Ni, B. Eysenbach, and R. Salakhutdinov.  Recurrent model-free rl can be a strong baseline for many pomdps.  _arXiv preprint arXiv:2110.05038_ , 2021.
  * Packer et al. (2018) C. Packer, K. Gao, J. Kos, P. Krähenbühl, V. Koltun, and D. Song.  Assessing generalization in deep reinforcement learning.  _arXiv preprint arXiv:1810.12282_ , 2018.
  * Piaget and Cook (1952) J. Piaget and M. T. Cook.  _The origins of intelligence in children._ WW Norton & Co, 1952.
  * Preston and Eichenbaum (2013) A. R. Preston and H. Eichenbaum.  Interplay of hippocampus and prefrontal cortex in memory.  _Curr. Biol._ , 23(17):R764–73, Sept. 2013.
  * Raju et al. (2022) R. V. Raju, J. S. Guntupalli, G. Zhou, M. Lázaro-Gredilla, and D. George.  Space is a latent sequence: Structured sequence learning as a unified theory of representation in the hippocampus, 2022.
  * Rakelly et al. (2019) K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen.  Efficient off-policy meta-reinforcement learning via probabilistic context variables.  In _International conference on machine learning_ , pages 5331–5340. PMLR, 2019.
  * Ritter et al. (2020) S. Ritter, R. Faulkner, L. Sartran, A. Santoro, M. Botvinick, and D. Raposo.  Rapid task-solving in novel environments, 2020.  URL https://arxiv.org/abs/2006.03662.
  * Samborska et al. (2022) V. Samborska, J. L. Butler, M. E. Walton, T. E. J. Behrens, and T. Akam.  Complementary task representations in hippocampus and prefrontal cortex for generalizing the structure of problems.  _Nat. Neurosci._ , pages 1–13, Sept. 2022.
  * Savinov et al. (2018) N. Savinov, A. Dosovitskiy, and V. Koltun.  Semi-parametric topological memory for navigation.  _arXiv preprint arXiv:1803.00653_ , 2018.
  * Shanahan and Mitchell (2022) M. Shanahan and M. Mitchell.  Abstraction for deep reinforcement learning, 2022.  URL https://arxiv.org/abs/2202.05839.
  * Sharan et al. (2017) V. Sharan, S. M. Kakade, P. S. Liang, and G. Valiant.  Learning overcomplete hmms.  In _Advances in Neural Information Processing Systems_ , pages 940–949, 2017.
  * Sharma et al. (2021) S. Sharma, A. Curtis, M. Kryven, J. Tenenbaum, and I. Fiete.  Map induction: Compositional spatial submap learning for efficient exploration in novel environments, 2021.  URL https://arxiv.org/abs/2110.12301.
  * Stachenfeld et al. (2017) K. L. Stachenfeld, M. M. Botvinick, and S. J. Gershman.  The hippocampus as a predictive map.  _Nat. Neurosci._ , 20, 2017.
  * Tenenbaum et al. (2011) J. B. Tenenbaum, C. Kemp, T. L. Griffiths, and N. D. Goodman.  How to grow a mind: statistics, structure, and abstraction.  _Science_ , 331(6022):1279–1285, Mar. 2011.
  * Tse et al. (2007) D. Tse, R. F. Langston, M. Kakeyama, I. Bethus, P. A. Spooner, E. R. Wood, M. P. Witter, and R. G. M. Morris.  Schemas and memory consolidation.  _Science_ , 316(5821):76–82, Apr. 2007.
  * Verbeek et al. (2003) J. J. Verbeek, N. Vlassis, and B. Kröse.  Efficient greedy learning of gaussian mixture models.  _Neural computation_ , 15(2):469–485, 2003.
  * Wang et al. (2016) J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick.  Learning to reinforcement learn, 2016.  URL https://arxiv.org/abs/1611.05763.
  * Webb et al. (2021) T. W. Webb, I. Sinha, and J. Cohen.  Emergent symbols through binding in external memory.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=LSFCEb3GYU7.
  * Whitehead and Ballard (1991) S. D. Whitehead and D. H. Ballard.  Learning to perceive and act by trial and error.  _Machine Learning_ , 7:45–83, 1991.
  * Whittington et al. (2020) J. C. R. Whittington, T. H. Muller, S. Mark, G. Chen, C. Barry, N. Burgess, and T. E. J. Behrens.  The Tolman-Eichenbaum machine: Unifying space and relational memory through generalization in the hippocampal formation.  _Cell_ , 183(5):1249–1263.e23, Nov. 2020.
  * Whittington et al. (2021) J. C. R. Whittington, J. Warren, and T. E. J. Behrens.  Relating transformers to models and neural representations of the hippocampal formation, 2021.  URL https://arxiv.org/abs/2112.04035.
  * Wu et al. (1983) C. J. Wu et al.  On the convergence properties of the em algorithm.  _The Annals of statistics_ , 11(1):95–103, 1983\.
  * Zhang et al. (2021) L. Zhang, G. Yang, and B. C. Stadie.  World model as a graph: Learning latent landmarks for planning.  In _International Conference on Machine Learning_ , pages 12611–12620. PMLR, 2021.
  * Zhou et al. (2021) J. Zhou, C. Jia, M. Montesinos-Cartagena, M. P. H. Gardner, W. Zong, and G. Schoenbaum.  Evolving schema representations in orbitofrontal ensembles during learning.  _Nature_ , 590(7847):606–611, Feb. 2021.

## Supplementary material

##  Appendix A Expectation-Maximization learning of CSCGs

Cloned Hidden Markov Models (HMMs), first introduced in (Dedieu et al., 2019),
are a sparse restriction of overcomplete HMMs (Sharan et al., 2017) that can
overcome many of the training shortcomings of dynamic Markov coding (Cormack
and Horspool, 1987). Similar to HMMs, cloned HMMs assume the observed sequence
$\boldsymbol{x}\equiv\\{x_{1}\ldots,x_{N}\\}$ is generated from a hidden
process $\boldsymbol{z}\equiv\\{z_{1}\ldots,z_{N}\\}$ that obeys the Markovian
property,

$P(\boldsymbol{x})=\sum_{\boldsymbol{z}}P(z_{1})\prod_{n=2}^{N}P(z_{n}|z_{n-1})\prod_{n=1}^{N}P(x_{n}|z_{n}).$
(2)

Here $P(z_{1})$ is the initial hidden state distribution, $P(z_{n}|z_{n-1})$
is the transition probability from $z_{n-1}$ to $z_{n}$, and $P(x_{n}|z_{n})$
is the probability of emitting $x_{n}$ from the hidden state $z_{n}$.

In contrast to HMMs, cloned HMMs assume that each hidden state maps
deterministically to a single observation. Further, cloned HMMs allow multiple
hidden states to emit the same observation. All the hidden states that emit
the same observation are called the clones of that observation. With this
constraint, the joint distribution of the observed data can be expressed as,

$P(\boldsymbol{x})=\sum_{z_{1}\in C(x_{1})}\ldots\sum_{z_{N}\in
C(x_{N})}P(z_{1})\prod_{n=2}^{N}P(z_{n}|z_{n-1})$ (3)

where $C(x_{n})$ corresponds to the hidden states (clones) of the emission
$x_{n}$.

Clone-structured cognitive graphs (CSCGs) build on top of cloned HMMs by
augmenting the model with the actions of an agent
$\boldsymbol{a}\equiv\\{a_{1}\ldots,a_{N}\\}$ (George et al., 2021),

$P(\boldsymbol{x}|\boldsymbol{a})=\sum_{z_{1}\in C(x_{1})}\ldots\sum_{z_{N}\in
C(x_{N})}P(z_{1})\prod_{n=2}^{N}P(z_{n}|z_{n-1},a_{n-1})$ (4)

###  A.1 Learning the transition matrix with emissions fixed

Learning a CSCG in a new environment requires optimizing the vector of prior
probabilities $\pi$: $\pi_{k}=P(z_{1}=k)$ and the action-augmented transition
tensor $T$: $T_{ijk}=P(z_{n}=k|z_{n-1}=j,a_{n-1}=i)$. The emission matrix $E$,
which encodes the particular allocation of clones to observations, is kept
fixed throughout learning.

The standard algorithm to train HMMs is the expectation-maximization (EM)
algorithm (Wu et al., 1983), which in this context is known as the Baum-Welch
algorithm. Learning a CSCG using the Baum-Welch algorithm requires a few
simple modifications: the sparsity of the emission matrix can be exploited to
only use small blocks of the transition matrix both in the Expectation (E) and
Maximization (M) steps.

To this end, we assume the hidden states are indexed such that all the clones
of the first emission appear first, all the clones of the second emission
appear next, etc. Let $N_{\textrm{obs}}$ and $N_{\textrm{actions}}$ be the
total number of emissions and actions, respectively. The transition matrix $T$
can then be broken down into smaller submatrices $T(u,v,w)$, where
$v,w\in\\{1,\ldots,N_{\textrm{obs}}\\}$ and
$u\in\\{1,\ldots,N_{\textrm{actions}}\\}$. The submatrix $T(u,v,w)$ contains
the transition probabilities $P(z_{n}|z_{n-1},a_{n-1}=u)$ for $z_{n-1}\in
C(v)$ and $z_{n}\in C(w)$, where $C(v)$ and $C(w)$ correspond to the clones of
emissions $v$ and $w$ respectively.

The standard Baum-Welch equations can then be expressed in a simpler form in
the case of a CSCG. The E-step recursively computes the forward and backward
probabilities and then updates the posterior probabilities. The M-step updates
the transition matrix via row normalization.

E-Step:

$\displaystyle\alpha(1)$ $\displaystyle=\pi(z_{1})$
$\displaystyle\alpha(n)^{\top}$
$\displaystyle=\alpha(n-1)^{\top}T(a_{n-1},C(x_{n-1}),C(x_{n}))$ (5)
$\displaystyle\beta(N)$ $\displaystyle=1(z_{N})$ $\displaystyle\beta(n)$
$\displaystyle=T(a_{n},C(x_{n}),C(x_{n+1}))\beta(n+1)$ (6)
$\displaystyle\xi_{uvw}(n)$ $\displaystyle=\frac{\alpha(n)\circ
T(a_{n},v,w)\circ\beta(n+1)^{\top}}{\alpha(n)^{\top}T(a_{n},v,w)\beta(n+1)}$
(7) $\displaystyle\gamma(n)$
$\displaystyle=\frac{\alpha(n)\circ\beta(n)}{\alpha(n)^{\top}\beta(n)}.$ (8)

M-Step:

$\displaystyle\pi(z_{1})$ $\displaystyle=\gamma(1)$ (9) $\displaystyle
T(u,v,w)$
$\displaystyle=\sum_{n=1}^{N}\xi_{uvw}(n)\oslash\sum_{u=1}^{N_{\textrm{actions}}}\sum_{w=1}^{N_{\textrm{obs}}}\sum_{n=1}^{N}\xi_{uvw}(n)$
(10)

where $\circ$ and $\oslash$ denote element-wise product and division,
respectively (with broadcasting where needed). All vectors in the E and
M-steps are $N_{\textrm{cpe}}\times 1$ column vectors, where
$N_{\textrm{cpe}}$ is the number of clones per emission. We use a constant
number of clones per emission for simplicity of description, but the number of
clones can be selected independently per emission.

Importantly, CSCGs exploit the sparsity pattern in the emission matrix when
performing training updates and inference, and achieve significant
computational savings when compared with HMMs (George et al., 2021).

###  A.2 Learning the emission matrix with transitions fixed

With a CSCG, transfer learning between different environments can be
accomplished by keeping its transition probabilities $T$ fixed and learning
the emissions associated to its nodes $E$ in the new environment. Further, if
we know that the new environment preserves the emission structure, then we can
further restrict the learning of $E$, with all the rows of $E$ that correspond
to the same observation in the original environment sharing the same
parameters.

The EM algorithm can be used to learn the emission matrix as follows. The
E-step recursively computes the forward and backward probabilities and then
updates the posterior probabilities. The M-step updates only the emission
matrix.

E-Step:

$\displaystyle\tilde{\alpha}(n)$
$\displaystyle=\left(T(a_{n-1})^{\top}\alpha(n-1)\right)\circ E(x_{n})$
$\displaystyle p_{\alpha}(n)$
$\displaystyle=\sum_{k=1}^{|Z|}\tilde{\alpha}_{k}(n)$
$\displaystyle\alpha(n)=\tilde{\alpha}(n)/p_{\alpha}(n)$ (11)
$\displaystyle\tilde{\beta}(n)$ $\displaystyle=T(a_{n})\left(\beta(n+1)\circ
E(x_{n+1})\right)$ $\displaystyle p_{\beta}(n)$
$\displaystyle=\sum_{k=1}^{|Z|}\tilde{\beta}_{k}(n)$
$\displaystyle\beta(n)=\tilde{\beta}(n)/p_{\beta}(n)$ (12)

$\gamma(n)=\frac{\alpha(n)\circ\beta(n)}{\alpha(n)^{\top}\beta(n)}$ (14)

M-Step:

$E(j)=\sum_{n=1}^{N}1_{x_{n}=j}\gamma(n)\oslash\sum_{n=1}^{N}\gamma(n)$ (15)

Note that the updates here involve $|Z|\times 1$ vectors, where $|Z|$
($=\sum{N_{cpe}}$) is the total number of hidden states in the model.
$T(a_{n})$ corresponds to the transition matrix for the action $a_{n}$ and is
of size $|Z|\times|Z|$, $E(x_{n})$ is a column of the emission matrix
corresponding to the emission $x_{n}$, and $1_{x_{n}=j}$ is an indicator
function. The forward message is initialized as $\tilde{\alpha}(1)=\pi\circ
E(x_{1})$; and the backward message is initialized as a vector of all 1s. The
emission matrix can be initialized randomly or with equal probabilities for
all observations.

When the clone structure is to be preserved, the $\gamma(n)$ term in the
E-step is modified as follows. For each observation
$j\in\\{1,\ldots,N_{\textrm{obs}}\\}$, we set the posterior probability for
all clones of $j$ to be the same:

$\bar{\gamma}_{k}(n)=\sum_{l\in C(j)}\gamma_{l}(n)\;\;\forall k\in C(j).$ (16)

For a given $T$ and $E$, the normalization term $p_{\alpha}(n)$ for the
forward messages in the E-step can be used to compute the negative log-
likelihood (NLL) of a sequence of observation-action pairs as follows:

${\rm NLL}=-\frac{1}{N}\sum_{n=1}^{N}\log p_{\alpha}(n)$ (17)

###  A.3 Learning the continuous CSCG

When each observation at timestep $n$ is a continuous vector $x_{n}$ (e.g., an
image), we want to maximize the following log-likelihood

$\log
p(\boldsymbol{x}|\boldsymbol{a};T,\boldsymbol{\mu})=\log\sum_{z,y}P(z_{1})\prod_{n=2}^{N}P(z_{n}|z_{n-1},a_{n-1})\prod_{n=1}^{N}p(y_{n}|z_{n})\mathcal{N}(x_{n}|\mu_{y_{n}},\sigma^{2}I)$
(18)

w.r.t. $T$, where $T_{ijk}=P(z_{n}=k|z_{n-1}=j,a_{n-1}=i)$, and
$p(y_{n}|z_{n})$ is 1 for $z_{n}\in C(y_{n})$ and 0 for all others.

This graphical model is tree-structured, so its value can be computed exactly.
Parameter optimization can be done via EM, where both the expectation and
maximization steps are exact. Rather than learning both $T$ and
$\boldsymbol{\mu}$ simultaneously, we find it simpler to proceed in two steps.

In the first step we fix $P(z_{1})=T_{ijk}=1/|Z|$, where $|Z|$ is the number
of states in the hidden space, and learn $\boldsymbol{\mu}$ only. The problem
thus simplifies to maximizing

$\displaystyle\log p(\boldsymbol{x}|\boldsymbol{a};T=T_{\text{uniform}},\mu)$
$\displaystyle=-N\log|Z|+\log\sum_{\boldsymbol{z},\boldsymbol{y}}\prod_{n=1}^{N}p(y_{n}|z_{n})\mathcal{N}(x_{n}|\mu_{y_{n}},\sigma^{2}I)$
$\displaystyle=-N\log|Z|+\sum_{n=1}^{N}\log\sum_{y_{n}=1}^{K}n_{C}(y_{n})\mathcal{N}(x_{n}|\mu_{y_{n}},\sigma^{2}I)$
$\displaystyle=\sum_{n=1}^{N}\log\sum_{k=1}^{K}\frac{n_{C}(k)}{|Z|}\mathcal{N}(x_{n}|\hat{\mu}_{k},\sigma^{2}I)$
(19)

w.r.t. $\boldsymbol{\mu}$. The number of clones of a given centroid
$\hat{\mu}_{k}$ is denoted by $n_{C}(k)$ (thus, $\sum_{k=1}^{K}n_{C}(k)=|Z|$).
In the last equality we collect all the Gaussians that are known to have the
same mean (i.e., all the clones with the same centroid). The computation is
more efficient, since it no longer scales with the total number of hidden
states $|Z|$, but only with the number of distinct means $K$.

The astute reader will recognize Eq. 19 as the log-likelihood of an isotropic
Gaussian mixture model, which can be optimized greedily using $K$-means. The
centroids $\hat{\mu}_{k}$ are the cluster centers and $\frac{n_{C}(k)}{|Z|}$
corresponds to the prior probabilities of each center. In other words, to
maximize $\log
p(\boldsymbol{x}|\boldsymbol{a};T=T_{\text{uniform}},\boldsymbol{\mu})$ w.r.t.
$\boldsymbol{\mu}$ we simply need to run $K$-means on the input data
$\boldsymbol{x}$ (temporal ordering becomes irrelevant) and then assign a
number of clones to each centroid that is proportional to the prior
probability of that cluster. When using $K$-means, the value of $\sigma^{2}$
does not affect the optimization of the centroids. Once the centroids have
been chosen, one can set $\sigma^{2}$ based on the hard assignments from
$K$-means, where the maximum likelihood estimate is closed form and simply
matches the average distortion. Directly optimizing Eq. 19 w.r.t.
$\boldsymbol{\mu}$ and $\sigma^{2}$ is also possible using EM (and in
principle, more precise), but will be more computationally expensive.

In the second step, we keep $\boldsymbol{\mu}$ fixed as per the previous step,
and learn $T$ by maximizing Eq. 18 w.r.t. $T$ using EM. While it would be
possible to learn $\boldsymbol{\mu}$ and $T$ simultaneously in the same EM
loop, we find that proceeding in these two steps results in faster convergence
without significantly deteriorating the quality of the found local optimum.

Two approaches are possible for the learning of $T$:

##### Soft-evidence

This would be the vanilla approach of maximizing Eq. 18 w.r.t. $T$ using EM:
in the E-step, we obtain the exact posterior over each $z_{n}$ by running sum-
product inference, and in the M-step we find the maximizing $T$ in closed form
for that posterior. This is computationally demanding, since the cost per
learning iteration scales with ${\cal O}(N|Z|^{2})$.

##### Hard-evidence

We can apply here the same idea that converts EM for Gaussian mixture
modelling into $K$-means (sometimes known as hard-EM). Instead of considering
all the centroids as partially _responsible_ Verbeek et al. (2003) for each
sample $x_{n}$, we can assign all the responsibility to the dominating
Gaussian. This is often a very precise approximation, since the dominating
Gaussian typically takes most of the responsibility. In turn, this means that
at each timestep, only the clones of the centroid that is closest (in the
Euclidean sense) can be active. With this change, the cost per learning
iteration now scales as ${\cal O}(Nn_{C}^{2})\ll{\cal O}(N|Z|^{2})$, where we
have assumed a constant number of clones $n_{C}=n_{C}(k)$ for all centroids
for ease of comparison. When doing this, the observations become effectively
discrete and learning $T$ reduces to the procedure described in Section A.1.

In our experiments we use the second approach, which is much more efficient.
Putting all together, the final learning procedure can then be summarized as
follows:

  1. 1.

Run $K$-means on the training data for a given number of centroids $K$.

  2. 2.

Assign to each centroid a number of clones proportional to the prior
probablity of that centroid, as obtained by $K$-means.

  3. 3.

Vector-quantize the data (train and test) by assigning to each data point the
label of the closest centroid.

  4. 4.

Learn the action-conditional transition matrix as if this was a discrete CSCG,
as explained in §A.1.

###  A.4 Transfer learning in continuous CSCGs

In the transfer learning setting we know $T$ for a given environment, and need
to learn $\boldsymbol{\mu}$. The principled approach is simple: maximize Eq.
18 w.r.t. $\boldsymbol{\mu}$ while keeping $T$ fixed. This can be done using
EM, where both expectation and maximization will be exact steps. Clone
structure can be enforced exactly as well, by tying the parameters of the
centroids that correspond to the same clone during the M-step.

The above approach uses soft-evidence and runs into the same computational
issues as we discussed in §A.3. If we do not need to keep the same clone
structure, we can resort to the same hard-evidence trick to significantly
bring down computation time, again from ${\cal O}(N|Z|^{2})$ to ${\cal
O}(Nn_{C}^{2})$. The process again reduces to using $K$-means and then doing
transfer learning in the discrete space:

  1. 1.

Run $K$-means on data from the new environment. $K$ can be the number of
centroids from the environment we are transferring from, but does not need to.
The only restriction is that it is not larger than the number of states in
$T$.

  2. 2.

Vector-quantize the data from the new environment by assigning to each data
point the label of the closest centroid.

  3. 3.

Learn the emission matrix as if this was a discrete CSCG, as explained in §A.2
(without preserving the clone structure).

####  A.4.1 Schemas for transfer learning

Grounded schema

A grounded schema is a 3-tuple $S_{g}\equiv\\{T,{\cal C},E\\}$ containing an
action-conditional transition tensor $T$, a clone structure ${\cal C}$, and an
emission matrix $E$ that respects the clone structure.

The transition tensor $T$ encodes the transition probabilities, with elements

$T_{ijk}=P(z_{n}=k|z_{n-1}=j,a_{n-1}=i)$

We refer to $T$ as the schema portion of the grounded schema.

Latent variables $z$ take values in the range $1,\ldots,|Z|$. This range is
partitioned in $K\leq|Z|$ groups of cloned states. The group index of a given
latent state is given by $1\leq{\cal C}(\cdot)\leq K$. Thus, two latent states
$z=j$ and $z=j^{\prime}$ belong to the same clone group iff ${\cal C}(j)={\cal
C}(j^{\prime})$.

The emission matrix $E$ encodes the emission probabilities, with elements

$E_{ij}=P(y_{n}=j|z_{n}=i)$

An emission matrix is said to respect the clone structure if the condition
${\cal C}(j)={\cal C}(j^{\prime})\implies
E_{ij}=E_{ij^{\prime}}\leavevmode\nobreak\ \forall\leavevmode\nobreak\
i,j,j^{\prime}$ is satisfied.

Ungrounded schema with clone structure An ungrounded schema with clone
structure is a tuple $S_{u}\equiv\\{T,{\cal C}\\}$. It can be derived from a
grounded schema $S_{g}$ by removing the emission matrix $E$. Conversely, an
ungrounded schema with clone structure can produce a grounded schema when it
is combined with an emission matrix $E$ that respects its clone structure.
This process is called grounding and corresponds to transfer learning, since
the same schema is reused.

Ungrounded schema An ungrounded schema, or simply a schema, is noted as $T$,
and is obtained from a grounded schema in which we have removed both the clone
structure ${\cal C}$, and the emission matrix $E$. As before, a schema can be
combined with an arbitrary clone structure ${\cal C}$ and an emission matrix
$E$ that respects it, to produce a grounded schema. This is an even more
flexible form of transfer learning.

##  Appendix B Experiment details

In George et al. (2021), it was observed that the convergence of EM for
learning the parameters of a CSCG can be improved by using a smoothing
parameter called the pseudocount. The pseudocount is a small constant that is
added to the accumulated counts statistics matrix, which ensures that any
transition under any action has a non-zero probability. This ensures that the
model does not have zero probability for any sequence of observations at test
time. For our experiments, unless otherwise specified, we use a psuedocount of
$2\times 10^{-3}$ for schema learning (learning $T$, with $E$ fixed), and
$10^{-7}$ for schema matching (learning $E$, with $T$ fixed).

###  B.1 Memory & Planning Game

The game environment is a $4\times 4$ grid of symbols in which the agent can
navigate in the four canonical directions by one grid step (up, down, left,
right), and collect reward of 1 at a goal location in the grid (Fig. 2A).
Reward is 0 otherwise. Once the agent collects the reward at the current goal,
the agent is placed in a new random position and a new goal symbol is sampled
to which the agent must navigate to collect the next reward. All grid
positions have unique symbols and the symbol-position mapping is randomized at
the start of each episode, which lasts for 100 actions. The agent’s
observation is a tuple of the symbol in its current position and the goal
symbol. See Ritter et al. (2020) for more details. We assume knowledge of the
collect action function and execute it only when the goal symbol is reached.

We employ 3 different hard coded exploration policies to cover the
observations in a new episode: (i) random navigation actions, (ii) random
navigation actions but limited to (up, right), and (iii) an optimal set of
actions to cover the learned transition graph – a Hamiltonian cycle that
repeats the sequence {up, up, up, right} four times. For a new episode, we
take actions based on one of these policies, learn the emission matrix $E$ at
every step, and plan to reach the current goal using the current estimate of
$E$. Planning a path from the current position to a goal position is achieved
via Viterbi decoding in the CSCG. If the probability of reaching the goal is
above a certain threshold, we find the series of actions that lead to the
goal.

In the beginning, the agent is uncertain about both its own position and the
goal position, so it executes actions based on one of the hard-coded policies
until we are confident of reaching the goal state from the current state above
a probability threshold. If the current goal symbol is not yet observed in
this episode, we execute actions based on the policy, otherwise we compute the
probability of our planned set of actions reaching the goal symbol. As we
navigate, the estimate of $E$ becomes better and the plans are more likely to
succeed. Our planning algorithm takes both these uncertainties into account.
We evaluated for a total of 100 episodes ($10^{4}$ steps) and the average
reward stays the same after the 9th episode once we learn the schema. Average
reward per episode are reported for 100 episodes from 10th episode onwards
i.e., after learning.

###  B.2 One-Shot StreetLearn

In each episode of this task, the agent is placed in a new neighborhood of a
city and needs to move to a goal location and direct itself to the target,
specified by the goal street view image, to collect the reward. The agent can
take one of four actions: move forward, turn left, turn right, collect reward.
After every reward collection, the agent is placed in a new location with a
new heading direction and provided a new goal street view image. Each episode
lasts 200 steps and a new neighborhood is sampled for each new episode. The
agent needs to collect as many rewards as possible during each episode. The
StreetLearn environment represents the agent’s perception, a street view
image, as a unique compressed string that we use as our agent’s observation.
For the CSCG agent, we follow an explore and exploit strategy. The agent first
navigates every action from every observation it encounters while collecting
the rewards as it encounters the goals. This is guided exploration to cover
every possible edge in the transition graph. A CSCG model is then learned from
the sequence of observations and actions experienced in this episode. For any
subsequent step, we find the closest path from the agent’s current observation
to the goal observation and execute the actions, re-planning at every step to
account for any mismatch between the model expectation and reality.

###  B.3 Schema matching

####  B.3.1 Schema matching in simulated 3D environments

In our simulated 3D environments, the agent receives egocentric observations
that are $64\times 64$ RGB images. The agent’s action space is discretized to
3 egocentric actions: go forward by $1$m, rotate $90^{\circ}$ left, and rotate
$90^{\circ}$ right. We convert the input RGB images into categorical
observations using a $K$-means vector quantizer.

Figure S1:  The top row shows the isometric views of example 3D simulated
environments used in our schema matching experiments. The bottom row shows the
transition graphs of CSCGs trained on the respective 3D environments. Note
that each location in the room is represented by four nodes in the transition
graph corresponding to the four possible heading directions of the agent.
Figure S2: Each column shows an example test room (top) and the negative log-
likelihood (NLL) of observations from that room under different schemas
(bottom). Error bars are SEM. Note that the correct schema was identified
(lowest NLL) in all 6 test rooms.

For the schema matching experiments with 3D environments, we used the six
rooms shown in Fig. S1. These were constructed with a variety of features,
including ramps and platforms that introduce one-way transitions (e.g.
dropping off of a ledge). For each of these environments, we trained a CSCG
using the following steps. First, we collected action-observation pairs from a
random walk of length $100,000$ steps. Then, we clustered the egocentric
observations from the walk using a $K$-means quantizer, with $K=50$. We used
EM to train a CSCG on these discretized observations and actions as described
in §A.1. The learned transition graphs for the six example rooms are also
shown in Fig. S1. Note that learning is not always perfect. Some of the
transition graphs have spurious edges. These imperfections, however, do not
affect our schema matching results.

Fig. S2 shows the schema recognition results in these 3D rooms. Given
observations from a test room and a hypothesized schema (ungrounded and
without clone structure), we learn the emission bindings as described in §A.4.
We compute the negative log likelihood of the observations according to Eq. 17
in §A.2. We observed that it takes typically between $500$ to $1000$ steps to
correctly identify the matching schema in each room, although this is higher
in a couple cases. In particular, disambiguation of schemas 5 and 6 required
closer to 2000 steps.

Figure S3: Rooms of different types and sizes used in the schema matching
experiments with size variations.

####  B.3.2 Schema matching in 2D rooms with size variation

We generated rooms of five different types, and three size variations per
type, as shown in Fig. S3. We selected the room and barrier dimensions such
that the rooms have similar number of accessible states. For the torus and
cylinder rooms, we use the same observation map as in the rectangular room
case. However, actions wrap around the top-bottom edges for the cylinder, and
the top-bottom and left-right pair of edges for the torus, as indicated by the
blue arrows in Fig. S3. We learned schemas on the medium size variation of
these room types. For each training room, we used action-observation sequences
from a random walk of length $50,000$ steps to train the respective CSCG.

For evaluating schema matching, we used test rooms of all three size
variations. Importantly, we generated the test rooms by permuting the
observations from the original room. As a consequence, at test time, our model
would have to relearn the observation mapping.

In Fig. S4, we show the schema matching performance with the use of clone-
structure. We observe that it takes at most $95$ steps to identify the correct
schema for all test rooms when using clone structure. We obtained these
results using the paired T-test over 25 random walks, per test room, and a
p-value threshold of $0.05$. We also evaluated the schema matching performance
without clone structure and found that it takes at most $255$ steps to find
the best matching schema when the clone structure is ignored.

Figure S4: Schema matching in novel environments with size and observation
variations. We measure the negative log likelihood (NLL) of observation
sequences in a given room under different schemas as the agent does random
walks. The schema with the least negative log likelihood (conversely the
highest likelihood) is considered the best match. Each panel corresponds to a
room of type and size indexed by the column and row headers, respectively.
Schemas are learned in medium sized rooms and tested on two other size
variations. Error bars are 95% CI of the SEM.  Figure S5: A. MNIST rooms
constructed using the binarized-MNIST dataset, with the observations mapped as
colors. B. Negative log likelihood (NLL) of observation sequences from a given
MNIST room under different schemas. Error bars are 95% CI of the SEM.

####  B.3.3 Schema matching in MNIST digit rooms

In this experiment, we used a set of ten digits from the binarized MNIST
dataset as our room layouts as they provide an interesting variety of shapes
and topologies. All the foreground pixels are treated as accessible locations;
background pixels are treated as obstacles. Observations at each pixel are
sampled in a similar way to the 2D rooms with aliased interiors (Fig. S5A). We
learn schemas for 10 selected MNIST-digits using action-observation pairs from
random walks of length $100,000$ steps in each room. We repeated our schema
matching procedure as before in this new set of rooms by keeping the room
structures constant but with new observations. Fig. S5 shows the negative log
likelihoods of observation sequences in test rooms under different schemas.
The correct schema was identified in at most $50$ steps in all test rooms.

###  B.4 Schema matching in composed environments

To demonstrate how schemas could be used in more complex environments composed
of multiple schemas, we constructed test rooms composed of pairs of
overlapping MNIST digits. As an agent randomly explores a composed room, we
compute the likelihood of a sliding window of action-observation pairs under
each schema. Note that we use a sliding window here, instead of the entire
history of observations, because we are interested in identifying schema the
agent is currently in. The probability of being in a particular schema can
then be computed using the softmax of the log-likelihoods. The details of
schema matching in composed MNIST rooms is presented in Algorithm 1. Heatmaps
in Fig. S6 show the estimated probability of being in one of the two schemas
at every location in a composed room, as the agent performs a random walk of
length $20,000$ steps and sliding window of length $200$ steps. We also
compute the accuracy of identifying the correct schema at all locations by
thresholding these probabilities. The average accuracy of identifying the
correct schema at all locations in all composed test rooms is $88.4\pm 0.9\%$.
See Table S1 for results on individual composed rooms.

Algorithm 1 Schema selection for composed rooms

Input List of schemas $S_{1},\ldots S_{H}$, observations $\boldsymbol{x}$ and
actions $\boldsymbol{a}$ of length $N$ in the test room, window length $w$.

1:$n\leftarrow w$

2:while $n<=N$ do

3: Estimate $E_{j}$ using $x_{n-w},\ldots,x_{n}$, $a_{n-w},\ldots,a_{n}$ for
each schema, $j\in\\{1,\ldots,H\\}$

4: Compute $L_{n,j}$: log likelihood of $x_{n-w},\ldots,x_{n}$ under schema
$j$

5: Compute $p_{n,j}=\mathrm{softmax}(L_{n})_{j}$ : the probability of being in
schema $j$ at time step $n$.

6: Select schema with highest $p_{n,j}$.

7:end while

Digits | 2, 3 | 5, 1 | 5, 3 | 2, 7 | 0, 9 | 9, 3 | 4, 6 | 8, 7
---|---|---|---|---|---|---|---|---
Acc. | ${87.6}\pm{4}$ | ${84.3}\pm{3}$ | ${91.1}\pm{3}$ | ${92.5}\pm{2}$ | ${87.9}\pm{3}$ | ${85.2}\pm{4}$ | ${85.9}\pm{3}$ | ${90.7}\pm{4}$
Table S1: Accuracy (Mean $\pm$ Standard deviation, in %) of the schema
predictions as an agent moves across various locations in composed MNIST
rooms. We use the ground truth locations to compute this accuracy, which is
averaged over five trials of random walks in each room. Note that for the
composed room experiments, the agent gets two schemas as input. Figure S6:
Each heatmap represents the probability of an agent being in one of two
schemas at each location in an MNIST composed room, averaged over all
instances of the agent being at that location during a random walk of length
$20,000$ steps.

###  B.5 Learning using schema compositionality

We generated novel environments by composing rooms with known schemas. For
each individual schema, we specify its frontiers as potential exits and
entries: (i) list of (state, action) tuples that specify how the agent can
exit the schema, and (ii) list of entry states from which an agent can enter
the schema. We pre-specify schemas and frontiers for our experiments, but
schemas and their frontiers can potentially be learned from training data
(e.g. as repeating motifs on a set of training environments using graph
clustering methods and metrics such as betweenness centrality). Once the
schemas and frontiers are specified, a novel composed environment is thus a
combination of different regions corresponding to known schemas that are
connected by their frontiers.

Algorithm 2 Learning a new CSCG model with schema compositionality

Input List of schemas $S_{1},\ldots,S_{H}$, frontiers
$FSA_{1},\ldots,FSA_{H}$, observations $\boldsymbol{x}$, actions
$\boldsymbol{a}$, and pseudocount $\alpha$
Output CSCG model of the room as $T$ and $E$.

1:$T\leftarrow{\rm block\ diagonal\ matrix}(S_{1},\ldots S_{H})$
$\triangleright$ Initialize $T$

2:for $i$ in $1,\ldots,H$ do $\triangleright$ Add frontier connections

3: for exit frontier tuple $(s,a)$ in $FSA_{i}$ do
$T[a,s,s_{k}]\leftarrow\frac{1}{|\\{s_{k}\in FSA\setminus
FSA_{i}\\}|}\,\,\forall{s_{k}\in\\{FSA\setminus FSA_{i}\\}}$

4: end for

5:end for

6:$E\leftarrow{\rm Uniform\ distribution}$ $\triangleright$ Initialize $E$

7:Use EM with fixed $T$ to learn $E$ from $(\boldsymbol{x},\boldsymbol{a})$
$\triangleright$ Learn $E$ with initial $T$ fixed

8:Use EM with $E$ fixed to learn $T$ with pseudocount $\alpha$
$\triangleright$ Learn $T$ with learned $E$ fixed

9:Use Viterbi learning with current $E$ fixed to refine $T$ $\triangleright$
Refine $T$

To learn a model of the composed environments, we use the observations and
actions collected during random walks. We first create a new CSCG model with a
joint transition tensor composed of known schemas as a block-diagonal tensor,
in which each block corresponds to a copy of one schema. We then connect exit
frontier states of a given schema copy to every other entry frontier (state)
in every other schema with a uniform probability of such transitions
conditioned by the corresponding frontier actions for those exit states. See
Fig. 3 (main text) for examples of graphs representing such joint transition
tensors. The procedure for learning the emission matrix and the joint
transition tensor using action-observation pairs in a composed environment is
described in Algorithm 2.

For the schema composition experiment, we used schemas for the six
environments that were used in the 3D schema matching experiment ($B.3.1, and
tested learning in a larger environment that was a composition of four of
these rooms. Because this complex environment, with one-way transitions in
some places, was difficult to cover effectively with a random walk, we opted
to use a more efficient walk generation procedure to obtain data for learning.
This procedure utilized ground-truth position knowledge to navigate, on every
step, towards the nearest unexplored (or least explored) transition. Using
this procedure, we were able to cover the entire composed environment in about
$10,000$ steps and generate data that was relatively balanced in its coverage
of the different parts of the environment for a wide range of walk lengths.

We first collected efficient walk data of $50,000$ steps for the multi-room
environment, and clustered image observations using $K$-means with $K=80$. We
constructed exact schemas for each of the six smaller rooms, and used these to
construct the initial joint transition tensor with six blocks, connecting
schema exit states to entrance states of every other schema (always for the
forward action only). The joint transition tensor was then used as the
starting point for the learning algorithm (Algorithm 2), which involved
learning both $E$ and $T$ for the multi-room environment.

We experimented with learning over different random walk lengths, from $1,000$
up to $30,000$. These walks were subsets of the full efficient walk data that
was collected. We computed mean NLL across 10 test random walks, each of
length $10,000$ steps, extracted from a separate random walk dataset that was
not used for training. As seen in Fig. 3B, with $5,000$ steps, we observed
several spurious edges in the learned composite transition graph. The NLL of
test data using this model was 0.73. With $10,000$ steps, we were able to
perfectly recover the joint transition graph. The NLL of test data using this
learned model was 0.17. In contrast, learning the CSCG for the multi-room
environment without schemas, with only $10,000$ steps (test NLL = 4.27) or
even $30,000$ steps (test NLL = 1.21) was not feasible. A comparison of the
NLL of test data using CSCGs trained with and without schemas as a function of
training data size is shown in Fig. 3B.

We also performed the schema composition experiment using 2D rooms from
§B.3.2. In Fig. S7, we visualize the models learnt with and without schemas
for different walk lengths.

Figure S7: Visual comparison of CSCG training with (schema composition) and
without (full CSCG learning) schemas, as a function of the length of the
training sequence, in an example composed 2D room. Figure S8:  Shortcut
finding performance in rectangle and square with a hole room layouts as
measured by the Manhattan distance from the goal and the number of plans/re-
plans to reach the goal, as a function of room size. Error bars are 95% CI of
the SEM.

###  B.6 Rapid inference and planning using schemas

Rapid schema matching and binding enables planning in novel environments with
limited experience. We evaluate planning to reach goals using schemas in test
rooms with different size and shape variations. We considered two test room
types – rectangle and square with a hole. The schemas correspond to the medium
size version of these rooms (same as in §B.3.2). We add an additional
smoothing term $\lambda\max(T)$ to the diagonal elements of the transition
tensors of each schema. This diagonal smoothing term controls the self-
transition probability at each node for each action. This is especially
important for generalization to size and shape variations as the previously
learned schemas can be rigid. For these experiments, we used $\lambda=0.2$.

For each room size and shape, we generate a sequence of random actions and
repeat every action 3 times. We select a starting location corresponding to a
unique observation in the room and execute these actions up to 200 steps and
collect the observations. Same action repetition creates trajectories where
the agent walks farther from the start position with less opportunity to fully
explore the surrounding area. The agent is then tasked with going back to the
start location of the walk by planning a shortest path back to the start
location from the current location. Based on the matched schema, we first
learn the emission bindings from the observation-action pairs. We then use the
grounded schema (i.e., schema + learned emissions) to Viterbi decode the the
current state and the start state, and plan the path back to start state in
the model using max-product message-passing. After executing every action of
this plan, we collect the observations and re-plan when the current plan fails
to take us the goal state. We repeat this for $50$ such random walks in each
test room. Fig. S8 shows the shortcut planning performance for the two example
room types.