Differentiable Clone-Structured Causal Graphs for End-to-End Cognitive Map Learning from Image Sequences

# Differentiable Clone-Structured Causal Graphs for End-to-End Cognitive Map Learning from Image Sequences
Source: [https://arxiv.org/html/2607.12382](https://arxiv.org/html/2607.12382)
Differentiable Clone\-Structured Causal Graphs for End\-to\-End Cognitive Map Learning from Image Sequences

Arash Nikzad1, Sasan Sarbishegi2, Ali Dasmeh3, Muhammad Asif4, Parsa Gharavi1, Erik Husom5, Sagar Sen5, Andrew B\. Lehr6,7,11, Olivier Penacchio8,9, Ana Clemente4, Tristan M\. Stöber7,10,11,\*

1Goethe University Frankfurt, Frankfurt, Germany;2Independent Researcher, Teheran, Iran;3Max Planck Institute for Human Development, Berlin, Germany;4Department of Cognitive Neuropsychology, Max Planck Institute for Empirical Aesthetics, Frankfurt, Germany;5SINTEF, Oslo, Norway;6Department of Neuro\- and Sensory Physiology, University Medical Center Göttingen, Göttingen, Germany;7Institute of Computer Science and Campus Institute Data Science, University Göttingen, Germany;8Bridging Research in AI and Neuroscience \(brAIN\), Computer Vision Center, Bellaterra, Spain;9Computer Science Department, Universitat Autònoma de Barcelona, Bellaterra, Spain;10Epilepsy Center Frankfurt Rhine\-Main, Department of Neurology, Goethe University Frankfurt, Frankfurt, Germany;11Circulant Labs, Bensheim, Germany\. \*Corresponding author : tristan\.stoeber@posteo\.net

Abstract\.How can an agent build a structured map of its world from nothing but an ongoing sequence of raw sensory input and its own movements, especially when natural variation means exact sensory patterns rarely repeat? The Clone\-Structured Causal Graph algorithm \(CSCG\), a normative hippocampus model, shows how an interpretable map can be learned from aliased observations\. However, CSCG requires a predefined discrete alphabet, and its expectation\-maximization formulation is not easily combined with existing neural network modules, preventing the end\-to\-end processing of raw image sequences\. We remove this barrier by reformulating CSCG as a single, fully differentiable module, gradCSCG, and coupling it to a learned vector\-quantized variational autoencoder \(VQ\-VAE\) perceptual front\-end\. A soft emission forward pass allows the map\-learning objective to flow back into perception, while a set of loss\-balancing mechanisms mitigates module collapse during joint training\. We demonstrate, first, that gradient training reproduces CSCG’s results on original symbolic grid worlds by recovering room topology from heavily aliased observations\. Second, we show that map recovery remains robust on MNIST image sequences, where each visit to a location yields a newly sampled image of its assigned digit\. Across four heavily aliased environments, the end\-to\-end pipeline successfully uncovers the underlying adjacency graph with high edge precision and recall, directly from visual input\. This work provides a proof of principle that CSCG can serve as a composable building block in a deep learning architecture\.

Keywords:cognitive maps, clone\-structured cognitive graph \(CSCG\), differentiable sequence models, vector\-quantized representation learning, hippocampus, NeuroAI, topology recovery\.

## 1Introduction

Well\-structured internal representations allow biological and artificial agents to pick shortcuts on routes never taken and to provide guidance in situations where trial\-and\-error learning would be fatal\. Thus, understanding and reengineering the emergence of such representations is a fundamental research frontier both in neuroscience \[1–5\] and artificial intelligence \(AI\) \[6–8\]\.

TheClone\-Structured Causal Graph\(CSCG\) algorithm \[9, 10\], a normative hippocampus model, explains how well\-structured representations can emerge from experience\. Technically, CSCG is an overcomplete hidden Markov model \(HMM\) with a fixed emission matrix that creates a statistical model from sequential observations\. Compressing a series of observation–action pairs into a higher\-order representation of its environment, CSCG learns to disambiguate contexts from aliased observations\. Obeying the Markov property — i\.e\. any subsequent state depends only on the current state — CSCG is forced to represent a novel context with a new clone among its hidden nodes\. The graph emerging from this cloning operation provides a condensed map of the environment and is suitable for planning, consolidation and abstraction\. However, while elegantly creating well\-structured representations in static and relatively small environments, it is unclear how to extend this approach to richer, perceptually complex observations\. The bottleneck is that CSCG is trained by Expectation–Maximization over a fixed, discrete observation alphabet, which prevents seamless composition with gradient\-trained neural modules\.

We resolve this by reimplementing CSCG as a single differentiable, gradient\-trained computation \[11\] in TensorFlow — a model we callgradCSCG\. This innovation allows us to co\-train gradCSCG with a vector\-quantized variational autoencoder \(VQ\-VAE\)\. We demonstrate that this approach preserves CSCG’s expressivity while enabling it to handle sensory variability and complexity in environments composed of MNIST digits\.

Contributions

1. 1\.Anend\-to\-end trainable pipeline composed of a gradient\-based CSGG and a VQ\-VAEable to create a topological map from sequences of images\.
2. 2\.Loss\-balancing for stable joint training— length normalization, weight annealing, a diversity penalty, and anti\-collapse safeguards \(Section 3\.6\)\.
3. 3\.A formal, reusabletopology\-recovery evaluation suite\(Section 3\.10\)\.
4. 4\.An empirical study on four MNIST grid\-world environments with strong aliasing \(Sections 4–5\)\.

### 1\.1Related Work

CSCG toolkit\.A differentiable cloned\-HMM forward pass with a soft\-observation interface and encoder\-gradient flow has also been developed in concurrent, independent open\-source work \[21\], whose image experiment couples a convolutional network with a gradient\-based CSCG for supervised digit classification\. Relative to this work, our specific contributions are the learned VQ\-VAE discretizer trained jointly with the sequence model, the loss\-balancing that keeps that joint training stable, and — crucially — the use of this gradient\-based setup toactually recover and evaluate an environment’s map: we deduce the physical adjacency graph from the learned transitions and score it against the ground\-truth topology \(Section 5\)\. Their image experiment, by contrast, uses a convolutional–CSCG coupling for supervised digit classification only\. They do not attempt to recover the topology, which is the central question our pipeline is built to answer\.

Neural sequence models for cognitive maps\.A complementary line learns maps from the latent codes of a neural predictor rather than from a cloned HMM\. Dedieu et al\. \[22\] train a transformer with discrete bottlenecks on next\-observation prediction and read interpretable cognitive maps off its bottleneck indices for planning in partially observed environments\. As in our pipeline, a neural representation is discretized and a map is recovered from observation–action streams; the difference iswhere the map lives\. In their approach the cognitive map is decoded from the transformer’s bottleneck codes as a separate, post\-hoc analysis and handed to an external solver for planning, so the network’s own working representation stays dense and is not itself an interpretable, directly usable map\. In ours the clone\-graph sequence model is co\-trained with the perceptual discretizer and its transition matrixisthe map: an intrinsically interpretable, on\-line structure that the model computes with and can be queried directly for planning and for understanding the environment\.

Dynamic and expanding maps\.Closer to the open problems we raise in Section 7, de Tinguy et al\. \[23\] grow a cognitive map online — dynamically expanding it over predicted poses within an active\-inference agent — and benchmark against CSCG on grid environments\. Their focus on a map thatexpandsas the agent explores is precisely the dynamic\-allocation capability that our fixed clone budget lacks\. However, their model learns via discrete counting operations that are not natively differentiable, making it difficult to couple seamlessly with deep neural networks for visual processing\. By contrast, our fully differentiable formulation allows the map to be co\-trained with a neural front\-end, learning the discretization end\-to\-end directly from raw pixels\.

## 2Methods

### 2\.1Problem formulation

An agent produces an episode of lengthTT: observationsx1:Tx\_\{1:T\}withxt∈𝒳⊂ℝH×W×Cx\_\{t\}\\in\\mathcal\{X\}\\subset\\mathbb\{R\}^\{H\\times W\\times C\}and actionsa1:T−1a\_\{1:T\-1\}withat∈𝒜=\{1,…,A\}a\_\{t\}\\in\\mathcal\{A\}=\\\{1,\\dots,A\\\}, whereata\_\{t\}is taken between timesttandt\+1t\{\+\}1\. Each observation is emitted at an underlying physical placegt∈𝒢g\_\{t\}\\in\\mathcal\{G\}; the environment has a true undirected adjacency graphℳ=\(𝒢,ℰ\)\\mathcal\{M\}=\(\\mathcal\{G\},\\mathcal\{E\}\)\. The placesgtg\_\{t\}and edgesℰ\\mathcal\{E\}are usedonly for evaluationand are never seen during training\. The goal is to learn, from\(x1:T,a1:T−1\)\(x\_\{1:T\},a\_\{1:T\-1\}\)alone, a latent model whose transition structure recoversℳ\\mathcal\{M\}\.

The pipeline \(Figure 1\) has two modules: a VQ\-VAE that maps each image to a discrete token, and an action\-conditioned cloned HMM over those tokens whose latent graph is the learned map\.

Imagextx\_\{t\}EncoderEϕE\_\{\\phi\}Codebook\{ek\}\\\{e\_\{k\}\\\}DecoderDψD\_\{\\psi\}ℒrec\\mathcal\{L\}\_\{\\mathrm\{rec\}\}Soft posteriorρt\\rho\_\{t\}overKKcodesgradCSCGforward \(soft\)ℒgradCSCG\\mathcal\{L\}\_\{\\mathrm\{gradCSCG\}\}ztz\_\{t\}q~t\\tilde\{q\}\_\{t\}log⁡ρt\\log\\rho\_\{t\}Figure 1:The gradCSCG pipeline\.Solid arrows: forward computation\. Dashed arrows: gradient flow\. The encoder feeds both a reconstruction branch \(hard quantizationq~t\\tilde\{q\}\_\{t\}\+ decoder\) and a sequence branch \(soft codebook posteriorρt\\rho\_\{t\}\+ differentiable cloned\-HMM forward pass\)\. Because the gradCSCG likelihood is differentiable inρt\\rho\_\{t\}, the topological objectiveℒgradCSCG\\mathcal\{L\}\_\{\\mathrm\{gradCSCG\}\}shapes the encoder\. The codebook itself is updated by an Exponential Moving Average \(EMA\), not by gradients\.
### 2\.2Perceptual front\-end: VQ\-VAE

Encoder and quantization\.A convolutional encoderEϕ:𝒳→ℝDE\_\{\\phi\}:\\mathcal\{X\}\\to\\mathbb\{R\}^\{D\}maps each image to a latentzt=Eϕ​\(xt\)z\_\{t\}=E\_\{\\phi\}\(x\_\{t\}\)\. A codebook\{ek\}k=1K\\\{e\_\{k\}\\\}\_\{k=1\}^\{K\},ek∈ℝDe\_\{k\}\\in\\mathbb\{R\}^\{D\}, whereKKdenotes the number of discrete latent codes andDDis the dimensionality of each codebook embedding, defines a discrete token by nearest\-neighbour assignment,

kt=argmink∈\{1,…,K\}∥zt−ek∥22,qt=ekt,k\_\{t\}=\\arg\\min\_\{k\\in\\\{1,\\dots,K\\\}\}\\;\\lVert z\_\{t\}\-e\_\{k\}\\rVert\_\{2\}^\{2\},\\qquad q\_\{t\}=e\_\{k\_\{t\}\},where the squared distance is computed as∥z−ek∥22=∥z∥22−2​⟨z,ek⟩\+∥ek∥22\\lVert z\-e\_\{k\}\\rVert\_\{2\}^\{2\}=\\lVert z\\rVert\_\{2\}^\{2\}\-2\\langle z,e\_\{k\}\\rangle\+\\lVert e\_\{k\}\\rVert\_\{2\}^\{2\}\. Gradients cross the non\-differentiablearg⁡min\\arg\\minby the straight\-through estimator \[18\],

q~t=zt\+sg​\(qt−zt\),\\tilde\{q\}\_\{t\}=z\_\{t\}\+\\mathrm\{sg\}\(q\_\{t\}\-z\_\{t\}\),wheresg​\(⋅\)\\mathrm\{sg\}\(\\cdot\)is the stop\-gradient operator, so the forward value isqtq\_\{t\}while∂q~t/∂zt=I\\partial\\tilde\{q\}\_\{t\}/\\partial z\_\{t\}=I\. A decoderDψD\_\{\\psi\}reconstructsx^t=Dψ​\(q~t\)\\hat\{x\}\_\{t\}=D\_\{\\psi\}\(\\tilde\{q\}\_\{t\}\)\.

Losses\.Over a minibatchℬ\\mathcal\{B\},

ℒrec=1\|ℬ\|​∑t∥xt−x^t∥22,\\mathcal\{L\}\_\{\\mathrm\{rec\}\}=\\tfrac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\_\{t\}\\lVert x\_\{t\}\-\\hat\{x\}\_\{t\}\\rVert\_\{2\}^\{2\},ℒcommit=1\|ℬ\|​∑t∥sg​\(qt\)−zt∥22,\\mathcal\{L\}\_\{\\mathrm\{commit\}\}=\\tfrac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\_\{t\}\\lVert\\mathrm\{sg\}\(q\_\{t\}\)\-z\_\{t\}\\rVert\_\{2\}^\{2\},the second term pulling encoder outputs toward their assigned code\.

EMA codebook updates\.The codebook isnottrained by gradient descent\. With decayγ∈\(0,1\)\\gamma\\in\(0,1\), per batch we accumulate, for every codekk, a cluster sizenkn\_\{k\}and a vector summkm\_\{k\},

nk\\displaystyle n\_\{k\}←γ​nk\+\(1−γ\)​∑t𝟏​\[kt=k\],\\displaystyle\\leftarrow\\gamma\\,n\_\{k\}\+\(1\-\\gamma\)\\sum\_\{t\}\\mathbf\{1\}\[k\_\{t\}=k\],mk\\displaystyle m\_\{k\}←γ​mk\+\(1−γ\)​∑t𝟏​\[kt=k\]​zt,\\displaystyle\\leftarrow\\gamma\\,m\_\{k\}\+\(1\-\\gamma\)\\sum\_\{t\}\\mathbf\{1\}\[k\_\{t\}=k\]\\,z\_\{t\},and setek←mk/n^ke\_\{k\}\\leftarrow m\_\{k\}/\\hat\{n\}\_\{k\}with Laplace\-smoothed size

n^k=nk\+ϵ\(∑k′nk′\)\+K​ϵ​\(∑k′nk′\)\.\\hat\{n\}\_\{k\}=\\frac\{n\_\{k\}\+\\epsilon\}\{\\left\(\\sum\_\{k^\{\\prime\}\}n\_\{k^\{\\prime\}\}\\right\)\+K\\epsilon\}\\left\(\\sum\_\{k^\{\\prime\}\}n\_\{k^\{\\prime\}\}\\right\)\.
Soft codebook posterior\.For the differentiable coupling \(Section 3\.5\) the encoder also emits a temperature\-controlled posterior over the codebook,

log⁡ρt​\(k\)=log⁡softmaxk​\(−∥zt−ek∥22/τ\),\\log\\rho\_\{t\}\(k\)=\\log\\mathrm\{softmax\}\_\{k\}\\left\(\-\\lVert z\_\{t\}\-e\_\{k\}\\rVert\_\{2\}^\{2\}/\\tau\\right\),which is differentiable inztz\_\{t\}\. Asτ→0\\tau\\to 0,ρt\\rho\_\{t\}concentrates onktk\_\{t\}and this recovers the hard assignment\.

### 2\.3Sequence model: action\-conditioned cloned HMM

State space and clone structure\.Each tokenkkis assignedCk≥1C\_\{k\}\\geq 1clones— latent states that all emit tokenkkbut participate in different transition contexts\. With a trailingsinkstate⊥\\bot, the state space is𝒮=\{1,…,N\}\\mathcal\{S\}=\\\{1,\\dots,N\\\}withN=1\+∑k=1KCkN=1\+\\sum\_\{k=1\}^\{K\}C\_\{k\}\. The sink state is a special terminal state that does not correspond to any visual token; it is used to absorb probability mass at the end of a sequence and to make sequence termination explicit in the HMM formulation\. A fixed mapω:𝒮∖\{⊥\}→\{1,…,K\}\\omega:\\mathcal\{S\}\\setminus\\\{\\bot\\\}\\to\\\{1,\\dots,K\\\}gives the token each state emits; in the uniform caseCk≡CC\_\{k\}\\equiv Candω​\(s\)=⌈s/C⌉\\omega\(s\)=\\lceil s/C\\rceil\. Emissions aredeterministic:

Bs,o=𝟏​\[ω​\(s\)=o\],log⁡Bs,o=\{0,ω​\(s\)=o−∞,otherwise,B\_\{s,o\}=\\mathbf\{1\}\[\\omega\(s\)=o\],\\qquad\\log B\_\{s,o\}=\\begin\{cases\}0,&\\omega\(s\)=o\\\\ \-\\infty,&\\text\{otherwise,\}\\end\{cases\}and the sink emits no real token\. Clones are exactly the mechanism that disambiguates aliasing: one token observed at two places is explained by two clones with distinct transition rows\.

Parameters\.The model has initial\-state logitsπ∈ℝN\\pi\\in\\mathbb\{R\}^\{N\}and action\-conditioned transition logitsΘ∈ℝA×N×N\\Theta\\in\\mathbb\{R\}^\{A\\times N\\times N\}, yielding

π¯=softmax​\(π\),Ta,i,j=softmaxj​\(Θa,i,⋅\)\.\\bar\{\\pi\}=\\mathrm\{softmax\}\(\\pi\),\\qquad T\_\{a,i,j\}=\\mathrm\{softmax\}\_\{j\}\\left\(\\Theta\_\{a,i,\\cdot\}\\right\)\.All learning resides in\(π,Θ\)\(\\pi,\\Theta\); emissions are fixed\.

Forward likelihood\.Writinglogsumexpiui=log​∑ieui\\mathop\{\\mathrm\{logsumexp\}\}\_\{i\}u\_\{i\}=\\log\\sum\_\{i\}e^\{u\_\{i\}\}, the log\-forward messages for an episode\(o1:T,a1:T−1\)\(o\_\{1:T\},a\_\{1:T\-1\}\)obey

log⁡α1​\(j\)=\\displaystyle\\log\\alpha\_\{1\}\(j\)=\{\}log⁡π¯j\+log⁡Bj,o1,\\displaystyle\\log\\bar\{\\pi\}\_\{j\}\+\\log B\_\{j,o\_\{1\}\},log⁡αt\+1​\(j\)=\\displaystyle\\log\\alpha\_\{t\+1\}\(j\)=\{\}logsumexpi\[log⁡αt​\(i\)\+log⁡Tat,i,j\]\+log⁡Bj,ot\+1,\\displaystyle\\mathop\{\\mathrm\{logsumexp\}\}\_\{i\}\\left\[\\log\\alpha\_\{t\}\(i\)\+\\log T\_\{a\_\{t\},i,j\}\\right\]\+\\log B\_\{j,o\_\{t\+1\}\},and the episode log\-likelihood isℓ​\(o1:T∣a1:T−1\)=logsumexpjlog⁡αT​\(j\)\\ell\(o\_\{1:T\}\\mid a\_\{1:T\-1\}\)=\\mathop\{\\mathrm\{logsumexp\}\}\_\{j\}\\log\\alpha\_\{T\}\(j\)\. The training loss is the mean negative log\-likelihood \(NLL\)

ℒgradCSCG=−1\|ℬ\|​∑\(o,a\)∈ℬℓ​\(o1:T∣a1:T−1\)\.\\mathcal\{L\}\_\{\\mathrm\{gradCSCG\}\}=\-\\tfrac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\_\{\(o,a\)\\in\\mathcal\{B\}\}\\ell\(o\_\{1:T\}\\mid a\_\{1:T\-1\}\)\.

### 2\.4Gradient\-based training of the cloned HMM

Unlike the classical EM training of cloned HMMs, we evaluate the forward recursion and NLL loss as a single differentiable, log\-space computational graph \(a masked time recursion over padded minibatches\) and optimize\(π,Θ\)\(\\pi,\\Theta\)directly by stochastic gradient descent with Adam \[19\]\. Log\-space arithmetic withlogsumexp\\mathop\{\\mathrm\{logsumexp\}\}keeps the recursion numerically stable over long episodes\. This gradient formulation is what makes the sequence model composable with a neural front\-end\.

### 2\.5Differentiable soft\-emission coupling

To let the topological objective shape perception, we replace the hard emission termlog⁡Bj,ot\\log B\_\{j,o\_\{t\}\}in the forward pass by the soft log\-posterior of the token that statejjemits:

log⁡α1s​\(j\)=\\displaystyle\\log\\alpha^\{s\}\_\{1\}\(j\)=\{\}log⁡π¯j\+log⁡ρ1​\(ω​\(j\)\),\\displaystyle\\log\\bar\{\\pi\}\_\{j\}\+\\log\\rho\_\{1\}\\left\(\\omega\(j\)\\right\),log⁡αt\+1s​\(j\)=\\displaystyle\\log\\alpha^\{s\}\_\{t\+1\}\(j\)=\{\}logsumexpi\[log⁡αts​\(i\)\+log⁡Tat,i,j\]\+log⁡ρt\+1​\(ω​\(j\)\),\\displaystyle\\mathop\{\\mathrm\{logsumexp\}\}\_\{i\}\\left\[\\log\\alpha^\{s\}\_\{t\}\(i\)\+\\log T\_\{a\_\{t\},i,j\}\\right\]\+\\log\\rho\_\{t\+1\}\\left\(\\omega\(j\)\\right\),with the sink assignedlog⁡ρt​\(⊥\)=−∞\\log\\rho\_\{t\}\(\\bot\)=\-\\infty\. The resulting log\-likelihoodℓs\\ell^\{s\}is differentiable inρ1:T\\rho\_\{1:T\}and hence, through the soft posterior, in the encoder parametersϕ\\phi\. Two properties hold by construction:

1. 1\.Consistency\.Ifρt\\rho\_\{t\}is the one\-hot distribution on the observed tokenoto\_\{t\}, thenlog⁡ρt​\(ω​\(j\)\)=log⁡Bj,ot\\log\\rho\_\{t\}\(\\omega\(j\)\)=\\log B\_\{j,o\_\{t\}\}and the soft recursion reduces exactly to the hard forward pass\. Asτ→0\\tau\\to 0the soft pipeline thus recovers the hard pipeline\.
2. 2\.End\-to\-end differentiability\.Gradients ofℒgradCSCG\\mathcal\{L\}\_\{\\mathrm\{gradCSCG\}\}propagateℓs→ρ→z→ϕ\\ell^\{s\}\\\!\\to\\\!\\rho\\\!\\to\\\!z\\\!\\to\\\!\\phi, so the encoder is trained, in part, to produce tokens that make the action\-conditioned sequenceexplainable\.

### 2\.6Joint objective and loss balancing

We train the model in three phases\. First, the VQ\-VAE is trained independently to initialize the visual encoder and decoder and obtain stable discrete representations\. Second, the complete model, comprising both the VQ\-VAE and gradCSCG, is trained jointly in an end\-to\-end manner with stochastic gradient descent\. Third, the VQ\-VAE is frozen and gradCSCG is fine\-tuned using hard emissions, allowing the temporal model to refine its transition structure while operating on discrete visual assignments\.

Combined objective\.At joint steptt,

ℒjoint=ℒrec\+β​ℒcommit\+λt​ℒ~gradCSCG\+αdiv​ℒdiv\.\\mathcal\{L\}\_\{\\mathrm\{joint\}\}=\\mathcal\{L\}\_\{\\mathrm\{rec\}\}\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{commit\}\}\+\\lambda\_\{t\}\\,\\widetilde\{\\mathcal\{L\}\}\_\{\\mathrm\{gradCSCG\}\}\+\\alpha\_\{\\mathrm\{div\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{div\}\}\.
Length normalization\.The raw NLL grows asO​\(T\)O\(T\), which on long episodes dwarfs theO​\(1\)O\(1\)reconstruction term and collapses the codebook\. We therefore use the per\-step NLL

ℒ~gradCSCG=−1\|ℬ\|​∑1T​ℓs​\(o1:T∣a1:T−1\)\.\\widetilde\{\\mathcal\{L\}\}\_\{\\mathrm\{gradCSCG\}\}=\-\\tfrac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\\tfrac\{1\}\{T\}\\,\\ell^\{s\}\(o\_\{1:T\}\\mid a\_\{1:T\-1\}\)\.
Weight annealing\.The sequence\-loss weight is ramped linearly so the codebook stabilizes under reconstruction before topological pressure turns on:

λt=λ⋅min⁡\(1,t/Tanneal\)\.\\lambda\_\{t\}=\\lambda\\cdot\\min\\left\(1,\\;t/T\_\{\\mathrm\{anneal\}\}\\right\)\.
Diversity penalty\.Letρ¯​\(k\)=1\|ℬ\|​∑tρt​\(k\)\\bar\{\\rho\}\(k\)=\\frac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\_\{t\}\\rho\_\{t\}\(k\)be mean codebook usage andH​\(ρ¯\)=−∑kρ¯​\(k\)​log⁡ρ¯​\(k\)H\(\\bar\{\\rho\}\)=\-\\sum\_\{k\}\\bar\{\\rho\}\(k\)\\log\\bar\{\\rho\}\(k\)its entropy\. The penalty

ℒdiv=log⁡K−H​\(ρ¯\)≥0\\mathcal\{L\}\_\{\\mathrm\{div\}\}=\\log K\-H\(\\bar\{\\rho\}\)\\;\\geq\\;0vanishes only at uniform usage and counteracts codebook collapse\.

Anti\-collapse safeguards and finalization\.During Phase 2 we monitor codebook perplexity and keep the highest\-perplexity checkpoint; optionalλ\\lambda\-throttling, rollback, and dead\-code revival provide further protection\. A shortfinalizationphase then freezes the encoder and refines\(π,Θ\)\(\\pi,\\Theta\)on hard tokens with the unnormalized loss, optionally with a transition\-entropy regularizer−η​∑a,iH​\(Ta,i,⋅\)\-\\eta\\sum\_\{a,i\}H\(T\_\{a,i,\\cdot\}\)to sharpen transition rows\. The following algorithm summarizes one joint step\.

Algorithm 1 — One joint training step Require:image chunkx1:Tx\_\{1:T\}, actionsa1:T−1a\_\{1:T\-1\}, weightsβ,λt,αdiv\\beta,\\lambda\_\{t\},\\alpha\_\{\\mathrm\{div\}\}, temperatureτ\\tau1\.zt←Eϕ​\(xt\)z\_\{t\}\\leftarrow E\_\{\\phi\}\(x\_\{t\}\)\(encode\)2\.qt,q~t←quantize​\(zt\)q\_\{t\},\\tilde\{q\}\_\{t\}\\leftarrow\\mathrm\{quantize\}\(z\_\{t\}\); update codebook by EMA3\.x^t←Dψ​\(q~t\)\\hat\{x\}\_\{t\}\\leftarrow D\_\{\\psi\}\(\\tilde\{q\}\_\{t\}\);ℒrec,ℒcommit←\\mathcal\{L\}\_\{\\mathrm\{rec\}\},\\mathcal\{L\}\_\{\\mathrm\{commit\}\}\\leftarrowreconstruction and commitment losses \(Section 3\.2\)4\.log⁡ρt←log⁡softmax​\(−∥zt−e⋅∥2/τ\)\\log\\rho\_\{t\}\\leftarrow\\log\\mathrm\{softmax\}\(\-\\lVert z\_\{t\}\-e\_\{\\cdot\}\\rVert^\{2\}/\\tau\)5\.ℓs←\\ell^\{s\}\\leftarrowsoft forward pass \(Section 3\.5\)6\.ℒ~gradCSCG←−ℓs/T\\widetilde\{\\mathcal\{L\}\}\_\{\\mathrm\{gradCSCG\}\}\\leftarrow\-\\ell^\{s\}/T;ℒdiv←log⁡K−H​\(ρ¯\)\\mathcal\{L\}\_\{\\mathrm\{div\}\}\\leftarrow\\log K\-H\(\\bar\{\\rho\}\)7\.ℒjoint←\\mathcal\{L\}\_\{\\mathrm\{joint\}\}\\leftarrowcombined objective \(Section 3\.6\)8\.update\(ϕ,ψ,π,Θ\)\(\\phi,\\psi,\\pi,\\Theta\)with Adam on∇ℒjoint\\nabla\\mathcal\{L\}\_\{\\mathrm\{joint\}\}

### 2\.7Decoding

The maximum\-a\-posteriori state path is obtained by the Viterbi recursion \[15\]:

δ1​\(j\)\\displaystyle\\delta\_\{1\}\(j\)=log⁡π¯j\+log⁡Bj,o1,\\displaystyle=\\log\\bar\{\\pi\}\_\{j\}\+\\log B\_\{j,o\_\{1\}\},δt\+1​\(j\)\\displaystyle\\delta\_\{t\+1\}\(j\)=maxi⁡\[δt​\(i\)\+log⁡Tat,i,j\]\+log⁡Bj,ot\+1,\\displaystyle=\\max\_\{i\}\\left\[\\delta\_\{t\}\(i\)\+\\log T\_\{a\_\{t\},i,j\}\\right\]\+\\log B\_\{j,o\_\{t\+1\}\},with backpointersψt\+1​\(j\)=arg⁡maxi⁡\[δt​\(i\)\+log⁡Tat,i,j\]\\psi\_\{t\+1\}\(j\)=\\arg\\max\_\{i\}\\left\[\\delta\_\{t\}\(i\)\+\\log T\_\{a\_\{t\},i,j\}\\right\]and tracebacksT⋆=arg⁡maxj⁡δT​\(j\)s^\{\\star\}\_\{T\}=\\arg\\max\_\{j\}\\delta\_\{T\}\(j\),st⋆=ψt\+1​\(st\+1⋆\)s^\{\\star\}\_\{t\}=\\psi\_\{t\+1\}\(s^\{\\star\}\_\{t\+1\}\), giving the most likely hidden state at each timett\. The decoded paths1:T⋆s^\{\\star\}\_\{1:T\}is the basis of all evaluation\.

### 2\.8Token compaction and clone allocation

In Phase 1, codes never emitted by the trained encoder are pruned and the alphabet is relabelled to the active set\. Clone countsCkC\_\{k\}may be uniform or set per observationby hand\(a fixed, manually chosen per\-observation clone budget\); the deterministic emission structure of Section 3\.3 is unchanged\. This is a static hyperparameter chosen before training,nota learned, on\-demand allocation of clones\.

### 2\.9Evaluation metrics

All metrics are computed from the decoded paths1:T⋆s^\{\\star\}\_\{1:T\}of a held\-out episode together with the ground\-truth placesg1:Tg\_\{1:T\}and edgesℰ\\mathcal\{E\}\(used only here, never in training\)\. Let𝒱\\mathcal\{V\}be the set of visited states andn​\(s,g\)=♯​\{t:st⋆=s,gt=g\}n\(s,g\)=\\sharp\\\{t:s^\{\\star\}\_\{t\}=s,\\;g\_\{t\}=g\\\}\.

State\-to\-place assignment\.Each visited state is mapped to its majority place,χ​\(s\)=arg⁡maxg⁡n​\(s,g\)\\chi\(s\)=\\arg\\max\_\{g\}n\(s,g\)\.

Clone purity\.

Purity=1\|𝒱\|​∑s∈𝒱maxg⁡n​\(s,g\)∑gn​\(s,g\),\\mathrm\{Purity\}=\\frac\{1\}\{\|\\mathcal\{V\}\|\}\\sum\_\{s\\in\\mathcal\{V\}\}\\frac\{\\max\_\{g\}n\(s,g\)\}\{\\sum\_\{g\}n\(s,g\)\},with the visit\-weighted variant∑smaxg⁡n​\(s,g\)/T\\sum\_\{s\}\\max\_\{g\}n\(s,g\)\\,/\\,T\. High purity means clone states correspond cleanly to single physical places\.

Projected map and edge F1\.Latent transitions are projected onto a place graph by

W​\(g,g′\)=maxa⁡maxi,j∈𝒱χ​\(i\)=g,χ​\(j\)=g′⁡Ta,i,j,W\(g,g^\{\\prime\}\)=\\max\_\{a\}\\;\\max\_\{\\begin\{subarray\}\{c\}i,j\\in\\mathcal\{V\}\\\\ \\chi\(i\)=g,\\;\\chi\(j\)=g^\{\\prime\}\\end\{subarray\}\}T\_\{a,i,j\},and thresholded into a learned edge setℰ^​\(η\)=\{\(g,g′\):g≠g′,W​\(g,g′\)\>η\}\\widehat\{\\mathcal\{E\}\}\(\\eta\)=\\\{\(g,g^\{\\prime\}\):g\\neq g^\{\\prime\},\\,W\(g,g^\{\\prime\}\)\>\\eta\\\}\. Withtp=\|ℰ^∩ℰ\|\\mathrm\{tp\}=\\lvert\\widehat\{\\mathcal\{E\}\}\\cap\\mathcal\{E\}\\rvert, precision, recall and F1 are defined in the usual way\. We report F1 over a threshold sweepη∈\{0\.01,0\.05,0\.1,0\.2,0\.3\}\\eta\\in\\\{0\.01,0\.05,0\.1,0\.2,0\.3\\\}\.

Viterbi\-path map and edge F1\.The projected map above thresholdseverylearned transition and therefore admits weak false edges\. As a sharper read\-out we keep only the transitions the decoded path actually takes: froms1:T⋆s^\{\\star\}\_\{1:T\}we tally consecutive state transitions, project each to a place edge\(χ​\(st⋆\),χ​\(st\+1⋆\)\)\(\\chi\(s^\{\\star\}\_\{t\}\),\\chi\(s^\{\\star\}\_\{t\+1\}\)\)withχ​\(st⋆\)≠χ​\(st\+1⋆\)\\chi\(s^\{\\star\}\_\{t\}\)\\neq\\chi\(s^\{\\star\}\_\{t\+1\}\), and retain an edge once it is traversed in more than a small fraction of the episode \(we use0\.2%0\.2\\%of theTTsteps\)\. Precision, recall and F1 againstℰ\\mathcal\{E\}are then computed as above\. This Viterbi\-path read\-out is the “Map F1” reported in Table 3 and the Viterbi\-path columns of Table 4 \(Section 5\.3\)\.

Action\-next\-cell accuracy\.For every represented placeggand actionaa, the predicted next placearg⁡maxg′⁡Wa​\(g,g′\)\\arg\\max\_\{g^\{\\prime\}\}W\_\{a\}\(g,g^\{\\prime\}\)is compared with the environment’s true successornext​\(g,a\)\\mathrm\{next\}\(g,a\)\.

Token–place entropies\.From the empirical token/place co\-occurrence we reportH​\(token∣place\)H\(\\text\{token\}\\mid\\text\{place\}\)andH​\(place∣token\)H\(\\text\{place\}\\mid\\text\{token\}\); the former is small when perception is consistent, the latter reflects the \(irreducible\) aliasing of the environment\.

### 2\.10Implementation and hyperparameters

The encoder is composed of three convolutional layers with strides2,2,12,2,1, followed by global average pooling and a dense projection toℝD\\mathbb\{R\}^\{D\}; the decoder uses the corresponding mirrored architecture\. The gradCSCG forward pass, soft forward pass, and training steps are implemented as compiled TensorFlow graphs\. Viterbi decoding is performed outside the compiled graph during inference\. Transition logits are initialized with a bias toward the sink state so probability mass is well\-defined before training\. Table 1 lists all hyperparameters\.

Table 1:Hyperparameters\. Ranges span the four environments of Section 4\.
### 2\.11Environments

Our experiments take place in a family of controlled navigation environments\. Each environment is a set of discrete locations —places— connected into a graph that is the hidden ground\-truth map; in this work the graph is a 2\-D grid, so each place is a cell linked to its immediate neighbours\. An agent explores by arandom walk: at every step it occupies one place, receives anobservationproduced by that place, and takes one of four movement actions \(up, down, left, right\) that moves it to an adjacent place, with walls and boundaries sticky \(an invalid move leaves it where it is\)\. Everything the learner ever sees is this stream of alternating observations and actions; the agent’s true location and the graph’s adjacency are never exposed, and are kept only for evaluation\. The task is to reconstruct the map — which places border which — from the stream alone\.

This is exactly the setting in which the original CSCG was validated, with one simplification on the perception side: there, each place emits asingle discrete symbolfrom a small alphabet, so the observation vocabulary is given in advance \[9\]\. We keep everything else — the known ground\-truth topology, the strong perceptual aliasing, and the hidden position — but replace the symbol at each place with araw image, so the vocabulary is no longer given and must be learned from pixels\. This lets us ask, on the very maps George et al\. used, whether a cognitive map can still be recovered once perception is itself part of the learning problem\.

Formally, writevt∈𝒱v\_\{t\}\\in\\mathcal\{V\}for the place visited at timett,at∈\{up,down,left,right\}a\_\{t\}\\in\\\{\\text\{up\},\\text\{down\},\\text\{left\},\\text\{right\}\\\}for the action taken, andxtx\_\{t\}for the observation received; an episode is the resulting stream

\(x1,a1\),\(x2,a2\),…,\(xT,aT\)\.\(x\_\{1\},a\_\{1\}\),\(x\_\{2\},a\_\{2\}\),\\dots,\(x\_\{T\},a\_\{T\}\)\.Each cellvvcarries a digit labeld​\(v\)∈\{0,…,9\}d\(v\)\\in\\\{0,\\dots,9\\\}, and we use two observation models:

- •Symbolic\(the original\-CSCG control \[9\]\): the observation is the digit token itself,xt=d​\(vt\)x\_\{t\}=d\(v\_\{t\}\)\.
- •Image\(our benchmark\): the observation is a freshly drawn MNIST image of that digit,xt∼𝒟d​\(vt\)x\_\{t\}\\sim\\mathcal\{D\}\_\{d\(v\_\{t\}\)\}, where𝒟d\\mathcal\{D\}\_\{d\}is the empirical MNIST distribution for digitdd\.

Under the image model, repeated visits to one cell never produce identical pixels \(non\-stationary appearance\), while distinct cells that share a digit produce visually similar observations \(aliasing\)\. The model must therefore turn high\-dimensional, variable images intostablediscrete tokens before action context can resolve the remaining aliasing — a strictly harder problem than the symbolic one, and a more realistic test of whether temporal\-structure learning can be coupled to learned perception\.

The four environments \(Table 2; Figure 2\) are MNIST analogues of the canonical demonstrations of the original CSCG \[9\], each isolating a different facet of map\-learning under aliasing:

- •aliased— a4×44\\times 4room with four digit classes arranged so that each recurs four times\. With only four distinct observations across sixteen places, appearance alone is almost uninformative; this mirrors George et al\.’s room with four unique observations \(Fig\. 2a,b of \[9\]\)\.
- •corridors— a5×55\\times 5layout whose interior walls carve narrow corridors joined by repeated digits, a walled\-maze variant that stresses recovery through bottlenecks\.
- •room— a6×66\\times 6room in which a ring of distinct border digits surrounds a4×44\\times 4interior of asinglerepeated digit, so sixteen interior cells look identical\. This mirrors George et al\.’s uniform\-interior room \(Fig\. 2c,d of \[9\]\); the large aliased core is the hardest case for a fixed clone budget, and we use it to illustrate that the per\-observation clone budget can simply be set by hand \(Section 5\.3\)\.
- •two\_rooms— a13×913\\times 9map of two offset rooms that share a3×33\\times 3patch, so one local appearance occurs at two globally distinct places \(aconfounder\)\. This mirrors the two overlapping rooms George et al\. use to probe transitive inference \(Fig\. 2e,f of \[9\]\); it is the largest and most aliased benchmark\.

Table 2:Benchmark environments\.KKis the number of discrete observation tokens \(codebook size\); “Places” counts walkable cells; the last column names the corresponding experiment in the original CSCG paper \[9\]\.![Refer to caption](https://arxiv.org/html/2607.12382v1/figures/environments.png)Figure 2:The four benchmark environments, each shown as its grid of per\-cell digit classes \(walls in grey\)\. Cells that share a digit are perceptually aliased; in the image benchmark every visit to a cell returns a different MNIST sample of its digit\. Note the uniform interior ofroomand the shared corner that links the two halves oftwo\_rooms\.For each environment we collect action\-conditioned uniform random walks, 4\-10 episodes of 10,000 steps, and never expose the agent’s position or the adjacency graph to the model; both are retained only to score a held\-out episode with the metrics of Section 3\.9\. Each run executes the three\-stage pipeline of Section 3\.6 \(VQ\-VAE warm\-up→\\rightarrowjoint training→\\rightarrowfinalization\)\. During warm\-up an optional, benchmark\-only digit\-classification loss may be applied to the encoder; since object classes are unknown in a general environment it is disableable, and the perceptual front\-end learns to discretize the observations without it\.

## 3Results

We evaluate gradCSCG on the four grid\-worlds of Section 4\. Training consumes only the observation–action stream; the agent’s position and the true adjacency graph are withheld and used solely to score a held\-out episode \(Section 3\.9\)\. The argument runs in three steps: gradient training reproduces the cognitive maps of the original, EM\-trained CSCG \(Section 5\.1\); that recovery survives the replacement of the given symbol by a raw, never\-repeating image \(Section 5\.2\); and the gradient\-trained formulation affords further properties of practical value \(Section 5\.3\)\. Table 3 summarises the headline metrics; the subsections establish each claim in turn\.

Table 3:Main results across the four environments \(image observations;roomuses a hand\-set per\-observation clone budget\)\. State\-to\-place purity is visit\-weighted; map recall and F1 are read from the Viterbi\-path graph \(Section 5\.3, Table 4\), with recall 1\.00 throughout\. Entries are means over repeated runs; the seed\-to\-seed s\.d\. is≤0\.01\\leq 0\.01for every environment excepttwo\_rooms\(≤0\.03\\leq 0\.03\)\.### 3\.1Gradient descent recovers the cognitive maps of the original CSCG

The learned latent graph contains the true map\.Trained only on observation–action walks, gradCSCG concentrates its transition probability on the edges that mirror the environment’s physical adjacency, and map\-edgerecall is 1\.00in every environment\. This is the central result of the original CSCG \[9\], obtained here by back\-propagation rather than expectation–maximization\. We first reproduce it in theoriginal symbolic settingof \[9\], in which each cell emits its integer observation directly with no images: trained by gradient descent, gradCSCG recovers the same maps George et al\. obtain by EM — a square room laid out as a 2\-D grid, a rectangular room whose uniform interior is distinguished only by its walls, and three disjoint rooms stitched into a single coherent graph \(Figure 3, cf\. Fig\. 2 of \[9\]\)\. The recovered transition graph reproduces the room topology, and the node colours \(one per observation\) show that perceptually identical cells are split into separate, correctly\-placed clones\. This isolates the change of training rule from the change of observation model; the same recovery then carries over to raw images \(Section 5\.2\)\.

![Refer to caption](https://arxiv.org/html/2607.12382v1/figures/section51_integer_cscg.png)Figure 3:gradCSCG reproduces the original CSCG on integer \(symbolic\) grid\-worlds, trained by gradient descent rather than EM\.Top row:the ground\-truth map of each environment, every cell \(or place\) coloured by its observation\.Bottom row:the transition graph recovered by gradCSCG, with nodes coloured by the observation they emit and positioned by the learned connectivity — colours are matched cell↔\\leftrightarrownode, so the correspondence is direct\.Left:a square room recovered as a 2\-D grid from a heavily aliased symbol stream \(cf\. George et al\. \[9\], Fig\. 2a,b\)\.Middle:a rectangular room whose uniform interior — the large block of a single repeated observation — is disambiguated only by the surrounding walls \(cf\. Fig\. 2c,d\)\.Right:three pentagonal rooms with identical local observations, stitched into one coherent graph by transitive inference \(cf\. Fig\. 2e,f\)\. In every case the recovered graph matches the ground\-truth topology\.Clones absorb the aliasing into place\-specific states\.Recovering topology from aliased observations requires splitting one observation into context\-specific latent states — the clone mechanism\. Decoding a held\-out episode and assigning each visited clone to its majority place, clones are almost perfectly place\-specific: visit\-weighted state\-to\-place purity is 0\.95–1\.00, every place is represented by at least one clone, and a typical place is covered by only 1\.5–1\.7 clones \(Table 3; the clone\-in\-layout panels of Figure 5\)\. Gradient training therefore reproduces the context\-split, place\-cell\-like representations that make the original CSCG interpretable\.

The recovered map answers “where does this action lead?”\.A map is useful only if it can be queried for the outcome of an action\. Reading the most likely successor place for each \(place, action\) pair off the learned transitions and comparing it with the true successor, action\-outcome accuracy is 0\.97–1\.00 \(Table 3\) — the graph is consistent enough to support the one\-step, inference\-based planning of \[9\]\.

### 3\.2The maps survive raw, non\-stationary, aliased images

A learned front\-end supplies CSCG\-grade tokens from pixels\.Where the original CSCG is handed clean symbols, our model must manufacture them\. The VQ\-VAE does so almost deterministically: a place emits a single token on 1\.0–1\.1 of its visits on average, the conditional entropyH​\(token∣place\)H\(\\text\{token\}\\mid\\text\{place\}\)is only 0\.09–0\.18 nats, and codebook perplexity tracks the number of digit classes — even though that place is never seen twice\. Reconstructions stay faithful throughout joint training \(Figure 4\), confirming that the topological objective sharpens rather than collapses the codebook\.

Topology recovery holds end\-to\-end and at scale\.Driving gradCSCG with theselearnedtokens — not the ground\-truth digits — the full pixel\-to\-map pipeline still recovers every adjacency \(recall 1\.00\) up to the largest, most aliased environment,two\_rooms\(13×913\\times 9, 87 places, 151 latent states\), at state\-to\-place purity 0\.98 and map F1 1\.00 \(Figure 5; Table 3\)\. The map that George et al\. recover from symbols, we recover from images\.

![Refer to caption](https://arxiv.org/html/2607.12382v1/figures/reconstructions.png)Figure 4:Perception is stable under joint training \(two\_rooms\)\. Top: input MNIST observations; bottom: VQ\-VAE reconstructions, each annotated with the discrete token assigned to it\. Every visit to a place yields a different image, yet the assigned token is consistent\.![Refer to caption](https://arxiv.org/html/2607.12382v1/figures/section52_mnist_rooms.png)Figure 5:gradCSCG on raw MNIST observations, one column per environment \(aliased,corridors,room,two\_rooms\)\.Top row:ground\-truth layout \(cells coloured by digit class\)\.Middle row:the decoded clone states drawn at their assigned physical cells \(transitions omitted for clarity\) — perceptually identical cells \(same colour, repeated across the grid\) are split into separate clones that sit at the correct places\.Bottom row:the recovered transition graph, nodes drawn as the MNIST digit they emit\. The topology is recovered from images in every environment, up to the13×913\\times 9two\_rooms\.
### 3\.3Additional properties of the gradient\-trained model

A minor convenience: per\-observation clone budgets\.The clone count is a hyperparameter that can be set per observation rather than uniformly\. Giving the heavily\-aliased interior digit ofroommore clones \(20\) than its rarely\-repeated border digits \(4\) recovers the same map with 57 latent states instead of 201, at no cost to recall\. This is astatic, hand\-setallocation fixed before training — not a learned, on\-demand assignment of clones — so we note it only as a practical knob, not a contribution\.

Reading the map off the Viterbi path recovers it almost exactly\.Projectingeveryabove\-threshold latent transition into the place graph floods it with weak false edges, so the naive precision is low \(0\.16–0\.44 atη=0\.01\\eta=0\.01; Table 4, left\)\. But those weak edges are decoding noise, not routes the agent ever takes\. We therefore keep only the place edges that the Viterbi\-decoded path actually traverses more than a small fraction of the episode —0\.2% of theTTsteps\(20 transitions atT=10,000T=10\{,\}000\)\. This is aminimum\-trafficcutoff that asks for genuine, repeated movement rather than a single mis\-step, and it scales with episode length instead of being a fixed magic number\. It removes essentially all false edges while leaving recall untouched: precision rises to0\.95–1\.00and map F1 to0\.98–1\.00across all four environments, at recall 1\.00 \(Table 4\)\. The read\-out is robust to the exact fraction: because true adjacencies are traversed hundreds of times per episode and spurious ones only a handful, any cutoff from∼\\sim0\.1% to∼\\sim0\.5% ofTT\(10–50 transitions here\) leaves recall at 1\.00 while precision saturates by∼\\sim0\.2%\.

Table 4:Map\-edge scores from the projected transition graph \(every edge with weight\>0\.01\>0\.01\) versus the Viterbi\-path graph \(edges traversed in more than 0\.2% of theTTsteps, i\.e\.\>20\>20atT=10,000T=10\{,\}000\)\. Recall is 1\.00 in both columns; the Viterbi\-path read\-out removes the weak false edges that depress the projected\-graph precision\.Joint training is stable\.Through the joint phase the reconstruction loss stays flat while the length\-normalized gradCSCG term falls once its weightλt\\lambda\_\{t\}ramps in; the finalization phase then sharply reduces the gradCSCG NLL with the encoder frozen \(Figure 6\) — the behaviour the loss\-balancing terms of Section 3\.6 are designed to produce\.

![Refer to caption](https://arxiv.org/html/2607.12382v1/figures/loss_curves.png)Figure 6:Training dynamics \(two\_rooms\)\. VQ\-VAE warm\-up, the joint phase \(reconstruction, commitment, length\-normalized gradCSCG, and diversity terms\), and the pure\-gradCSCG finalization phase with the encoder frozen\.

## 4Discussion

Our central result is aproof of principle: the Clone\-Structured Cognitive Graph \(CSCG\), a normative account of how the hippocampus builds cognitive maps \[9, 10\], can be reformulated to live inside a gradient\-trained pipeline\. Re\-deriving its forward algorithm as a differentiable, log\-space computation lets the model be optimized by backpropagation rather than expectation–maximization\. On the original symbolic grid\-worlds the gradient\-trained model reproduces the hallmark behaviour of George et al\. — recovering room topology from heavily aliased observations and splitting identical\-looking places into context\-specific clones \(Section 5\.1\)\. Crucially, the reformulation also makes CSCGcomposable: co\-trained with a VQ\-VAE perceptual front\-end, it recovers cognitive maps of pixel\-based environments — where every place looks different on each visit and many places look alike — that the symbolic model cannot even ingest \(Sections 5\.2–5\.3\)\. The contribution is therefore less a new map\-learning result than a change ofsubstrate: CSCG becomes a module that can be wired into, and trained jointly with, neural networks\.

This change of substrate is worth making because the two research traditions that bear on map learning each pay a price the other avoids\. Interpretable cognitive\-map models — CSCG \[9, 10\], the Tolman–Eichenbaum Machine \[14\], the successor/predictive map \[3\] — recover sparse, relational structure that supports planning and generalization, but they presuppose a discrete,givenobservation alphabet and so cannot, by themselves, look at pixels\. Neural world models — the Dreamer family of recurrent state\-space models \[6\] and the vector\-quantized representation learners they build on \[12\] — consume raw perception readily, but their latent dynamics are dense and hard to read as a map\. gradCSCG sits in the gap between them: it keeps CSCG’s sparse, aliasing\-resolving clone structure while delegating perception to a learned front\-end\. The enabling fact is not itself new — the HMM forward recursion is a differentiable computation \[16\], and neural emissions for HMMs have been studied in language modelling \[17\] — but applying it to theaction\-augmentedcloned HMM, and letting the topological objective propagate back to shape perception, is what turns an ordinary encoder into asequence\-awareone\.

That coupling is not free, and the way it fails is instructive\. Trained naively, the joint objective collapses the codebook, because the sequence likelihood grows with episode length and swamps reconstruction; the loss\-balancing of Section 3\.6 — length normalization, weight annealing, and a diversity penalty — is what makes the two modules cooperate rather than compete, and in our experiments it is the difference between learning a map and learning a single token\. We regard this as the main practical lesson for any future CSCG\-plus\-network hybrid: the perceptual and sequence objectives must be explicitly balanced, or the stronger one consumes the other\.

## 5Limitations and Outlook

Limitations\.Two caveats bound the present results directly\.\(i\) Read\-out threshold\.High precision requires reading the map off the Viterbi\-decoded path with a minimum\-traffic cutoff \(Section 5\.3\); the cutoff is adaptive \(a fixed fraction of episode length\) and recall is insensitive to it, but a fully parameter\-free read\-out — e\.g\. stronger transition\-entropy regularization or edge calibration — is still open\.\(ii\) Scope\.The environments are controlled, static, 2\-D MNIST grid\-worlds; we have not yet tested 3\-D, partially observed, non\-stationary, or visually richer settings\.

Outlook\.The deeper limitation is conceptual, and it points to where this work goes next\. By construction, CSCG captures thestatic, relationalskeleton of an environment — which places exist and how they connect — buttrue behavioural flexibility needs more than a static map\. A realistic agent must also track thedynamicelements of a scene \(objects that move, appear, or change\), handle far more complex and higher\-dimensional observations, and adaptquicklyas the world changes — none of which a clone graph over fixed places is designed to do\. We therefore see gradCSCG not as a complete agent but as thestructuralcomponent of a larger system, with genuine flexibility coming from pairing it with complementary modules trained in the same differentiable framework\. The most natural partner is arecurrent network \(RNN\): where the clone graph holds the slow, stable map, an RNN can flexibly carry the fast\-changing, dynamic context the map omits \[6\] — a division of labour that mirrors how real neural circuits combine stable spatial codes with rapidly updating population activity\. Richer input — 3\-D or egocentric vision — would in turn call for a heavier perceptual front\-end that pre\-processes raw views before they are discretized\. Making CSCG gradient\-trained is precisely what makes such hybrids buildable, since every part can then be optimized by the same backpropagation\.

Many questions remain open\. The most immediate isdynamic clone recruitment: here the clone budget is a static hyperparameter fixed before training \(Section 5\.3\), whereas the cloning principle ultimately calls forrecruiting a new clone whenever a novel context appears\. A differentiable mechanism that grows or prunes clones on demand would let the model size itself to an environment’s true complexity — and learn it faster and more incrementally — a natural next step now that training is gradient\-based\. Threshold\-free map read\-out, scaling to larger and partially observed worlds, and coupling the learned map to a planning or reinforcement\-learning loop \[6\] are further directions\.

## 6Conclusion

We have shown that the Clone\-Structured Cognitive Graph can be trained by gradient descent while faithfully reproducing the behaviour of its EM\-trained original, and that, so trained, it can be grounded in a learned VQ\-VAE front\-end to recover cognitive maps of pixel\-based environments\. We close on what this isfor\. The hippocampus — of which CSCG is a normative model \[9, 10\] — does not work in isolation; it operates in close interaction with the entorhinal cortex, prefrontal cortex, and many other regions \[14\], each carrying out a different kind of computation, and the cognitive map is useful precisely because it is embedded in that larger system\. We read our contribution in the same spirit: a gradient\-trained CSCG is not a monolithic solution but acomposablemodule, and rephrasing it in the language of deep learning is what lets it be wired together with the complementary machinery it needs — neural perception, recurrent tracking of dynamics, and planning\.

Because our formulation is differentiable, the benefit runs in both directions: the same module that provides explicit representations to artificial agents, may also become a more scalable normative model for studying how the brain builds and uses such representations\. Unlocking the power of world models, in brains and in machines alike, is therefore less about the module on its own than about a clever division of labour among specialized parts\.

Code and data availability\.The implementation, benchmark environments, training scripts and evaluation suite are available in the project repository\.

## Code and Data Availability

The complete implementation of thegradCSCGpipeline, VQ\-VAE discretization modules, synthetic training environments, and evaluation suite is publicly available on GitHub at[https://github\.com/tristanstoeber/gradCSCG](https://github.com/tristanstoeber/gradCSCG)\. The MNIST datasets used in this work are open\-source and automatically fetched by the provided training scripts\.

## Acknowledgements

AN, PG acknowledge support by the Goethe Research Exchange Program\. TMS acknowledges funding from the “Advancing World model Learning with Neural Cloned\-Structured Causal Graphs” project, funded by the European Union, via the oc1\-2024\-TES\-01 issued and implemented by the ENFIELD project, under the grant agreement No 101120657; and from the CIDAS Fellowship 2025, University of Göttingen\.

## References

- \[1\]E\. C\. Tolman\. Cognitive maps in rats and men\.Psychological Review, 55\(4\):189–208, 1948\.
- \[2\]J\. O’Keefe and L\. Nadel\.The Hippocampus as a Cognitive Map\. Clarendon Press, 1978\.
- \[3\]K\. L\. Stachenfeld, M\. M\. Botvinick, and S\. J\. Gershman\. The hippocampus as a predictive map\.Nature Neuroscience, 20:1643–1653, 2017\.
- \[4\]J\. C\. R\. Whittington, D\. McCaffary, J\. J\. W\. Bakermans, and T\. E\. J\. Behrens\. How to build a cognitive map: insights from models of the hippocampal formation\.Nature Neuroscience, 25:1257–1272, 2022\.
- \[5\]D\. Ha and J\. Schmidhuber\. World models\.arXiv:1803\.10122, 2018\.
- \[6\]D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap\. Mastering diverse control tasks through world models\.Nature, 640\(8059\):647–653, 2025\.
- \[7\]J\. Pasukonis, T\. Lillicrap, and D\. Hafner\. Evaluating long\-term memory in 3D mazes\.arXiv:2210\.13383, 2022\.
- \[8\]M\. R\. Samsami, A\. Zholus, J\. Rajendran, and S\. Chandar\. Mastering memory tasks with world models\. InInternational Conference on Learning Representations \(ICLR\), 2024\.
- \[9\]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\.Nature Communications, 12:2392, 2021\.
- \[10\]R\. V\. Raju, J\. S\. Guntupalli, G\. Zhou, M\. Lázaro\-Gredilla, and D\. George\. Space is a latent sequence: a theory of the hippocampus\.Science Advances, 10:eadm8470, 2024\.
- \[11\]W\. Sun, J\. Winnubst, M\. Natrajan, et al\. Learning produces an orthogonalized state machine in the hippocampus\.Nature, 640:165–175, 2025\.
- \[12\]A\. van den Oord, O\. Vinyals, and K\. Kavukcuoglu\. Neural discrete representation learning\. InAdvances in Neural Information Processing Systems \(NeurIPS\), 2017\.
- \[13\]A\. Razavi, A\. van den Oord, and O\. Vinyals\. Generating diverse high\-fidelity images with VQ\-VAE\-2\. InAdvances in Neural Information Processing Systems \(NeurIPS\), 2019\.
- \[14\]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, 2020\.
- \[15\]L\. R\. Rabiner\. A tutorial on hidden Markov models and selected applications in speech recognition\.Proceedings of the IEEE, 77\(2\):257–286, 1989\.
- \[16\]J\. Eisner\. Inside\-outside and forward\-backward algorithms are just backprop\. InProceedings of the Workshop on Structured Prediction for NLP, pages 1–17, 2016\.
- \[17\]K\. M\. Tran, Y\. Bisk, A\. Vaswani, D\. Marcu, and K\. Knight\. Unsupervised neural hidden Markov models\. InProceedings of the Workshop on Structured Prediction for NLP, pages 63–71, 2016\.
- \[18\]Y\. Bengio, N\. Léonard, and A\. Courville\. Estimating or propagating gradients through stochastic neurons for conditional computation\.arXiv:1308\.3432, 2013\.
- \[19\]D\. P\. Kingma and J\. Ba\. Adam: a method for stochastic optimization\. InInternational Conference on Learning Representations \(ICLR\), 2015\.
- \[20\]Y\. LeCun, L\. Bottou, Y\. Bengio, and P\. Haffner\. Gradient\-based learning applied to document recognition\.Proceedings of the IEEE, 86\(11\):2278–2324, 1998\.
- \[21\]R\. Young\. cscg\_toolkit: JAX/PyTorch and Julia implementations of Clone\-Structured Cognitive Graphs\. GitHub repository,[https://github\.com/SynapticSage/cscg\_toolkit](https://github.com/SynapticSage/cscg_toolkit), 2025\.
- \[22\]A\. Dedieu, W\. Lehrach, G\. Zhou, D\. George, and M\. Lázaro\-Gredilla\. Learning cognitive maps from transformer representations for efficient planning in partially observed environments\.arXiv:2401\.05946, 2024\.
- \[23\]D\. de Tinguy, T\. Verbelen, and B\. Dhoedt\. Learning dynamic cognitive map with autonomous navigation\.arXiv:2411\.08447, 2024\.