Beyond Static Graph World Models: Learning Stochastic Latent Dynamics over Evolving Topologies

1 INTRODUCTION

• Graph environments have evolving relational structure, and graph-based models may generalise from smaller training graphs to larger graphs where direct training is infeasible.The paper identifies social networks, navigation, and infrastructure as example contexts.Source: §1 · p. 1
• Existing graph world models are mostly limited to fixed topology, fail to exploit topology for message passing, or work only in deterministic, fully observable settings.GDM uses a sparse recurrent adjacency matrix with a recurrent state-space architecture for stochastic transitions in latent space.Source: §1 · p. 1
• GDD uses MMD to compare joint next-state distributions over graph topology, node features, and graph features.With a characteristic kernel, GDD is zero if and only if predicted and true distributions match.Source: §1 · p. 2
• GDM generalises to unseen graph sizes and outperforms state-of-the-art baselines on multi-step rollouts.The study also ablates the model and finds recurrent adjacency sampling important.Source: §1 · p. 2

2 RELATED WORK

• RSSM and other latent world models represent dynamics through recurrent states, variational inference, Gaussian-process dynamics, transformers, structured state-space models, or object-centric latents.The cited frameworks include RSSM (Hafner et al., 2019; 2025), DVBF, and PRSSM.Source: §2 · p. 2
• Graph world models include direct state-space transitions and latent-space RSSM variants, but the former are less effective in stochastic settings and G-RSSM does not use graph topology for message passing.Song & Cai (2026) study state-space transitions; Karacelebi et al. (2026) propose G-RSSM.Source: §2 · pp. 2–3
• Application-specific graph world models rely on domain knowledge and are not applicable to general graph-based environments.Examples include origami folding, satellite networks, and structural dynamics.Source: §2 · p. 3
• Some world models infer relational structure when observations are not graphs, while a survey identifies insufficient treatment of topological plasticity and probabilistic modelling in graph world models.The cited approaches include dynamic Bayesian networks and learned entity-pair relations.Source: §2 · p. 3
• World-model evaluation uses downstream planning or non-distributional metrics, while distributional metrics are less common in model-based RL.Planning can be computationally expensive, and a stochastic benchmark finds world models dramatically underestimate environmental variance.Source: §2 · p. 3

3 BACKGROUND

• A graph-based POMDP has graph observations and evolving edges; topology is fully observable, while node and graph features may be partially observable.Actions are sets of valid node selections, and complex edge changes can be composed from multiple selections.Source: §3 · p. 3
• MMD measures distance between distributions through their kernel mean embeddings in a reproducing kernel Hilbert space.For a characteristic kernel, MMD is zero if and only if the distributions match.Source: §3 · pp. 3–4
• An RSSM represents state with deterministic and stochastic latent variables, sampling stochastic state from a learned prior and correcting it with observations during training.Training minimises KL divergence between prior and posterior alongside reconstruction loss.Source: §3 · p. 4

4.1 ARCHITECTURE

• GDM assigns deterministic and stochastic latent vectors to each node and to the graph, and samples recurrent adjacency so predicted topology participates in message passing.The stochastic latent components and adjacency are factorised but their dependencies are modelled through conditional distributions.Source: §4.1 · p. 4
• The encoder creates node and graph embeddings by transforming each feature with a learned feature-specific linear map and averaging the resulting embeddings.Source: §4.1 · p. 4
• During training, posterior observations also update deterministic states because the ordinary RSSM update was less effective for dense observation features.The deterministic prior states are updated recurrently from previous state and action embeddings.Source: §4.1 · p. 4
• The model samples graph stochastic state and predicts edge additions or removals for each ordered node pair using Hard Concrete sampling.The adjacency update assumes incremental changes and uses differentiable edge weights for learning.Source: §4.1 · pp. 4–5
• SparseGAT-style message passing uses the predicted adjacency, while posterior training teacher-forces adjacency from the observed graph.The decoder uses prior deterministic state so observation-informed deterministic state cannot bypass stochastic state, then predicts graph and node features and other signals.Source: §4.1 · p. 5

4.2 TRAINING

• Training minimises the standard RSSM loss with prediction and KL terms, plus an alignment term for deterministic prior and posterior states.Prediction losses cover features, adjacency, action mask, reward, and continuation; KL terms combine node, graph, and adjacency components.Source: §4.2 · p. 5

5 GRAPH DISTRIBUTION DISTANCE

• The graph kernel combines topology similarity with node- and graph-feature kernels, using Jaccard for topology, multi-bandwidth RBF for continuous features, and weighted Hamming for categorical features.Source: §5 · p. 6
• Under finite aligned node sets and finite feature index and categorical value sets, the joint kernel is characteristic on attributed graphs.Source: §5 · p. 6
• A poor factor can drive the joint kernel toward zero and compress the metric's range, motivating additive terms that assess factors independently.Source: §5 · pp. 6–7
• The factorised kernel remains characteristic for any epsilon > 0; the implementation uses equal weights and biased MMD, with near-zero estimates for in-distribution samples.Source: §5 · p. 7
• GDD is the MMD between predicted and ground-truth graph distributions, comparing the full joint state rather than evaluating graph factors only separately.The joint state includes topology, node features, and graph features.Source: §5 · pp. 6–7

6 EXPERIMENTS

• The study compares adapted GWM, EA-GWM, and G-RSSM baselines across four environments, using random-rollout GDD as its primary evaluation metric.Models train for 5,000 epochs on random-policy rollouts; the environments span deterministic and stochastic, fully and partially observable settings.Source: §6 · p. 7

6.1 SIZE GENERALISATION FOR NEXT-STATE DISTRIBUTION PREDICTIONS

• GDM matches or beats all baselines on stochastic environments in single-step GDD tests across graph sizes.RGC was not evaluated on large graphs because episode lengths grow quadratically with graph size.Source: Table 1 · p. 8, §6.1 · p. 8
• GDD agrees with per-factor metric trends while additionally capturing joint dependence, and GDM achieves extremely low adjacency-matrix BCE.The paper attributes the low adjacency BCE to GDM's explicit recurrent adjacency matrix.Source: §6.1 · p. 8

6.2 MULTI-STEP ROLLOUTS

• GDM consistently outperforms baselines on 2–5-step rollouts over 50-node graphs; Vanilla GWM and EA-GWM differ little.The authors say EA-GWM's regularisation terms do not significantly improve multi-step rollout performance in these environments.Source: Figure 2 · p. 8, §6.2 · p. 8

6.3 ABLATIONS

• Recurrent adjacency has a significant effect on GDD across all environments, while other ablation effects vary by environment.The posterior correction matters significantly in OD but less elsewhere; SparseGAT and Hard Concrete matter more in some environments than others.Source: Figure 3 · p. 9, §6.3 · p. 9

• Independent edge sampling cannot capture inherently dependent edge events, even when individual edge probabilities are accurate.This limits fidelity in environments whose edge events are dependent.Source: §7 · p. 9
• GDM computes edge-update logits for every node pair and omits edge features, constraining scalability and applicable environments.The paper identifies adding edge features without O(|V|^2) computation as future work.Source: §7 · p. 9
• GDD's MMD component depends on the kernel, scales quadratically with sample count, and does not establish downstream-task performance.The metric also adds a weighted term because a characteristic kernel can compress the range of poor predictions.Source: §7 · p. 9

8 CONCLUSION

• GDM outperforms the baselines on GDD across environments and graph sizes outside the training distribution.The conclusion identifies edge features and the relationship between GDD and downstream-task performance as future directions.Source: §8 · p. 9
• GDM performs message passing on a sparse recurrent adjacency matrix and uses a recurrent state-space model to capture stochastic transitions in latent space.Its topology predictions are described as accurate.Source: §8 · p. 9
• GDD uses MMD with a characteristic graph kernel to compare joint next-state distributions across topology, node features, and graph features.The metric captures dependencies among these state components and provably distinguishes different distributions.Source: §8 · p. 9

A.1 PROOF OF LEMMA 1

• Continuous-feature RBF mixtures are C0-universal, while positive categorical weights yield a strictly positive-definite categorical kernel.The proof uses positive diagonal rescaling for continuous features and weighted Hamming distance for categorical features.Source: §A.1 · p. 13
• The Jaccard kernel is universal on the finite space of graph topologies, making the tensor-product joint kernel C0-universal and characteristic.The proof handles the empty graph as an isolated diagonal entry.Source: §A.1 · p. 13

A.2 PROOF OF LEMMA 2

• The additive kernel is positive definite because its continuous, categorical, and adjacency components are positive definite and have nonnegative weights.Source: §A.2 · p. 13
• Because the joint kernel is characteristic, zero squared MMD between two probability measures implies that the measures are equal.Source: §A.2 · p. 13

B ENVIRONMENT DESCRIPTIONS

• Environment observations include graph adjacency, categorical and continuous node and graph features, the number of node selections, and the action mask.Source: B ENVIRONMENT DESCRIPTIONS · p. 13

B.1 SEARCH AND RESCUE

• Search and Rescue is stochastic and partially observable: visits can reduce node stability and remove edges, while survivor presence is revealed only within two hops.Source: §B.1 · p. 14

B.2 OPINION DYNAMICS

• Opinion Dynamics is stochastic and fully observable; the environment rewires edges or copies opinions before the agent selects nodes for intervention.Source: §B.2 · pp. 14–15

B.3 CASCADING FAILURES

• Cascading Failures adds stochastic failures and partial observability to a power-grid failure model, with node failure governed by load relative to an unobserved safety factor.Source: §B.3 · p. 15

B.4 ROBUST GRAPH CONSTRUCTION

• Robust Graph Construction is deterministic and fully observable, with the agent adding edges to increase graph robustness under targeted node removal.Source: §B.4 · p. 16

C BASELINE IMPLEMENTATIONS

• The adapted G-RSSM uses environment-specific encoders and decoders, and a shared node-wise adjacency decoder whose parameter count is independent of graph size.Graph-level features are broadcast to nodes during encoding, and node-set actions are represented by node-level indicators.Source: C BASELINE IMPLEMENTATIONS · p. 16
• EA-GWM and Vanilla GWM use feature-specific encoders and heads for multimodal variable-size graphs, but differ in training objectives.EA-GWM adds rollout-consistency, spectral, and critical-node losses; GDM uses Hard Concrete sampling for sparse adjacency updates, unlike the Bernoulli adjacency sampling of the GWM baselines.Source: C BASELINE IMPLEMENTATIONS · pp. 16–17

D EVALUATION DETAILS

• GDD evaluation estimates transition distributions with 30 environment and model samples per reference state-action pair, then averages transition scores across episodes.Reward, continuation, and action mask are included as state features, except after terminal transitions when only reward and continuation are evaluated.Source: D EVALUATION DETAILS · p. 17

D.1 BANDWIDTH FITTING

• Categorical-feature weights are fitted from empirical change rates, while continuous-feature bandwidths use median nonzero one-step changes.Fitting uses sampled training transitions; stable categorical features receive greater weight than frequently varying ones.Source: §D.1 · p. 17

E.1 FEATURE SCALING

• Scalar targets are standardized using training-set statistics for reconstruction loss, then decoder predictions are transformed back to environment units.The same procedure is applied to the reward head.Source: §E.1 · p. 18

E.2 TRAINING DETAILS

• 5000-epoch training uses random-policy rollouts from graphs with 15, 18, 20, 23, and 25 nodes, while validation uses 20 graphs each at 20 and 30 nodes.The reported procedure trains 10 seeds per model.Source: §E.2 · p. 18

E.3 HYPERPARAMETER SEARCH

• Hyperparameter combinations are trained with 10 seeds, and the combination with the lowest mean validation GDD is selected for each model-environment pair.The search ranges are provided per model.Source: §E.3 · p. 18, Table 2 · p. 18

E.3.1 FINAL HYPERPARAMETERS

• Final hyperparameters are reported for each model and environment, with asterisks marking parameters that were not searched.Hard Concrete parameters β = 2/3, γ = −0.1, and ζ = 1.1 are used across environments; mean pooling is used throughout.Source: §E.3.1 · p. 18, Table 3 · p. 19

F.1 GDD VERIFICATION

• Self-GDD bias is negligible across environments and graph sizes, and is exactly zero for deterministic RGC.Table 4 reports self-GDD using κ = 30 samples.Source: Table 4 · p. 20, §F.1 · p. 20

F.2 TRAINING CURVES

• Training-loss values are not directly comparable across methods because their loss equations and coefficients differ.In OD, GDM loss rises after an initially stable period as reward prediction improves at the expense of higher node-feature KL loss.Source: §F.2 · p. 20, Figure 4 · p. 20

F.3 FACTOR-WISE DISTRIBUTION METRICS

• GDM predicts adjacency with extremely low BCE; G-RSSM has high adjacency BCE across environments, while Vanilla GWM and EA-GWM struggle especially with static edges.Source: §F.3 · p. 20, Figure 5 · p. 21
• Vanilla GWM and EA-GWM capture some features more accurately; routing posterior information through sampled stochastic latents limits RSSM representations, though GDM’s deterministic posterior update partially mitigates this and information loss can remain.Source: §F.3 · pp. 20–21

F.4 MULTI-STEP ROLLOUT RESULTS

• Multi-step evaluation uses 100 samples across rollout lengths of one to five steps; invalid actions are handled as no-ops in the environment and passed through in the model.Source: §F.4 · p. 23, Figure 7 · p. 23

F.5 ABLATION RESULTS

• The ablation study reports single-step GDD by ablation and environment, including a combined removal of graph state, h+ correction, and recurrent adjacency.Source: §F.5 · p. 24, Table 5 · p. 24