BraiNCA - brain-inspired neural cellular automata and applications to morphogenesis and motor control

year: 2026
paper: brainca-brain-inspired-neural-cellular-automata-and-applications-to-morphogenesis-and-motor-control
website:
code: https://github.com/LPioL/BraiNCA | https://github.com/MaxWolf-01/brainca
connections: NCA, Benedikt Hartl, levin

---

Notation

${\mathbf{c}}_i^t\in {\mathbb{R}}^C$ … cell state of cell $i$ at time $t$
${\rho }_i^t$ … optional conditioning input (positional encoding, prev action, …)
${\mathbf{s}}_i^t=[{\mathbf{c}}_i^t;{\rho }_i^t]\in {\mathbb{R}}^S$ … extended state
${\mathcal{N}}_i$ … local neighborhood indices
${\mathcal{L}}_i$ … long-range neighborhood indices
$a_{ij}^t\in {\mathbb{R}}^A$ … raw local attention score
$b_{ij}^t\in {\mathbb{R}}^A$ … raw long-range attention score
${\alpha }_{ij}^t$ … normalized local attention weight
${\beta }_{ij}^t$ … normalized long-range attention weight
$f_{\text{attn}}^{(\mathcal{N})}$ … local attention MLP, ${\mathbb{R}}^{2S}\to {\mathbb{R}}^A$
$f_{\text{attn}}^{(\mathcal{L})}$ … long-range attention MLP, ${\mathbb{R}}^{2S}\to {\mathbb{R}}^A$
${\mathbf{h}}_k^t\in {\mathbb{R}}^C$ … GRU hidden state (values aggregated by attention); in the code this is the cell state ${\mathbf{c}}_k^t$
${\mathbf{n}}_i^t$ … aggregated local signal
${\mathbf{l}}_i^t$ … aggregated long-range signal
${\mathbf{z}}_i^t$ … interaction vector
$f_Z$ … composition operator for $\mathbf{z}$
${\mathbf{m}}_i^t\in {\mathbb{R}}^M$ … message MLP output, GRU input
$f_{\text{msg}}$ … message MLP
${\mathbf{u}}_t$ … system-level feedback (e.g. one-hot prev action)
${\mathbf{o}}_t$ … external observation
$f_{\text{GRU}}$ … GRU cell
${\mathbf{r}}_i^t\in {\mathbb{R}}^C$ … refinement residual
$f_{\text{refine}}$ … refinement MLP
${\mathbf{\ell }}_i^t$ … per-cell action logits (lunar lander)
$w_i^t$ … per-cell fire probability
${\mathcal{R}}_r$ … cells in action region $r$
$L_r^t$ … region-level action logit
$\pi (a_t=r\mid {\mathbf{o}}_t)$ … action policy
$N$ … number of cells
$C$ … cell state dim
$S$ … extended state dim
$A$ … attention score dim
$M$ … message dim
$T$ … final timestep

Limitations / Ideas / Obvious next steps

• GRU, not QKV-attention-memory
• With graph attn there’s also no any2any between the msgs per node (not sure if it matters - prlly not if u factor efficiency and diffusion over time)
• Hidden state is public / the same as the public message, not compressed / …
• Long-range connectivity is task-specific/engineered and fixed (Zipf’s law… tho it’s more complicated than that and different per task and… you get the point, it could be much simpler, bitter lesson pilled)
• All cells get (the same) observations…
• All cells are active every step
• Arrangement affects performance and learning speed - let the cells arrange themselves
• No growth and death; weights are the genome, evolved slowly over many generations of growth, growth happening during rollouts (i.e. genome only meta learns). This prlly (definitely) needs some pre-training to have any chance of working. Perhaps… pre-pre-training the NCA on CAs…

Unclarities

• Why does the morpho experiment use additive fusion $[\mathbf{c};\mathbf{n}+\mathbf{l}]$ for the interaction vector while lunar lander concatenates $[\mathbf{c};\mathbf{n};\mathbf{l}]$? Not motivated in the paper.
• why dont they do all to all comms in the attn agg? efficiency? unnecessary? (same info is spread out through timesteps) - maybe full attn only worth it for internal states
• Lunar lander readout: paper prose says ”$f_{\text{refine}}$ output is split into two heads” but the equations write both heads as reading from ${\mathbf{h}}_i^t$ directly. $f_{\text{act}}$: $M\to M\to 2$, but ${\mathbf{h}}_i^t\in {\mathbb{R}}^C$, $C\ne M$ suggests the prose is right and the equations are wrong. And it also makes more sense that way…

BraiNCA

Init:
${\mathbf{c}}_i^0\sim \mathcal{N}(0,0.1\cdot I)$ (morpho) or ${\mathbf{c}}_i^0=0$ (lunar)
Obs ${\mathbf{o}}_t$ and optional system-feedback ${\mathbf{u}}_{t-1}$ broadcast to all cells (last action taken by lander, one-hot; empty in morpho).
Connections
Repeat for $T$ steps (morpho: $T=35$ total; lunar: $n_{\text{steps}}=3$ per env step), $i\in \{1,\dots ,N\}$ cells ($N=256$):

$$\begin{array}{rl}{\mathbf{s}}_i^t& =[{\mathbf{c}}_i^t;{\rho }_i^t]\\ a_{ij}^t& =f_{\text{attn}}^{(\mathcal{N})}([{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]),{\alpha }_{ij}^t={\text{softmax}}_j(a_{ij}^t),{\mathbf{n}}_i^t={\sum }_{k\in {\mathcal{N}}_i}{\alpha }_{ik}^t{\mathbf{h}}_k^t\\ b_{ij}^t& =f_{\text{attn}}^{(\mathcal{L})}([{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]),{\beta }_{ij}^t={\text{softmax}}_j(b_{ij}^t),{\mathbf{l}}_i^t={\sum }_{k^{\prime }\in {\mathcal{L}}_i}{\beta }_{i{k}^{\prime }}^t{\mathbf{h}}_{k^{\prime }}^t\\ {\mathbf{z}}_i^t& =f_Z({\mathbf{s}}_i^t,{\mathbf{n}}_i^t,{\mathbf{l}}_i^t)\\ {\mathbf{m}}_i^t& =f_{\text{msg}}({\mathbf{z}}_i^t;{\mathbf{u}}_t,{\mathbf{o}}_t)\\ {\mathbf{h}}_i^t& =f_{\text{GRU}}({\mathbf{m}}_i^t,{\mathbf{h}}_i^{t-1})\\ {\tilde{\mathbf{h}}}_i^t& ={\mathbf{h}}_i^t+f_{\text{refine}}({\mathbf{h}}_i^t)\\ {\mathbf{c}}_i^{t+1}& =\{\begin{array}{ll}{\tilde{\mathbf{h}}}_i^t& \text{(morphogenesis)}\\ W_C{\tilde{\mathbf{h}}}_i^t& \text{(lunar lander)}\end{array}\\ {\mathbf{\ell }}_i^t& =f_{\text{act}}({\tilde{\mathbf{h}}}_i^t)\text{(lunar lander only)}\end{array}$$

Lunar lander

• cells are placed in regions, quadrants or somatotopic (T-shaped, like the body of a lander, left right thrusts up top, main thruster down bottom, noop action regions bottom left and right). The noop region is excluded from LR connections, which is not motivated.
• softmax over logits $({\ell }_{\text{noop}},{\ell }_{\text{fire}})$ → fire probability $w_i$
• majority vote, confidence weighted
• Within each region, compute a weighted average of fire logits , weighted by the fire probability/confidence $w_i$:

$$L_r=\frac{{\sum }_{i\in \mathcal{R}r}{w}_i\cdot {\ell }_{i,\text{fire}}}{{\sum }_{i\in {\mathcal{R}}_r}{w}_i}=L_{\text{NOOP}},L_{\text{LEFT}},L_{\text{MAIN}},L_{\text{RIGHT}}$$

• softmax over $L_r$ to get action dist to sample from
  

Detailed

Extended state ${\mathbf{s}}_i^t=[{\mathbf{c}}_i^t;{\rho }_i^t]$ ${\rho }_i^t$ bundles all non-cell-state inputs. What it contains is task-dependent: Morphogenesis: $\rho =$ nothing $\mathbf{s}=\mathbf{c}$, $S=C=9$ … 3 visible channels (cell type logits) + 6 hidden channels. (Code uses $C=16$: 3 visible + 13 hidden.) Lunar lander: ${\rho }_i^t=[{\ell }_i^{t-1};{\rho }_i]$ where ${\ell }_i^{t-1}\in \{0,1\}$ is the previous fire/noop decision, and ${\rho }_i\in {\mathbb{R}}^2$ is normalized grid position ¹ (static). $S=15$ … extended state. $C+1+2=12+1+2$. The 1 is the previous fire/noop decision, the 2 is the positional encoding. $C=12$ … latent dim (no "visible" channels, the readout is via a separate action head).

Neighbor scoring ( graph attention) $f_{\text{attn}}$: Linear($2S$, 64) → GELU → Linear(64, $A$), applied to the concatenated pair $[{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]$, with $A=1$ (a scalar score per neighbour). Each cell scores each of its neighbours independently, and local and long-range use separate $f_{\text{attn}}$. In the official implementation, the local scorer uses GELU but the long-range one uses ReLU, even though the paper says GELU for both.

\begin{align*}
a_{ij}^t &= f_{\text{attn}}^{(\mathcal{N})}([\mathbf{s}i^t;\mathbf{s}j^t]), \quad \alpha{ij}^t = \text{softmax}j(a{ij}^t), \quad \mathbf{n}i^t = \sum{k\in\mathcal{N}i}\alpha{ik}^t,\mathbf{h}k^t \
b{ij}^t &= f{\text{attn}}^{(\mathcal{L})}([\mathbf{s}i^t;\mathbf{s}j^t]), \quad \beta{ij}^t = \text{softmax}j(b{ij}^t), \quad \mathbf{l}i^t = \sum{k’\in\mathcal{L}i}\beta{ik’}^t,\mathbf{h}{k’}^t
\end{align*}

Scoring uses the extended states $\mathbf{s}$; the aggregated values are the hiddens $\mathbf{h}_k$, i.e. the cell states $\mathbf{c}_k$ (see GRU note). Grid edges: the original impl. zero-pads, so a boundary cell's missing neighbours enter the softmax as zero vectors instead of being dropped. That gives every cell a distance-from-edge signal, which for morphogenesis is its only positional information (besides the init … which is supposed to be random… according to the paper, contrary to the code, but non-random would be a silly memorization task so all my final experiments use random inits). I first tried masking the edges instead but morphogenesis convergence was much more fragile (lunar has positional encodings).

"Interaction vector" (cell input)

\mathbf{z}{i}^{t}=f{Z}(\mathbf{s}{i}^{t},\mathbf{n}{i}^{t},\mathbf{l}_{i}^{t})

Morphogenesis vanilla: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t]$
Morphogenesis long-range: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t+{\mathbf{l}}_i^t]$
Lunar lander long-range: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t;{\mathbf{l}}_i^t]$

"Message MLP"

$${\mathbf{m}}_i^t=f_{\text{msg}}({\mathbf{z}}_i^t;{\mathbf{u}}_t,{\mathbf{o}}_t)$$

In morphogenesis, ${\mathbf{u}}_t$ and ${\mathbf{o}}_t$ are empty so it’s just $f_{\text{msg}}({\mathbf{z}}_i^t)$. The code makes it a stack of 1×1 convs, $2C\to 64\to 32\to C$, so $\mathbf{m}$ comes out $C$-dim and the GRU input matches its hidden width. The paper never states these dims.
In lunar lander, they’re simply concatted to ${\mathbf{z}}_i^t$, all cells receive the same obs.

GRU update + readout ${\mathbf{h}}_i^{t-1}$ as its recurrent input, so it reads like $\mathbf{h}$ is a second state that persists on its own, next to $\mathbf{c}$. The code keeps no separate $\mathbf{h}$: it feeds ${\mathbf{c}}_i^t$ in as the hidden; wherever the paper writes ${\mathbf{h}}^{t-1}$ the implementation uses ${\mathbf{c}}^t$. That includes the ${\mathbf{h}}_k$ that attention aggregates, which are just the neighbours' cell states ${\mathbf{c}}_k$.

The paper’s GRU takes

\begin{align*}
\mathbf{h}{i}^{t} &= f{\textnormal{GRU}}(\mathbf{m}{i}^{t},;\mathbf{c}{i}^{t}) \
\tilde{\mathbf{h}}_i^t &= \mathbf{h}i^t + f{\textnormal{refine}}(\mathbf{h}_i^t)
\end{align*}

No $\mathbf{c}_t$ skip connection is needed because the gate already holds the old state. $\mathbf{h} = (1-\mathbf{z})\tilde{\mathbf{h}} + \mathbf{z}\,\mathbf{c}_t$ is an interpolation, so it's bounded, and $\mathbf{c}_{t+1} = \mathbf{h} + f_\text{refine}(\mathbf{h})$ doesn't blow up. If you take the paper literally instead (a separate $\mathbf{h}$) and add a $\mathbf{c}_t$ skip connection on top, which is what our first reimpl did, you get $\mathbf{c}_{t+1} = \mathbf{c}_t$ plus an $O(1)$ term that never goes to zero, and it blows up. $f_\text{refine}$: the "refinement MLP" (two dense layers, $C$ units, GELU), a [[ResNet|resblock]] on the GRU output. Morphogenesis: $\mathbf{c}_i^{t+1} = \tilde{\mathbf{h}}_i^t$ Lunar lander: $\mathbf{c}_i^{t+1} = W_C \, \tilde{\mathbf{h}}_i^t$ (state projection) and $\boldsymbol{\ell}_i^t = f_\text{act}(\tilde{\mathbf{h}}_i^t)$ (action logits). $f_\text{act}$: two-layer MLP (GELU), outputs 2D (noop, fire).