Schema Networks: Causal Generative Models

• Schema Networks are factor-graph-based generative models that encode relational and causal mechanisms to enhance interpretability and efficient planning.
• They use modular decompositions and disentangled cause–effect dynamics, enabling robust zero-shot transfer and outperforming common reinforcement learning baselines.
• Hidden Schema Networks extend these ideas by integrating neural models with explicit symbolic reasoning via schema-guided random walks for improved language modeling.

Schema Networks are structured, factor-graph-based generative models that explicitly encode relational and causal mechanisms in complex environments, facilitating generalization, interpretability, and efficient planning. Originating in the context of intuitive physics, these networks have since evolved into frameworks suitable for both reinforcement learning and language modeling, distinguished by their modular decomposition of state, disentangled cause–effect dynamics, and explicit handling of symbolic relational structures (Kansky et al., 2017, Sánchez et al., 2022).

1. Formal Structure and Schema Representation

Schema Networks (SNs) operate within the Markov Decision Process (MDP) formalism $M=(\mathcal{S}, \mathcal{A}, T, R, \gamma)$, where state is decomposed “object-orientedly” into $N$ entities $\{E_1, \dots, E_N\}$, each described by $M$ binary attributes $\alpha^{(t)}_{i,j}$. The global state vector is $s^{(t)} \in \{0,1\}^{N \cdot M}$. Transitions are factorized across entities and attributes:

$$p(s^{(t+1)}|s^{(t)}, a^{(t)}) = \prod_{i=1}^N \prod_{j=1}^M T_{i,j}\left(\alpha^{(t+1)}_{i,j} | s^{(t)}, a^{(t)}\right)$$

Each factor $T_{i,j}$ aggregates predictions via an OR over all schemas (logical rules), plus a self-transition:

$$T_{i,j}(\alpha_{i,j}^{(t+1)} = 1 | s, a) = 1 - \prod_{q=1}^Q [1 - \Pr(\phi^{k_q} = 1)] - \prod [?] [1 - \Pr(\Lambda_{i,j} = 1)]$$

A grounded schema $\phi^k$ is an AND over specific preconditions (attributes, action bits):

$N$0

Ungrounded (template) schemas parameterize the conjunction of any choice of entity indices, allowing for compositional and context-independent reuse (Kansky et al., 2017).

Hidden Schema Networks (HSNs) generalize this notion under an explicit global latent symbol graph $N$1 with $N$2 discrete nodes ($N$3, adjacency $N$4). Sentences map to sequences of symbols (schemas) via biased random walks on $N$5:

$N$6

This discrete compositionality supports linguistic or commonsense reasoning (Sánchez et al., 2022).

2. Learning Objectives and Training Algorithms

Learning in classic SNs targets the schema-template matrix $N$7 by minimizing a regularized empirical loss:

$N$8

The online addition of schemas employs a greedy, LP-relaxed procedure:

• Iteratively solve a sparse covering LP to maximally explain uncaptured positives without false positives and with minimal $N$9 complexity,
• Binarize via mixed-integer refinement,
• Expand the schema set until the budget or coverage constraints are met.

HSNs use variational inference with an ELBO augmented by mutual information:

$\{E_1, \dots, E_N\}$0

Discrete latent structures are handled via Gumbel–Softmax relaxations for differentiable sampling. All model parameters are co-trained using Adam (Sánchez et al., 2022).

3. Inference, Planning, and Reasoning

In SNs, the learned factor-graph defines a generative world model, supporting regression-style (backward-chaining) planning. Given a goal (e.g., desired rewards), planning is cast as a MAP-inference problem subject to logical constraints, reducible to 0/1-compatibility SAT. Approximate inference employs one round of max-product belief propagation to compute feasibility marginals, followed by a backward Viterbi or DFS-based action traceback and constraint augmentation to avoid negatives before action execution.

HSNs extend this reasoning paradigm to language and knowledge graphs. Posterior inference involves:

• Global graph induction via edgewise Bernoulli distributions,
• Sampling latent schema walks conditioned on input,
• Integrating these symbolic schemas via attention-based conditioning of pretrained LMs (BERT encoder, PSA-augmented GPT-2 decoder),
• Downstream autoregressive reasoning for combinatorial inference tasks (e.g., if-then commonsense knowledge generation).

4. Transfer, Generalization, and Experimental Findings

Schema Networks excel in zero-shot transfer owing to their modular, causality-disentangling architecture. In Breakout variations, training on standard layouts and immediate testing on novel scenarios ("Offset Paddle", "Middle Wall", "Random Target", etc.) yielded SN scores surpassing A3C and Progressive Networks by 10–30 points on average, often achieving human-level performance with no further adaptation (Kansky et al., 2017).

HSNs demonstrate strong generalization and interpretability in synthetic and natural language domains. On graph recovery tasks, ROC–AUC for edge prediction reached $\{E_1, \dots, E_N\}$1 for Barabási–Albert and $\{E_1, \dots, E_N\}$2 for Erdős–Rényi graphs. For language modeling, HSNs outperform fine-tuned GPT-2 and baselines in both perplexity and mutual information:

Model	PTB PPL	PTB MI	YAHOO PPL	YAHOO MI	YELP PPL	YELP MI
GPT2 FT	48.9	–	36.8	–	35.0	–
iVAE_MI	60.5	0.00	111.3	0.00	85.7	0.00
HSN (100,5)	45.7	2.8	32.9	2.2	31.4	2.0
HSN (100,20)	46.8	6.4	33.5	5.6	32.2	5.1

Shorter schemas optimize PPL, while longer schemas increase MI; small-world structure emerges in inferred graphs, and schema walks align with topical or sentiment clusters (Sánchez et al., 2022).

5. Symbolic–Neural Integration and Interpretability

SNs realize object-relational factorization, supporting interpretable causal reasoning: schemas codify the local dynamics (e.g., “if ball at $\{E_1, \dots, E_N\}$3 and paddle=LEFT, then ball shifts left”) and abstract across spatial layouts or object arrangements. The factor-graph generative transition offers visual and logical explanatory power (Kansky et al., 2017).

HSNs integrate discrete symbolic induction with neural models. The schema induction layer projects sentences into random walk sequences on the latent graph, labeling interpretable clusters (topics, sentiments, causal roles), and the graph structure exhibits high clustering coefficients and non-uniform degree distributions. Conditioned generation via the schema yields superior control and transparency in LLM outputs (Sánchez et al., 2022).

6. Applications: Planning, Language Modeling, Commonsense Reasoning

In reinforcement learning, Schema Networks provide generative world models that support planning by logical regression, achieving robust transfer across task variants. They act as both relational simulators and causal reasoners, separating dynamics from layout.

In language and knowledge domains, HSNs condition pretrained LMs on explicit symbol-paths to:

• Enhance perplexity and MI versus opaque LMs,
• Enable symbolic schema-guided generation,
• Support compositional “if-then” reasoning via structured random walk inference and auxiliary autoregressive models,
• Outperform COMET-GPT2 and COMET-BART on ATOMIC if-then generation (BLEU-2 and BERTScore: $\{E_1, \dots, E_N\}$4/$\{E_1, \dots, E_N\}$5 for HSN vs. $\{E_1, \dots, E_N\}$6/$\{E_1, \dots, E_N\}$7 for COMET-BART).

A plausible implication is that explicit relational inductive biases, factor-graph causal modeling, and symbolic–neural integration enable scalable, interpretable, and generalizable models across vision, planning, and language domains (Kansky et al., 2017, Sánchez et al., 2022).