Graph schemas as abstractions for transfer learning, inference, and planning
Authors: J Swaroop Guntupalli
Rajkumar Vasudeva Raju
Shrinu Kushagra
Danny Sawyer
Ishan Deshpande
Guangyao Zhou
Miguel Lazaro-Gredilla
Dileep George
Venue: Submitted to ICLR 2023 (2023)
Keywords: Schema learning
abstractions
higher order graphs
perceptual aliasing
aliased graphs
planning
spatial navigation
cognitive science

Abstract: We propose schemas as a model for abstractions that can be used for rapid transfer learning, inference, and planning. Common structured representations of concepts and behaviors---schemas---have been proposed as a powerful way to encode abstractions. Latent graph learning is emerging as a new computational model of the hippocampus to explain map learning and transitive inference. We build on this work to show that learned latent graphs in these models have a slot structure---schemas---that allow for quick knowledge transfer across environments. In a new environment, an agent can rapidly learn new bindings between the sensory stream to multiple latent schemas and select the best fitting one to guide behavior. To evaluate these graph schemas, we use two previously published challenging tasks: the memory \& planning game and one-shot StreetLearn, that 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 further demonstrate learning, matching, and reusing graph schemas in navigation tasks in more challenging environments with aliased observations and size variations, and show how different schemas can be composed to model larger environments.

Under review as a conference paper at ICLR 2023
GRAPH SCHEMAS AS ABSTRACTIONS FOR TRANSFER LEARNING, INFERENCE, AND PLANNING
Anonymous authors Paper under double-blind review
ABSTRACT
We propose schemas as model for abstractions that can be used for rapid transfer learning, inference, and planning. Common structured representations of concepts and behaviors-schemas have been proposed as powerful way to encode ab- stractions. Latent graph learning emerging as a new computational model of the hippocampus to explain map learning and transitive inference. We build on this work to show that learned latent graphs in these models have a slot structure- schemas-thal allow for quick knowledge transfer across environments. In a new environment, an agent can rapidly learn new bindings between the sensory stream multiple latent schemas and select the best fitting to guide behavior. To evaluate these graph schemas, we use two previously published challenging tasks: the memory & planning game and one-shot StreetLeean, that 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 few steps in novel variations of these tasks. We further demonstrate learning, matching, and reusing graph schemas in navigation tasks in more challenging environments with aliased observations and size variations, and show how different schemas can be composed to model larger environments.
INTRODUCTION
Discovering and using the right abstractions in new situations affords efficient transfer learning as well as quick inference and planning. Common reusable structured representations of concepts or behaviors sessssschhassssseeeeee -have been proposed as powerful way to encode abstractions (Tenenbaum al. 2011 Mitchell (Fig [ Humans excel at this ability, argued to be key factor behind intelligence and a fundamental limitation in current Al systems (Shanahan & Mitchell, [022]] A computational model with the ability discover and reuse previously- learredd structured knowledge to behave and plan in novel situations can advance Al systems.
Experimental evidence suggests that some animals have this ability. Rats and mice tend to learn new environments faster if they can reuse past schemas (Zhou et al.] 2021 T'e et al. (007). and there are cells in the macaque hippocampus encoding schemas for spatial abstractions (Baraduc et al. 2019)). Neural circuits in hippocampus (HPC) and prefrontal cortex (PFC) are both implicated in different aspects of schema learning, recognition, update, and maintenance. Recent work in mice suggested that PFC encodes the schema while HPC binds this schema to sensorimotor specifics of task (Sam-) borska et al.] 2022). New experiences are also shown to be learned within single trial if it fits an existing schema, and such knowledge rapidly becomes HPC-independent much faster than typically expected for memory consolidation. Kumaran et al. (016)) proposed an updated complementary learning systems (CLS) theory based on this evidence. but there is far no explicit demonstra- tion of such rapid learning with schema reuse as far as we know. Schema-basee based mechanisms also aid memory consolidation over longer time-scales into reusable knowledge structures (Gilboa & Marlatte] 2017). Even though spatial navigation and memory dominated research in HPC,, is also known to be important for non-spatial declarative memory and relational reasoning tasks. Structured relational representations have been proposed common mechanism that generalizes to spatial and non-spatial tasks and memory (Eichenbaum & Cohen 2014; Stachenfeld et al. 2017). Recent work cognitive maps in HPC model these representations using higher order latent graph struc- tures and show generalization to disparate HPC functions (Whittington et al. 2020: George et al.
2021 Sharma al. 2021 Whittington al. 2021

---

Under review as a conference paper at ICLR 2023
In this work, we take such model and provide concrete computational model of abstractions using graph schemas. We describe how graph schemas can be learned and then reused for transfer learning, quick inference, and planning for behavior in new situations by rapidly learning observa- tion bindings and discovering the best schema online. We build these schemas using clone-structured cognitive graphs (CSCG), computation model of cognitive maps in HPC (George al. [202T). In particular, CSCG work showed that latent high order graphs can be learned in highly aliased settings using a smooth, probabilistic, parameterization optimization. We use navigation as our setting, which requires handling perceptual aliasing, diffi- cult problem in which standard methods fail (Lajoie et al., 2018). Using higher order graphs helps with modeling aliased observations (Xu et al., 2016). CSCGs also generalize beyond spatial naviga- tion to non-spatial domains (Dedieu et al. 2019), so schemas on CSCG can extend beyond spatial navigation. This generalization of graph schemas neatly maps to the idea of useful abstractions as template learning and structure mapping as described Shanahan Mitchell (2022). Our compu- tational model provides a concrete implementation of some of the operations described in that work. For example, their operation seeing similarity can be interpreted as the process of finding which of the known templates best describes the observations in terms of likelihood in model, and binding different instantiations to nodes in a structured graph can be interpreted as learning different emission bindings in graph schema.
the graph learning problem using gradient-based
A Schema matching with binding
Cloned hidden Markov model
CSCG graphical model
experiences schemas
states
clones
observations
observations
Schema extraction
E Schema matching
F Planning goal
environment
model with
w/o bindling
new environment
start
goal
schema
schmmm 2
inferred plan
execution
matched schema with learned
plan
scheem
reapaan
Figure I: Overview of using graph schemas in CSCG model. A. Schemas serve as abstractions that can be matched with new experiences. Multiple schemas can be composed together to model ob- servations that could be aliased. Schematic of cloned HMM with multiple cloned states sharing the emissions to model them in different contexts. C. A graphical model of CSCG with action con- ditioning. Schemas in CSCG are extracted from learned models with transitions and emissions, and then unbinding the specific emissions (shown colors) but keeping the clone structure (shown as node shapes). In a novel environment with different observations, the agent navigates and finds the schema that best explains the observations by learning formed in a new room using the matched schema. As the plan is being executed, new observations used to update the beliefs about current location and re-plan needed.
bindings. F. Planning can be per-
RELATED WORK
Large language models (LLMs) have shown impressive demonstrations of few-shot generalization tasks without re-training, possibly using in-context learning (Brown et al. 2020 Wei al.
2

---

Under review as a conference paper at ICLR 2023
2021).
While LLMs ostensibly use some form of abstractions and rapid binding new imputs, the exact mechanism for this is an active research area (Min et al. 2022).
In navigation setting, rapid generalization for solving tasks in novel environments can be broken down into three stages: exploration, modeling, and planning. Different approaches treat these stages as either separate entangle them in various combinations. Some of the recent work on learning to generalize to novel tasks and environments is predominantly done in a reinforcement learning (RL) framework. Approaches using model-free RL show some generalization and behavior that mimics model-based RL agents in some settings, but are limited in their generalization (Wang et al.) 2016). More recent works such as episodic planning network (EPN) and hierarchical chunk attention memory (HCAM) added episodic memory storage with attention heads to selectively attend and reuse stored memories to rapidly adapt to solving tasks in new environments in few-shot settings (Ritter et al. 2020 Lampinen et a (021). These models match optimal planning only after training for billions of steps. Further, none of these approaches explicitly builds models of the environment and plans 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. In most, if not all, of these settings, observations are not aliased and therefore unique. Model-based RL systems do build and plan on explicit models, however they typically start with known model or require extensive training data build the model (Moerland et al. 2020). By adding more structure to the model and the ability to reuse the learned structures in novel environments, training can potentially be done in fewer samples.
Schema matching is another related area where some recent work focuses on finding correspondence between graphs in different contexts. Work by Crouse et al. (202T) was the first to use the neural network approach to structure mapping, but is restricted to the matching problem and has no active exploration 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 's Progressive Matrices (Webb et al., 2021 Kerg et al. 2022). 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. As far as we know, ours is the first model to learn explicit latent higher-order graph structure model the environment only from observation sequences, and to build schemas on top of that model. These allow rapid transfer, inference, and planning in problems that can be cast as (stochastic) graph navigation.
PRELIMINARIES
Consider an agent navigating in a directed graph. When the agent visits a node in that graph, the node emits an observation. However, multiple nodes may emit the same observation (i.e., they are aliased), so this not enough to disambiguate where in the graph the agent located. Additionally, actions do not have deterministic 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 call this graph G, which defines the environment.
the same action at the same node might result
When the agent performs a sequence of actions Q1,..., an (with discrete a, E {1,.... ... Nactions }), it will as result receive sequence of observations, obtaining the stream Ii...., ZN (with discrete In E (1, Notx )). The goal of learning is to recover the topology of the environment G from se- quences of actions and observations. The goal of transfer is to reuse the topology of the graph when the nodes are relabeled with new observations with or without ceuucture-preservate Additionally, we can exploit the learned graphs for further tasks, such goal-directed navigation.
3.1 CLONE-STRUUCURED COGNITIVE GRAPHS
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 a categorical hidden variable 21...., ZN to model the node of the graph that the agent is at in each time step. With this, it is possible to formulate graphical model for the sequence of observations given the actions. If we take the conditional version of their model. but write the emission probabilities explicitly, the connection with hidden Markoy models (HMMs)
3

---

Under review as a conference paper at ICLR 2023
becomes more obvious:
N
N P(an an-1) 1) II P(nn|zn) n=1
P(xlaa) - E P(z) II
where we use the shorthand {a) ..... ................. .NJ. {ai ...., aN..,N,, {21 ...,,... 2 Additionally, we define a transition matrix T with elements Typ = P(an = k) 2n-11 = J. an-1 = i); and an emission matrix E with elements Fij- P( žn i).
For an unrestricted E, this is an action-conditional version of an HMM. What makes this a CSCG the structure of the emission matrix: each row contains single 1 and all the remaining entries are zero-equivalently, emissions are deterministic. Multiple hidden states can emit the same ob- servation, however. All the hidden states that emit the same observation are called the clones of that observation. Also, unlike HMMs, learning proceeds by choosing particular allocation of clones to observations, encoding in E, and keeping E fixed throughout learning.
Learning a new graph G describing particular environment only involves learning T, i.c. max- imizing P(x a; T, E) wrt T' (See Appendix A.I for details). Transfer learning between different environments is possible too: if we know that new environment has the same topology as previ- ous one for which T has already been learned, but with different emissions associated to its nodes, we can maximize P(z T, E) wrt E while keeping 7 fixed to its known value from the previous environment (See Appendix A.2 for details). Further, if we know that the new environment preserves the emission structure (when two nodes emitted the same observation in the original environment, they also do so in the new environment). then we can further restrict the learning of E, with all the of E that correspond to the same observation in the original environment sharing the same pa- rameters. Learning can be efficiently accomplished using the EM algorithm, analogous to learning in HMMs, but with some computational advantages derived from the structure of E. The graphical model is depicted in Fig. IB-C.
3.2 SCHEMAS FOR TRANSFER LEARNING
As explained in the previous subsection, is possible reuse parts of a previously learned CSCG to learn new environment faster. The portion of the CSCG that gets transferred across environment is called schema, which consists of the transition matrix together with its unbinded emission structure Drig(DD) For example, in the room navigation setting, a schema models the locations in room and how actions move the agent between them, as well as the knowledge that the floor or door can look the same at multiple locations in the room. Schemas allow for rapid model learning in new environments that have matching topologies and emission structures with fast binding (Fig. [)). Inference can be performed with the matched schemas 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.
RESULTS
4.1 MEMORY PLANNING GAME
We first evaluate our model on the Memory & Planning Game (MPG) proposed by Ritter et al. 2020) to evaluate rapid adaptation game, the agent can navigate on a 4 x 4 grid, observing symbols, and the task is to collect rewards at a specified goal location that grid (Fig. ( After collecting a reward, the agent re-spawned in new location and a new goal is specified. All grid positions have unique symbols and the symbol- position mapping is randomized after each episode, which lasts for 100 steps. See B.T for details. This setup lets us evaluate and the agent needs to explore to collect the observations and bind them rapidly to the schema to maximize rewards with optimal navigation.
new environments and solve tasks in that environment. In the
model and schema reuse the structure is maintained across episodes
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 leam CSCG model that best explains the experiences across all episodes so far observed. We reuse the same schema (T) across all episodes and learn new binding (emission matrix F.) per episode. It takes
4

---

Under review as a conference paper at ICLR 2023
Memory & Planning game
30
Sample environment
Learned graph
No. of steps to goal
EPN
Aver rage reward
Mittin
10
tack
cotim
24
training steps
task no. in episode
start
goal
Learned graph
30
No of stops to goal
Average reaad
StreetLear
-00
EPN CSCG
10
10
20
training stops
X10'
task no. in opisade
Figure 2: Performance on MPG and One-shot StreetLLeam, A. An example MPG episode and the latent graph learned by the CSCG. B. Average reward in MPG. CSCG performance remains the same after 9 episodes (900 steps), so we plotted the average. EPN performance is shown on evaluation at different stages of training. C. Number of steps get the reward in subsequent tasks in an episode. CSCG explores first without collecting the reward, so our first task takes longer to complete but subsequent tasks do optimal planning. D. Example One-Shot StreetLearn environment, with agent's s observations. Right panel shows the latent graph learned by CSCG. E. Average reward across 100 episodes of Streetl Learn in the test city. Notation similar to MPG plot. Number of steps to get the reward in subsequent tasks in an episode. First few tasks take variable numbers of steps as we collect rewards incidentally. After that, the CSCG agent collects all subsequent rewards in optimal steps. EPN baselines and optimal values for both environments re-plotted using the data from Ritter et al. (020). Error bars are 95% CI of the SEM.
only 9 episodes (900 steps) to learn perfect schema of this 4x4 4 grid environment and in subsequent episodes, we rapidly learn the schema bindings and do planning to maximize the rewards. Learning an efficient exploration policy to cover all observations in a new episode is an interesting problem, but, since our focus in this work is schema learning and reuse, 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 + standard error of the mean (SEM) after learning the schema is: 17.3 + 0.57 for random and 26.4 + 0.17 for optimal exploration policy, which is comparable to Episodic planning network (EPN) (~it))Ritter et al.) (020). In contrast, EPN takes more than 10 million episodes (> 10 steps) of training reach its optimal performance (Fig. 2). Planning in our model is optimal, on par with EPN and the oracle. Fig. 2C shows the average number steps taken to collect rewards in subsequent tasks. 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 our case (18 0.09 steps) than EPN, but the average reward in an episode is comparable.
4.2 ONE-SHIOT STREETLEARN
One-Shot StreetLearn is another task proposed by Ritter et al. (2000) based on the StreetLearn task MMirowski et al., 2019) to evaluate rapid task solving ability a more challenging setting with varying connectivity structure across episodes (Fig.PP) In each episode, the agent is placed in neighborhood city and the task is to navigate to a goal, specified by the goal street view image, and collect the reward. After collecting reward, the agent is re-spawned in a new location and new goal is specified. Each episode lasts 200 steps and a new neighborhood is sampled for the next episode. Unlike the MPG, the transition graph changes every episode. We evaluate the ability of CSCG to rapidly learn the model an episode and to navigate optimally to the goals to maximize the rewards. Note that there is schema reuse in this setting: we learn model for every
5

---

Under review as a conference paper at ICLR 2023
episode. This showcases CSCG's ability to learn model rapidly within 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, follow an explore and exploit strategy with a hard coded exploration policy. During exploration, agent navi- gates every action from every observation 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, agent plans with CSCG and collects wards (See Appendix [32 for details). Average reward + SEM over 100 episodes is 21.7 + 3.7 (Fig. E). This performance lower than EPN (28.7) as our exploration strategy is not optimal. 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 0.03 steps, which matches with the optimal value (Ritter et al. 2020) 2 Note that we do not transfer any learning across episodes in this setting since the graph changes every episode, unlike the MPG. In cities with re-usable StreetLearn graph structures such as grid layout in Manhattan,
reuse CSCG schemas similar MPG and benefit from reuse. We evaluate this schema reuse in detail and in much harder setting in the following experiments on navigating in rooms with extensively aliased observations.
ecample time
corresponding sequiance of aliased observations
room type
size
cylinder
rectangle
square hole
torus
U-shape
NLL
schema
cylinder
250
250
0
750
750
250
rectangle
square torus U-shape
medium
WLLL
Schema size mediumm types.
large
NLL
stops 250
0
stops 250
250
of 250 % no. of stops 250
Figure 3: 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 the agent does random walks. The schema with the least negative log likelihood (conversely the highest likelihood) considered the best match. Each panel corresponds to a room of type and size indexed by the column and row headers, respectively. Schemas trained in medium sized rooms and tested on two other size variations. Error bars are 95% CI of the SEM.
4.3 SCHEMA MATCHING IN HIGHLY ALIASED ENVIRONMENTS
In the next set of experiments, we use a more challenging setting with larger environments and extensive aliasing, and evaluate schema matching novel environments varying in size and obser- vations from the trained environments We start with set of environments of different shapes and topologies with extensive aliasing. Analogous to a large empty arena, observations of these environments are identical (perceptually aliased: Lajoie et al.) 2018), see Fig. 3 top row As described in section 3.2 we learn the schemas in training rooms (Fig. IP). We evaluate schema matching test rooms with similar layout but with novel observation mapping and size variations. In a novel test room, the agent takes a random walk and evaluates the likelihood the observation sequences given the actions executed under different learned schemas (Fig. TE. Note that this
the interior

---

Under review as a conference paper at ICLR 2023
quires 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 different intervals during the random walks. We used the parameters specified in Appendix B in the following experiments unless otherwise specified.
CSCG schemas rapidly learn matching bindings in novel environments that differ in size and observations from the training environments. We generated rooms with different layouts in sizes: cylinder, rectangle, square with hole, torus, and U-shape (Fig. 3 We learn CSCGs for the medium-sizee sized versions of these rooms from a random walk of 50, 000 steps. the test rooms with novel observation mapping and 3 different sizes per shape, the agent takes random walk while we learn the new F. and evaluate the likelihoods of these observations under all schemas every steps. We compute these values up to 250 steps. We repeat this 25 times in each test room with different random walks, See Appendix [B.3] for details. Fig. 3 shows the negative log likelihoods of all test room under all schemas. By reusing clone structure, we able to correctly match the schema in all cases by 50 steps demonstrating a rapid matching and adaptation to novel environments with size and observation variations (See 8 for results without using clone structure).
Example MNIST room
B
recognition
Schema recognition composed rooms
room:
Probability of being room
0 9
0.7
cchama dex
probabiltty
0.6
ALE
00.5
0.4
0.3
no. of S1006
Figure 4: A. An example MNIST digit used to train schemas. B. Negative log-likelihood (NLL) vs no. of steps in the test room corresponding to digit O. We are able to recognize the correct schema in fewer than 50 steps. Error bars are 95% CI of the SEM. C. The heatmap represents the probability of being in schema 0 at each location the composed room, averaged over all instances of the agent landing at that location during a random walk of length 20, 000 steps.
Schema matching in MNIST rooms. In this experiment, we used a set of ten digits from the binarized MNIST dataset our room layouts they provide an interesting variety of shapes and topologies. All the foreground pixels are treated as accessible locations; background pixels are treated obstacles. Observations at each pixel are sampled in similar way the shaped rooms with aliased interiors (see Fig. 4 Appendix Fig. 9 for examples of the training rooms). We repeat our schema matching experiments as before in this new set of rooms. We first trained schemas on rooms corresponding to all ten digits. To evaluate schema matching, we generated test rooms with the same structure but with new observation mappings. Fig. 4 shows the negative log likelihoods of observation sequences in test room 0 under different schemas (see Appendix Fig, [0 for all rooms). The correct schema was identified in at most 50 steps in all test rooms. (See Appendix B.4 for more details).
4.4 SCHEMA MATCHING IN COMPOSED ENVIRONMENTS
To demonstrate how schemas could be used in more complex environments, we constructed test rooms composed of pairs of overlapping MNIST digits (Fig. [FC). 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 the correct schema at the current time-step. The probability of being in a particular schema can then be computed using the softmax of the log-likelihoods (see Appendix B.4] for more details). In Fig. 4c, we show the probability of being in room at all locations in a room composed of the digits 0 and 9, using random walk of length 20, 000 steps and sliding window of length 200 steps. We can also compute the accuracy of identifying the correct schema at all locations by thresholding these probabilites. The average accuracy of identifying the correct schema at all locations in all composed test rooms is 88.4 + 0.9%. See Appendix Table for results on all composed rooms.
7

---

Under review as a conference paper at ICLR 2023
Walk in test (14 14)
Matching schema
Planning/re-planning to reach goal
[9
plan
plan
plan
plan
# intbal final
geal
infereed
Rectangle schema of medium size
Square with hole schema of medium size
srouilingg wit ensenning
w/o smouti with smeething
Manhstaa distance to geal
nffo arsuilling wet amoot ting
smoutl ing with amoothing
8
No. of plans
Marhotan
13
()++))(-+))
100m (6++2)((8+4)
room (6+ (8++
sire (3+A/12) (3+8/77)
hale (3+R/22) (3+R)
Figure 5: Schemas enable rapid and robust navigation planning in new environments. A. Shortcut finding with re- re-planning in a larger room. Given partial observations in a novel room (14 x 14 square with a x 5 hole that is off-center), the best matching schema (9 x 9 square with 3 x hole in the center) can be used to rapidly navigate around obstacles and find the shortest path to goal. B. Shortcut finding performance as measured by the Manhattan distance from the goal and the number of plans/re-plan to reach the goal, as a function of room size. Error bars are 95% CI of the SEM.
4.5 RAPID PLANNING IN NOVEL ENVIRONMENTS WITH GRAPH SCHEMAS
Rapid schema matching and binding enables planning in novel environments with limited experi- ence. In this experiment, an agent first executes random walk in test room. At the end of the walk, the agent has to find the shortest path the starting location using the best matching schema. The agent first uses the observations to learn E in the new environment. Next, it uses the schema and the binding
infer the shortest path its goal. After executing this plan, if the agent estimates that it has not reached its goal after accounting for the new observations, it is possible that there agent initiates re-
schema mismatch or the estimated E inaccurate. When this occurs, planning. 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- re-planning steps (Fig. 5A).
We evaluate shorteut finding performance in terms of Manhattan distance from the goal location. We used 50 trials with initial random walks of length 200 steps. In Fig. 5B, we show how the per- formance degrades gracefully We also show how the number of planning attempts required to reach the goal increases as the size mismatch increases. We also evaluated shorteut finding without diagonal smoothing (described in Appendix B). The diagonal smoothing term adds self-transition probability at each node for each action, and allows the agent account for observation mismatches in aliased settings. Without this smoothing, the agent never reaches the goal in larger size variations of the test rooms.
the size mismatch between the schema and the test room increases.
4.6 SCHEMA-BASED NAVIGATION IN SIMULATED 3D ENVIRONMENTS
All our previous experiments focused on 2D grids, but we can use schema-based transfer learning to quickly navigate in novel simulated 3D environments. To demonstrate this, we use 3D environ- ment (Beattie et al. (016)) with a T-maze layout (Fig. (A). The agent can navigate with 3 discrete actions (move forward, turn left, turn right) and the observations are RGB images. We use vector- quantization (k-means) get cluster indices from the raw RGB observations and use these indices as input observations to learn a CSCG model of the training environment (Fig: 63). Fig. 6 shows latent graph learned by the CSCG after training using random walk. We generate test 3D environ- ment with the same layout but with different colored walls, floor and environment lighting (Fig.6p) This corresponds entirely new RGB observations for agent navigating in this room. The agent
8

---

Under review as a conference paper at ICLR 2023
A
Shortcut finding test
Input image
Train maze
Aoservationa
goal
Vecter Quantizer
start inferred path
T
Test maze
cluster cantar
Learned transition graph
used
Figure 6: A. Top views of the train and test mazes. B. Schematic of vector quantizer used to map input images to discrete space. C. The mapped observations then used learn CSCG for the train maze. D. At test time, the agent quickly learns the emission matrix in the new room and uses the schema to navigate to goal state.
walks in the test room and learns an emission matrix (new bindings for the learned schema) in fewer than 90 steps. Using this model, the agent can plan and navigate goals in the room (Fig. 6P).
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 (Zhou et al.) 2021 "Tee et al.) 2007 Kumaran et al., 2016). We have proposed a concrete computa- tional model for this 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 new environments.
CSCG schemas learned graph structure in far fewer episodes than deep RL agent and matched optimal planning MPG and One-Shot StreetLearn tasks. 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 different schemas. We showed successful planning to goals in rooms with shape and size variations from schema by re-planning and updating the model while walking to the goal and this can be adapted to simple 3D) environments. There clear directions for potential future works, some of which listed below.
Schema learning from experiences. We learned schemas independently and explicitly in this work, but the real world, it might not be feasible to have access differentiated experiences belonging to distinct schemas. Learning reusable schemas from a continuous stream of experiences could be an interesting future work.
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 match. This can thought as keeping specific memories versus using abstract schemas.
Schemas vs memories.
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 perceive and remember experiences that fit in their existing schemas and develop the flexibility later (Piaget & Cook 1952). Similarly, we could update schemas based on new experiences, and even make 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-conditionel beliefs on future observations, and by choosing actions that could optimally disambiguate among different schemas, we could potentially do better than random exploration. Similarly, instead of random exploration to learn new environ- ment and schema, we could direct the exploration policy by composing known schemas (Sharma al. 202T), and even actively learn them while exploring.

---

Under review as a conference paper at ICLR 2023
REFERENCES
P Baraduc. J-R Duhamel, and S Wirth. Schema cells in the macaque hippocampus. Science, 363 (6427):635 639, February 2019.
Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Ander- son, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. DeepMind lab. December 2016.
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 few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Gordon V Cormack and R Nigel S Horspool. Data compression using dynamic Markov modelling. The Computer Journal, 30(6):54--550 1987.
Maxwell Crouse, Constantine Nakos, Ibrahim Abdelaziz, and Ken Forbus. Neural analogical match- ing. AAAI, 35(1):809-8117 May 2021.
Antoine Dedieu, Nishad Gothoskar, Scott Swingle, Wolfgang Lehrach, Miguel Lázaro-Gredilla, and Dileep George. Learning higher-order sequential structure with cloned HMMs. May 2019.
Howard Eichenbaum and Neal J Cohen. Can we reconcile the declarative memory and spatial navigation views hippocampal function? Neuron, 83(4))764-770, August 2014.
Dileep George, Rajeev Rikhye, Nishad Gothoskar, Swaroop Guntupalli, Antoine Dedieu, and Miguel Eizzro-Gredilla.. Clone-structured graph representations enable flexible learning and vi- carious evaluation of cognitive maps. Nat. Commun., 12(1):23,2, April 2021.
Asaf Gilboa and Hannah Marlatte. Neurobiology of schemas and Schema-Mediated memory. Trends Cogn. Sci., 21(6):618-66 August 2017.
Saurabh Gupta, James Davidson, Sergey Levine, Rahul Sukthankar, and Jitendra Malik. Cognitive mapping and planning for visual navigation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 161662625,, 2017.
Giancarlo Kerg, Sarthak Mittal, David Rolnick, Yoshua Bengio, Blake Aaron Richards, and Guil- laume Lajoie. Inductive biases for relational tasks. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
Dharshan Kumaran, Demis Hassabis, and James McClelland. What learning systems do intelligent agents need? complementary learning systems theory updated. Trends Cogn. Sci., 20(7):512-534, July 2016.
Pierre Yves Lajoie, Siyi Hu, Giovanni Beltrame, and Luca Carlone. Modeling perceptual aliasing in SLAM via Discrete-Continune graphical models. October 2018.
Andrew Lampinen, Stephanie Chan, Andrea Banino, and Felix Hill. Towards mental time travel: hierarchical memory for reinforcement learning agents. Advances in Neural Information Pro- cessing Systems, 34:281 1 1111182221955 2021.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes In-Context learning work? February 2022.
Piotr Mirowski, Andras Banki-Horvath, Keith Anderson, Denis Teplyashin, Karl Moritz Hermann, Mateusz Malinowski, Matthew Koichi Grimes, Karen Simonyan, Koray Kavukcuoglu, Andrew Zisserman, et al. The streetlearn environment and dataset. arXiv preprint arXiv: NXXXv:9903001299 2019.
Melanie Mitchell. Abstraction and making in artificial intelligence. Ann. N. Acad. Sci., 1505(1):79-101, December 2021.
10

---

Under review as a conference paper at ICLR 2023
Thomas M Moerland, Joost Broekens, Aske Plaat, and Catholijn M Jonker. Model-based reinforce- ment learning: A survey. June 2020.
Jean Piaget and Margarer Trans Cook. The origins of intelligence in children. 1952.
Sam Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt Botvinick, and David Raposo. Rapid Task-Solving in novel environments. June 2020.
Veronika Samborska, James L Butler, Mark E Walton, Timothy E J Behrens, and Thomas Akam. Complementary task representations hippocampus and prefrontal cortex for generalizing the structure of problems. Nat. Neurosci., pp. 1-13, September 2022.
Murray Shanahan and Melanie Mitchell. Abstraction for deep reinforcement learning. February 2022.
Vatsal Sharan, Sham M Kakade, Percy S Liang, and Gregory Valiant. Learning overcomplete hmms. In Advances Neural Information Processing Systems, pp. 940-949, 2017.
Sugandha Sharma, Aidan Curtis, Marta Kryven, Josh Tenenbaum, and lla Fiete. Map induction: Compositional spatial submap learning for efficient exploration in novel environments. October 2021.
L. Stachenfeld, M M Botvinick, and $ Gershman. The hippocampus as a predictive map. Nat. Neurosci., 20, 2017.
Joshua B Tenenbaum, Charles Kemp, Thomas L Griffiths, and Noah D Goodman. How to grow mind: statistics, structure, and abstraction. Science, 13166022) March 201
Dorothy Tse, Rosamund F Langston, Masaki Kakeyama, Ingrid Bethus, Patrick A Spooner, Emma R Wood, Menno P Witter, and Richard M Morris. Schemas and memory consolidation. Science, 31665821):76-82, April 2007.
Jane Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. November 2016.
Taylor Whittington Webb, Ishan Sinha, and Jonathan Cohen. Emergent symbols through binding in external memory. In International Conference on Learning Representations, 2021. URL https:
Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe Lc. Finetuned language models are Zero-Shot learners. September 2021.
James C Whittington, Timothy Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy E J Behrens. The Toomman-Eichenbaum machine: Unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5): 1263.e23, November 2020.
James R Whittington, Joseph Warren, and Timothy E J Behrens. Relating transformers to models and neural representations of the hippocampal formation. December 2021.
CF Jeff Wu et al. On the convergence properties of the em algorithm. The Annals of statistics, 11 (1):95-103, 1983.
Jian Xu. Thanuka L Wickramarathne, and Nitesh V Chawla. Representing higher-order dependen- cies in networks. Sci Adv, 1600e26 May 2016.
Jingfeng Zhou, Chunying Jia, Marlian Montesinos Caatessnos-CCrttgenna. Matthew P Gardner, Wenhui Zong, and Geoffrey Schoenbaum. Evolving schema representations in orbitofrontal ensembles during learning. Nature, 590(7847):606 61 February 2021.

---

Under review as a conference paper at ICLR 2023
A
EXPECIATION-MAXIIIZATTOOI LEARNING OF CSCGS
Cloned Hidden Markov Models (HMMs), first introduced in (Dedieu et al., 2019), sparse restriction of overcomplete HMMs (Sharan et al. 2017) that can overcome many of the training shortcomings of dynamic Markov coding (Cormack & Horspool 1987). Similar HMMs, cloned HMMs assume the observed data x = (II.... ENn are generated from hidden process = ,,x...,N) that obeys the Markovian property,
N P(x)= = P(2) II P(znlan-1) 2=2
N II P( (2)) ==1
Here P(21) is the initial hidden state distribution, P(znlzn- is the transition probability from Zn-1 to Zns and [(.) is the probability of emitting Vn from the hidden state 2n-
In contrast to HMMs, cloned HMMs assume that each hidden state maps deterministically to a single observation. Further, cloned HMMs allow multiple hidden states 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(x) = Σ
E ExEC(xx)
P(a)) II ((zalza-1)
21EC(es)
where C(In) corresponds to the hidden states (clones) of the emission Zn.
Clone-structured cognitive graphs (CSCGs) build on top cloned HMMs by augmenting the model with the actions of an agent (01.... ,ay (George et al. 2021).
a(r,a)=
Σ
Σ
P(a)
II n
P(zn, an 1
21(((=1) ExEC(Ex)
A.1 LEARNING THE TRANSITION MATRIX WITH EMISSIONS FIXED
Learning CSCG in new environment requires optimizing the vector of prior probabilities The
P(z1 = k) and the action-augmented transition matrix T: Tijx- = P(z, = k| = an-1 = i). The emission matrix F, which encodes the particular allocation of clones observations, is kept fixed throughout learning.
The standard algorithm train HMMs is the expectation-maxiiizatin (EM) algorithm (Wu al. 1983), which in this context is known as the Baum-Weldh algorithm. Learning CSCG using the Baum-Welch
Welch algorithm requires 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 Nots and Naciions 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 E (1,...,No Nots and u. E {1...., ......,,ou. Nactions }. The submatrix Tuu, u, w) contains the transition probabilities P(z,, an-1 u1) for .n C(v) and 2n E C(w), where C(v) and C(w) correspond to the clones emissions e and w respectively.
The standard Baum- Welch equations then be expressed in simpler form in the case of CSCG. The E-step recursively computes the forward and backward probabilities and then updates the pos- terior probabilities. The M-step updates the transition matrix via normalization.
E-Step:
a(1) = T(x1)
a(n) T a(n - 1 (an-1, En-1, ((n) T(an, Fny In+1) (,
B(N) 1xx)
12

---

Under review as a conference paper at ICLR 2023
Suvu (n) a(n)) o T(an v. w) o 3(1 + + a(n)) T w + 1)
7(n) a(ni) o 3(n))
((m)T((m))
M-Step:
T(ボ1) 7(1)
N
N
Nacion Nots
N
Euer (n) 2 T(u, v, w)
Swew (12)
where and 2 denote the element- product and division, respectively (with broadcasting where needed). All vectors are Nap x column vectors, where Nope is the number of clones per emission. We use 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 CSCG, transfer learning between different environments can be accomplished by keeping its transition probabilities T fixed and learning the emissions associated to its nodes 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 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:
Namss
ā(n) = ((T(an11) an-1)nn) E(zn)
Pa(n) Σ (x(n)
a(n) a(n)/pa(n)
N
B(n) = T(an) o E(znt1))
p2(n)= = Σ (n)
8(n) 3(m)/pa(n)
k=1
7(1) a(n) B(n)
an)) (g)))))))))))))))))))))( (())
M-Step:
N EG) = Σ (n.= Ø (n)
N
n=1
Note that the updates here involve Nelones X vectors, where Nekones the total number of hidden states in the model. T(a,) corresponds to the transition matrix for the action a, and is of size Nelones x Nelones. E(I,) is column of the emission matrix corresponding to the emission In. and 1z,=) an indicator function. The forward message is initialized ā(1) π o E(x1); and the backward message is initialized as a vector of all Is.
When the clone structure is to be preserved, the 7(1) term in the E-step is modified follows. For each observation j E {1,...,,oe Noss,, we set the posterior probability for all clones of j to be the same:
nu (m) Σ n(n) Vk E C(j).
LECU)
13

---

Under review as a conference paper at ICLR 2023
For given T and E. the normalization term Pa(n) for the forward messages in the E-step can be used to compute the negative log-likelihood (NLL) of a sequence of Doservation-actionn pairs as follows:
NLL =
log Pa (n)
N 1=
B EXPERIMENT DETAILS
In George et al. (2021). b, it was observed that the convergence of EM for learning the parameters of a CSCG can improved by using a smoothing parameter called the pseudocount. The pseudocount is a small constant that is added to the accumulated counts statistics matrix (C), which ensures that any transition under any action has non-zero probability. This ensures that the model does not have zero probability for any sequence of observations test time. During the schema learning phase (learning T, with E fixed), we use a pseudocount of 2 x 10-3, During the schema matching phase (learning E, with T fixed),
pseudocount of 10-?.
In the schema matching phase, we apply an additional smoothing term À max(C") to the diagonal elements of C. This diagonal smoothing term controls the self- transition probability at each node for each action. This is especially important for tasks such as shortcut finding, where we want a schema to adapt to test rooms with size and shape mismatch. For all our experiments that involve learning E, with 7' fixed, we used À = 0.2, unless otherwise specified. These values for the pseudocount and the diagonal smoothing were obtained using hyper-parameter sweeps.
B. MEMORY & PLANNING GAME
The game environment is a 4 grid of symbols and the agent can navigate in the four canonical directions by one grid step (up, down, left, right) and collect reward of at goal location in the grid (Fig. A). Reward is otherwise. Once the agent collects the reward at the current goal, the agent is placed in new random position and
goal symbol is sampled to which the agent must navigate to collect the next reward. All grid positions have unique symbols and the symbol-ppsitinng mapping is randomized after each episode, which lasts for 100 actions. The agent's observation is tuple of symbol in its current position and the goal symbol. See Ritter et al. (2020) for more details. We assume knowledge of collect action function and execute only when the goal symbol reached.
We employ 3 different hard coded exploration policies to cover the observations in a new episode: random navigation actions, random navigation actions but limited (up. right), and optimal set of actions to cover the learned transition graph: Hamiltonian cycle that repeats (up, up. up, right) four times. For a new episode, we take actions based on one of these policies, learn E at every step, and plan to reach the current goal using the current estimate of F. Planning a path from the current position to a goal position achieved via Viterbi decoding in the CSCG. If the probability of reaching the goal is above 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 probability threshold. execute actions based on the policy, otherwise
the current goal symbol is not yet observed
in this episode, planned set of actions reaching the goal symbol. As we navigate, the estimate of E' becomes better and the plans are more likely succeed. Our planning algorithm takes both these uncertainties into account. We evaluated for total of 100 episodes (10* 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.c. after learning.
compute the probability of our
B.2 ONE-SHIO STREETLEARN
In each episode of this task, the agent is placed in neighborhood of a city and needs to move to a goal location and direct itself to the target direction, specified by the goal street view image, and 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 new goal street view image. The agent needs to collect as many
14

---

Under review as a conference paper at ICLR 2023
rewards as possible during the episode duration of 200 steps. after which neighborhood is sampled for the next episode. StreetLearn environment represents the agent's perception, a street unique compressed string that agent, we follow an explore and exploit strategy. The agent first navigates every action from every observation it encounters while collecting the rewards exploration to cover every possible edge in the transition graph. A CSCG model is then learned from this experienced sequence of observations and actions in this episode. For any subsequent step, 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 in model expectation and the reality.
view image.
agent's observation. For the CSCG
encounters the goals. This is a guided
B.3 SCHEMA RECOGNITION IN ROOMS WITH SIZE VARIATION
We generated rooms of five different types, and three size variations per type, as shown in Fig. 7 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-hottom edges for the cylinder, and the top-bottom and left-right pair of edges for the torus, as indicated by the blue arrows Fig. 7 We learned schemas on the medium size variation of these room types. For cach training room, we used aation-obseevattion sequences from a random walk of length 50, 000 steps to train the respective CSCG.
For evaluating schema recognition, used test rooms of all three size variations. Importantly, generated the test rooms by permuting the observations from the original room. As a consequence, test time, our model would have to relearn the observation mapping.
In figures 3 (main text) and [] we show the schema matching performance with and without the of clone-structure, respectively. We observe that it takes at most 95 steps to identify the correct schema for all test rooms when using clone structure. On the other hand, it takes at most 255 steps to find the best matching schema when the clone structure is ignored. We obtained these results using the paired T-test over 25 random walks, per test room, and a p-value threshold of 5%.
B.4 MNIST ROOM EXPERIMENTS
We learn schemas for 10 selected MNIST-digits (Fig. 9 using action-observation pairs from random walks of length 100, 000 steps in each room. For experiments with composed MNIST rooms, we selected the sliding window size of 200 by doing a hyper-paramete, sweep. The algorithm for schema selection is presented below.
Algorithm Schema selection for composed rooms Input:
Schemas SH.
Observations and actions of length N in the test room R.
Window size
n + w while <= N do
Estimate E, using Fn- wan an-win for each schema, E (1...., H}
Compute Ln log likelihood of 2n-w:ns an-w:n under jchema
Compute Pn.j = softmax( L,); the probability of being in schema j at time step n. Select schema with highest pn.j- end while
In Table I
we present the accuracy of schema selection for multiple composed rooms.
B.5 RAPID INFERENCE AND PLANNING USING SCHEMAS
For each room size and shape, generate sequence of random actions and repeat every action times. We select starting location corresponding to a unique observation the room and execute these actions up to 200 steps and collect the observations. Same action repetition creates trajectories
15

---

Under review as a conference paper at ICLR 2023
Composed digits
Mean accuracy
Standard deviation of accuracy
(averaged over five runs)
Digits 2 and
87.66
4.14
Digits 2 and
92.47
1.51
Digits 0 and
87.88
2.77
Digits 4 and 6
85.89
2.58
Digits 5 and
84.28
2.74
Digits 5 and 3
91.14
2.76
Digits and
85.21
3.87
Digits 8 and 7
90.86
4.04
Digits and
90.38
2.79
Table 1: Accuracy of the schema predictions the agent moves across various locations in com- posed MNIST rooms. We use the ground truth locations to compute this accuracy. Note that for the composed room experiments, the agent gets two schemas as input (H = 2). Further, the accuracy is averaged over five trials of random walks in each room.
where the agent walks farther from the start position with opportunity to fully explore the sur- rounding area. The agent is then tasked with going back to the start location of the walk by planning a shortest path back the start location from the current location. Based on the matched schema, first learn bindings, emission matrix E, from the observation-action pairs. We then use that model to Viterbi decode the the current state and the start state, and plan our path back to start the model using max- max-product observation and re- re-plan when the current plan fails to take us the goal state. We repeat this for 50 such walks for each test
passing. After executing every action of this plan, collect the
B.6 SIMULATED 3D ENVIRONMENTS
In the simulated 3D environments, observations are 64 x 64 pixel RGB images, and the action space discretized to 3 egocentric actions: go forward by 1m, rotate 90" left, rotate 90º right. Thus, the agent is only allowed to walk on grid locations 1m apart. We convert the input RGB images into categorical values using a k-means vector quantizer with & = 36. For training CSCG on the example T-maze, we collected a random walk of length 100, 000 steps in the maze.
We also tried other quantization techniques like the vector quantized variational auto encoder (VQ- VAE). We first trained VQ- VAE, with number of embeddings 8, images from multiple 3D rooms. An example transition graph of a CSCG trained using the VQ-VAE to generate observations is shown in I1 Even in this case, observe that able to correctly recover the structure of the maze.
16

---

Under review as a conference paper at ICLR 2023
room type
cylinder
rectangle
square with hole
torus
V shape
7x7
7x7
4
x
small
6
x
9
8
x
medium
13x13
x3
9x-11
9x-11
large
Figure 7: Rooms of different types and sizes used in the schema matching experiments.
Schema matching without using clone structure
room type
size
cylinder
rectangle
square hole
U-shape
NLL
type: cylinder
500
%
rectangle
square w. h.
medium
NLL
U-shape
500
500
500 %
500 0g
500
Schema size
mediumm
types.
large
NLL
% no. of steps
0 no. of steps
no. steps
no. steps
no. of steps
Figure 8: We measure the negative log likelihood (NLL) of observation sequences in a given room under different schemas, without using the clone structure. Each panel corresponds to a room of type, size indexed by the column and row headers, respectively. Error bars are 95% CI of the SEM.
/
7
9
Figure 9: Images from the binarized-MNNN -NNIST datasct used for our schema matching experiments.
17

---

Under review as a conference paper at ICLR 2023
room
room
room
room
room
schema digit
2
3
250
250
250
250
0
250
test room 5
test room
6
test room = 7
test room =
test room 9
5
6
8
of steps 250
of 250
of steps 250
of steps 250
of steps 250
Figure 10: Schema recognition for MNIST rooms. We measure the negative log likelihood (NLL) of observation sequences in given room under different schemas. Each panel corresponds to test room indexed by its header. Error bars are 95% CI of the SEM.
Figure 11: Learned transition graph of CSCG trained on the T-maze using observations from VQ-VAE quantizer.
18


Reviews:
[{"review_id": "oD7gG4KVM9J", "reviewer": "Reviewer_Zycs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes and analyzes a computational model for the learning and transfer of schemas as a mean for abstraction. A computational model that decouples learning a schema from observation/emission is suggested, and its performance across a range of empirical simulations is assessed.", "summary_of_the_review": "The paper provides a good empirical analysis of a computational model using graph schema for abstraction and learning. More discussion about the model and the sharing of the experimental data would enrich this contribution.", "strength_and_weaknesses": "The paper clearly illustrates the strenghts and the limitations of the model on a number of experimental settings designed to highlight different features of the model. I found the presentation of the model in Section 3.1 quite synthetic, and I would have appreciated a more detailed definition of it; also, in Equation (1) I am puzzled by the first \\sum whose range is not well specified and by P(z_1). ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lP-8kDbmsms", "reviewer": "Reviewer_RKUK", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors present a method which employs an instantiation of Clone-structured cognitive graphs to learn structured representations of environments, which can be used downstream for transfer, inference and planning. The performance of the proposed approach is evaluated using grid-world style navigation tasks in primarily 2D and in one 3D environment.", "summary_of_the_review": "Currently the paper has a ton of scope for improvement (see my comments above) and I cannot recommend acceptance. Hence, my overall score is going to be a weak reject.", "strength_and_weaknesses": "Strengths:\n\nThe paper is trying to address a very important research problem. The high-level idea of using graphs as schemas is a very cool concept and I like the overall direction a lot.\n\nImprovements:\n\n1. The work introduces CSCGs as graph schemas for general tasks but only tests it on task with primarily navigation objectives. If the scope of the work is to use graph schemas for only navigation tasks, then this needs to be explicitly stated early in the paper.\n\n2. The related work is very much lacking. References are completely lacking for many related works on using graph-structured schemas to represent long-duration temporal tasks, for much more complex tasks in the real-world in various (supervised, unsupervised and RL) settings. The paper also cites LLMs in first paragraph of section 2 and I'm unsure how they are relevant here. Examples of related works:\nA. Multi-modal Cooking Workflow Construction for Food Recipes (https://arxiv.org/abs/2008.09151)\nB. Predicting the Structure of Cooking Recipes (https://aclanthology.org/D15-1090.pdf)\nC. Procedure Planning in Instructional Videos (https://arxiv.org/abs/1907.01172)\nD. Procedure Planning in Instructional Videos via Contextual Modeling and Model-based Policy Learning (https://arxiv.org/abs/2110.01770)\nE. proScript: Partially Ordered Scripts Generation via Pre-trained Language Models (https://arxiv.org/abs/2104.08251)\nF. Neural Task Graphs: Generalizing to Unseen Tasks from a Single Video Demonstration (https://arxiv.org/abs/1807.03480)\n\n3. The CSCG structure described in section 3 (preliminaries) seems hard to scale. It is basically describing a POMDP setup with nodes as states and edges with transitions corresponding to taking actions. But in most practical problems, we have infinite states so this graph is going to be intractable to maintain, except in small grid worlds.\n\n3. Section 3.1 mentions matrices, so are the state space and observation state considered countably finite? Are the #nodes in graph G known a priori? What about the number of clones per observation?\n\n4. Section 3.1: Do we know the emission model E a priori? If not, how can we keep it fixed throughout learning? Seems like a restrictive and unrealistic assumption to me for most applications!\n\n6. The paper only presents experiments on navigation-style environments in 2D/3D grid-world settings. This implies countably finite states and observations and only 3-4 actions. Hence, the CSCGs are small and tractable. In the case of the 3D T-maze experiment, it is unclear if the vector quantizer used for the 3D T-maze experiments would generalize to more complex 3D environments with complex objects in it.\n\n7. I might have missed the first usage of EPN, but Section 4.1 suddenly brings it up with no prior citation or description. Please clarify if I missed its description in the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "kuaLY7CIZp7", "reviewer": "Reviewer_rzAg", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes Graph Schemas, a particular type of action-conditioned HMM with deterministic emissions. The uncertainty comes from the transition matrix, as well as observation \"clones\", i.e. multiple nodes that output the same observation. These schemas can then be transferred by keeping T fixed while changing E, enabling fast learning and planning in different simulated environments.", "summary_of_the_review": "The paper reads well and is well-motivated from a CogSci perspective. However, the realism of its final assumptions is not supported, neither by writing or formulation, nor by experiments. To improve the writing I would devote more time on the method section and less on the experiment section. For an ICLR-level bar, I believe there needs to be some non-toy experiment or a very strong motivation. Furthermore, the differences with the HMM literature need to be discussed in much more detail to justify the novelty of this approach.", "strength_and_weaknesses": "- Strengths\n    - I think the paper is clear both in its motivation, what it sets to achieve, and the future remaining directions.\n    - Within simulated discrete environments, I believe the types of environments was pretty diverse.\n- Weaknesses\n    - The overall setting feels quite contrived, the discreteness of the states and actions, the exact clones, etc, are motivated from a cognitive perspective, but the realism of those assumptions isn't well-motivated.\n    - Deriving from the previous point, the experiments were quite contrived as well, no realistic experiments were performed, which doesn't allow me to see a way of bringing these ideas to eventual useful systems.\n    - The model isn't explained that much and the intellectual novelty compared to standard HMM literature is not explained.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "RhbtLjM9z3", "reviewer": "Reviewer_XbR6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies ‘clone structured cognitive graphs’ (CSCG), a particular type of action-conditioned Hidden Markov Model (HMM) with specific assumptions on the emission function (deterministic emissions and partial observability, where multiple hidden states can generate the same observation). They investigate the ability of this model to transfer between tasks, primarily when the task topology remains the same but the emission functions change and need to be relearned. Experiments on a range of tasks show that the CSCG model is indeed capable of fast adaptation in transfer experiments. ", "summary_of_the_review": "This paper studies a relevant topic (fast adaptation in navigation tasks), has a lot of experiments, and shows good performance. However, there is one big issue, as mentioned above, which is the use of Hidden Markov Models. Previous literature has extensively shown how HMMs can be specified, estimated, and successfully used in smaller tasks. The authors indeed confirm these ideas, by showing that HMM structure can be efficiently reused between tasks, for example transferring the latent transition structure while re-estimating the emission function. \n\nHowever, the main issue that HMMs ran into was their inability to scale to more complex, high-dimensional tasks. The authors only touch upon this topic in their final experiment, where they use clustering on pixel observations to make the HMM still applicable in a basic T-Maze. I do not say we need a neural network in every ICLR paper, but I do think this is an issue the paper should address much more. Would it be possible to use the HMM architecture with a learned embedding function? Can you also jointly optimize high-capacity embedding function and latent dynamics, and still transfer one of them afterwards? There is much more work in this direction, see for example [1,2,3,4], and EPN itself of course. I think you need to address these issues, because when it purely comes to the properties of the HMM approach, then I do not think the current paper presents enough novelty. \n\n[1] Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances in Neural Information Processing Systems, pages 3604–3613, 2017.\n[2] Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. In Advances in neural information processing systems, pages 2946–2954, 2016.\n[3] Rahul G Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In AAAI, pages 2101–2109, 2017.\n[4] Rangapuram, Syama Sundar, et al. \"Deep state space models for time series forecasting.\" Advances in neural information processing systems 31 (2018).\n", "strength_and_weaknesses": "Strength:\n* Fast adaptation and transfer in navigation tasks is a relevant and important research topic. \n* The paper has extensive experiments. \n* The method shows quick adaptation in new settings. \n\nWeaknesses: \n* The approach is essentially based on a Hidden Markov Model, where the emission function gets relearned in a new task. This is useful and works well, but Hidden Markov Models and all their variants have of course been extensively studied in previous decades. A major challenge for HMMs is the inability to scale to higher-dimensional problems, and the authors do not really address these concerns (see discussion in conclusion of review).  \n* It is unclear to me how much prior information these methods need. For example, how is the number of z variables chosen per task? I suppose it usually matches the true number of nodes in the source task, or the number of clusters in the 3D task? This seems to be a strong form of prior knowledge, i.e., the designer already knowns how many nodes to specify.  \n* After the number of z variables has been specified, they can all possibly be connected in the T matrix right? In other words, the learned topologies that you visualize after learning eliminate all elements of T matrix that get estimated to zero? Or do you also put certain restrictions in the T matrix before training? \n* Your methods use hardcoded exploration strategies. You indicate exploration is not part of the current method, which is fine, but it will influence performance. It is unclear to me what exploration strategies the baseline (EPN) uses, and therefore hard to compare performance. In MPG (Fig 2) for example, to what extend is the performance of EPN influenced by a suboptimal exploration strategy? I think these are things you should address. \n* In general, you have a very extensive experiments section, which is good, but your methodology section is rather short.\n\n* In StreetLearn (Fig 2), it seems you use pixel based observation in the figure, but I do not read anything about it in the results? Do you use an emission function back to pixel level? I guess you would then need the clustering approach of your last experiment again? Otherwise the task looks much more complicated in the figure than it really is. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject