Graph networks as learnable physics engines for inference and control

Motivation

• Object- and relation-centric representations reuse the same computations across elements, enabling combinatorial generalization to unexperienced scenarios with familiar components and compositional rules.The introduction motivates this approach with systems whose possibilities grow rapidly with their elements, such as multi-link chains and bouncing balls.Source: §1 · p. 1

Graph-network representation

• Graph networks encode bodies as nodes and joints as edges, with node, edge, and global update functions representing body dynamics, interactions, and system properties.Knowledge is shared across system elements, supporting generalization to new systems built from the same types of body and joint components.Source: §1 · pp. 1–2

Results

• Across seven simulated physical systems and one real robotic system, the models make accurate, generalizable predictions, identify hidden properties, and achieve competitive control performance.The authors describe this as a general-purpose learnable physics engine for complex 3D systems: it learns to approximate physical laws from current-state/next-state pairs without specific prior knowledge of those laws.Source: §1 · p. 2

Technical contributions

• The framework makes three technical contributions: graph-network forward models, inference models, and control algorithms that use those models for planning and policy learning.Source: §1 · p. 2

• Cognitive scientists have treated rich generative models as central to perception, reasoning, and decision-making.Source: §2 · p. 2
• The core implementation builds on graph neural networks (GNNs), while the work on physical dynamics models follows earlier learning approaches (Scarselli et al., 2005; 2009a;b; Battaglia et al., 2016; Chang et al., 2016; Watters et al., 2017).Source: §2 · p. 2
• The inference model shares aims with approaches to system identification, robust policies under hidden property variation, and exploration under uncertainty (Yu et al., 2017; Peng et al., 2017; Rajeswaran et al., 2016; Sun et al., 2011).Source: §2 · p. 2
• Model-based planning relates to classic control with predefined models and learned model-based control, while joint model-policy learning and a concurrent GNN policy approach provide related directions (Li & Todorov, 2004; Tassa et al., 2008; 2014; Atkeson & Santamaria, 1997; Deisenroth & Rasmussen, 2011; Levine & Abbeel, 2014; Heess et al., 2015; Gu et al., 2016; Nagabandi et al., 2017; Wang et al., 2018).Wang et al. (2018) used a GNN to approximate a policy, complementing this work's use of a related architecture for forward and inference models.Source: §2 · p. 2

Graph representation of a physical system

• A directed graph represents bodies as nodes and joints as edges, while separate static and dynamic graphs encode system parameters and instantaneous states and actions.Static features include global, body, and joint parameters; dynamic features include body positions, orientations, linear and angular velocities, and joint-action magnitudes.Source: §3 · p. 2, Figure 2 · p. 3

Graph networks

• Each graph-network block updates edges, then nodes, then global features, forming graph-to-graph message-passing modules that can be composed in deep or recurrent architectures.The three update functions are implemented with multilayer perceptrons.Source: §3 · p. 3

Forward models

• The one-step forward model stacks two unshared GNs with a graph skip connection, while a recurrent variant applies edge-, node-, and global-feature GRUs before a GN.Both models predict state differences; repeatedly feeding predictions and supplied controls back into the model produces long-range rollouts.Source: §3 · p. 3

Inference models

• The recurrent inference model uses observed dynamic-state trajectories under control inputs to encode unobserved static properties in a latent graph, which informs a one-step forward model.The full architecture is trained jointly to infer properties from observed behavior and use them to improve predictions.Source: §3 · pp. 3–4

Control algorithms

• Differentiable learned forward and inference models support gradient-based model-predictive control, while Stochastic Value Gradients train a model and policy function jointly.The paper also develops an agent that learns the GN-based model and policy function simultaneously.Source: §3 · p. 4

Environments

• Forward-model data came from seven actuated MuJoCo environments and human-controlled JACO stacking trajectories.The simulated environments were Pendulum, Cartpole, Acrobot, Swimmer, Cheetah, Walker2d, and JACO; simulated random controls generated state transitions.Source: §4 · p. 4
• Generalization and system-identification data varied physical parameters and structure across Pendulum, Cartpole, Swimmer, Cheetah, and JACO.Variations included link lengths, body masses, and motor gears; Swimmer link counts ranged from 3 to 15.Source: §4 · p. 4

MPC planning

• For MPC, the GN predicted trajectories and rewards, then optimized action sequences by backpropagating reward gradients and applying gradient descent.The objective was to maximize a dynamic-state-dependent reward over a chosen planning horizon from a given initial state.Source: §4 · p. 4

Model-based reinforcement learning

• For model-based RL, a differentiable GN simulator supplied return gradients to a jointly trained stochastic policy using single-step SVG(1).Comparisons used model-free SVG(0) stochastic policies and DDPG deterministic policies (Lillicrap et al., 2016).Source: §4 · p. 4

Baseline comparisons

• Forward-model baselines were a constant state-copy predictor and a learned MLP; MPC baselines included ground-truth-model DDP (Tassa et al., 2008; 2014), SVG(0), and DDPG.The MLP used the GN training data and flattened graph inputs and targets; its architecture was selected from 20 hyperparameter combinations.Source: §4 · p. 5

Prediction performance evaluation

• Evaluation averaged normalized squared errors for position, orientation, linear velocity, and angular velocity across 1,000 100-step test sequences.The metrics were one-step dynamic state differences and trajectory rollout differences, normalized individually to the constant-prediction baseline.Source: §4 · p. 5

Learning a forward model for a single system

• The GN model had lower error than the MLP baseline in 6 of 7 simulated control systems, and predictions across most other systems beat the constant prediction baseline.The advantage was especially pronounced for systems with repeated structure, such as Swimmer, while Pendulum showed negligible difference.Source: §5 · p. 5
• For Swimmer6, the GN had smaller relative error increases than the MLP from training to test data and to data from a learned DDPG agent.The authors suggest parameter sharing across bodies and joints may regularize the GN and support its stronger generalization.Source: §5 · pp. 5–6

Learning a forward model for multiple systems

• The GN made accurate predictions when static system properties, including body and joint masses, lengths, and angles, varied.These properties were supplied as node and edge attributes in the input graphs.Source: §5 · p. 6
• A single GN predicted across systems with different numbers of bodies and joint structures, including held-out Swimmers with 7 and 10 links; performance degraded gradually for 11–15 links.Held-out structures were not seen during training, and the caption reports 100-step rollout error across 3–15-link Swimmers.Source: §5 · p. 6, Figure 6 · p. 6

Real robot data

• On real JACO robot data from a human-controlled stacking task, the feedforward GN was less accurate than the recurrent GN, while both beat the constant prediction baseline.The recurrent model closely tracked test-set orientations and angular velocities.Source: §5 · p. 6, Figure 7 · p. 6

6. Results: Inference

• Across three Pendulum systems with unobserved variable lengths, predictions were very accurate after an ID phase with control inputs but very poor without them.The controlled ID phase made the systems identifiable; predicted trajectories then tracked ground truth in the subsequent 80-step prediction phase.Source: §6 · p. 6, Figure 9 · p. 7
• Lag or hysteresis may induce long-range temporal dependencies that the feed-forward model cannot capture.Source: §6 · p. 6
• After a 20-step ID phase using only observed dynamic variables, the model generally predicted well, sometimes approaching true-parameter performance; mismatched hidden parameters reduced performance to the no-ID level or worse.The ID phase observed randomly controlled dynamics; the figure reports 100-step rollout error relative to a constant-prediction baseline.Source: §6 · p. 6, Figure 8 · p. 7
• After identifying a system, the inferred representation can be stored and reused to predict trajectories from different initial states without repeating system identification.The paper contrasts this with an RNN approach that would need to identify the same system from data each time a trajectory with different initial conditions is predicted.Source: §6 · p. 7

Model-predictive control for single systems

• GN-based MPC produced target-pose and target-palm JACO trajectories and reasonable Cheetah movements for rewards targeting forward motion, height, vertical speed, or torso angular speed.The JACO demonstrations used a 20-step planning horizon; the Cheetah experiments also used a 20-step horizon.Source: §7 · p. 7, Figure 10 · p. 7
• Swimmer6 MPC approached Mujoco-planner and DDPG performance, and matching their performance when trained on random and DDPG trajectories.The random-action-trained model used one-step predictions but controlled with a 100-step planning horizon.Source: Figure 11 · p. 8

Model-predictive control for multiple systems

• A single GN model supported MPC across Pendulum, Cartpole, Acrobot, Swimmer6, and Cheetah, while zero-shot Swimmer control was near the Mujoco baseline for many link counts but worse at 3-5 and 11-15 links.The model was trained on Swimmers with 3-6 and 8-9 links and tested on 7 and 10-15 links.Source: §7 · p. 8, Figure 11 · p. 8

Model-predictive control with partial observations

• With 20 random actions for system identification, a GN forward model controlled Pendulum, Cartpole, SwimmerN, Cheetah, and JACO under partial observations.The system-identification phase implicitly inferred hidden properties before control.Source: §7 · p. 8

Model-based reinforcement learning

• The jointly trained GN-based SVG(1) agent reached higher performance in fewer episodes than SVG(0), while SVG(2-4) showed no significant gain and SVG(5+) sometimes performed worse.The model and policy were trained simultaneously; the learning curves summarize six random seeds.Source: §7 · p. 8, Figure 12 · p. 8

8. Discussion

• Across challenging physical systems, graph-network models were accurate, robust, and generalizable for prediction, system identification, and planning.The authors describe their object- and relation-centric inductive bias as the basis of this model class.Source: §8 · p. 8
• The models worked best when bodies and joints shared structure, were less successful when opportunities for sharing were limited, and did not address compounding long-horizon prediction errors.Swimmers and Cheetah illustrate the stronger- and weaker-sharing cases, respectively.Source: §8 · p. 8
• Proposed directions include real-world control, simulation-to-real transfer, stochastic settings, and identifying system structure alongside parameters.The authors also suggest imagination-based planning (Hamrick et al., 2017; Pascanu et al., 2017) and architectures integrating GN-like policies.Source: §8 · p. 8
• The authors frame compositional representations within statistical learning as a step toward robust, efficient, general-purpose reasoning and decision-making.Source: §8 · p. 8
• Preliminary experiments found little benefit from model pre-training, though the authors say further exploration is warranted.Source: §8 · p. 8

A. Summary of prediction and control videos

• The video appendix catalogs representative prediction rollouts from random-control-trained models and MPC control trajectories across initial states.Video labels distinguish prediction or control and fixed, parameterized, or system-identification settings, with a system abbreviation.Source: Table A.1 · p. 11, Table A.2 · p. 11

B. Description of the simulated environments

• The descriptions specify action and state variables for a pendulum setting with 2 total bodies and a linked-head setting with N+2 total bodies, alongside a random-parametrization label.Body density is kept constant when size changes.Source: B. Description of the simulated… · p. 12
• Listed environment variations include pendulum length (0.2-1) and mass (0.5-3), pole dimensions, and link counts from 3 to 9 with varying link lengths and thickness.The descriptions also include body dimensions and base-angle variation.Source: B. Description of the simulated… · p. 12

C.1. Random control

• Individual body masses and motor gears were each varied from 0.5 to 1.5.Source: §C.1 · p. 12
• Random control inputs were selected at random time steps from spline interpolations of randomly generated values.This was the default unless otherwise indicated.Source: §C.1 · p. 12
• Figure C.1 shows sample random sequences drawn from the same distribution used to generate random system training data.Source: §C.1 · p. 12, Figure C.1 · p. 12

C.2. Datasets

• 10,000 fixed-system and 20,000 parametrized-environment 100-step sequences were generated, with 1,000 and 5,000 validation and test sequences, respectively.For fixed systems, 10,000 sequences were generated for each system; parametrized environments used 20,000 sequences.Source: §C.2 · p. 13

C.3. Real JACO

• Real JACO data comprised 2,000 human-controlled 100-step stacking trajectories; joint angles and velocities represented state, and simulated JACO rendered the trajectories.The split was 1,800 training, 100 validation, and 100 test trajectories, sampled at 40 ms.Source: §C.3 · p. 13

D.1. Framework; D.2. Graph network architectures

• The implementation used custom TensorFlow/Sonnet GNs with ReLU MLPs, 2–3 edge layers of 256–512 units, 2 node/global layers of 128–256, and 128-sized updated features.Source: §D.1 · p. 13, §D.2 · p. 13

D.3. Data normalization

• Inputs were normalized to zero mean and unit variance using shared node/edge transforms; forward outputs were inverse-normalized, while inferred static latent graphs remained unnormalized.Source: §D.3 · p. 13

D.4. System invariance; D.5. Prediction of dynamic state change

• Individual translation-invariant systems were recentered before prediction; the model predicted state changes, updating position and velocities additively and orientation with a Hamilton product.Source: §D.4 · p. 14, §D.5 · p. 14

D.6.1. ONE-STEP PREDICTION; D.6.2. ONE-STEP PREDICTION WITH SYSTEM ID

• The forward model maps system parameters, state, and actions to the next state; System ID uses an observed sequence to infer the instance before prediction, reusing that inference throughout rollouts.Source: §D.6.1 · p. 14, §D.6.2 · p. 15

D.7.1. ONE-STEP; D.7.2. ABSTRACT PARAMETER INFERENCE

• One-step training used supervised delta prediction and input noise to restore joint-valid configurations; System ID was trained end-to-end from a random 20-step context and a separate supervised timestep.Normalization statistics were estimated online by accumulating counts, sums, and squared sums of input features and state changes.Source: §D.7.1 · p. 15, §D.7.2 · p. 15

D.7.3. RECURRENT ONE-STEP PREDICTIONS; D.7.4. LOSS

• The Real JACO recurrent predictor used teacher forcing on 21-step sequences, with the first 20 graphs as inputs and last 20 as targets; losses used normalized L2 for motion deltas and angle distance for orientation.The orientation loss minimizes 1 − cos^2(qe·qp), accounting for equivalent quaternions q and −q.Source: §D.7.3 · p. 16, §D.7.4 · p. 16

D.8. Training details

• Models used batches of 200 graphs or graph sequences and Adam on a single GPU, with starting learning rates tuned at 1−4.Source: §D.8 · p. 16
• Learning rates decayed exponentially by a factor of 0.975 every 50,000 steps for fast training or 200,000 steps for slow training.Source: §D.8 · p. 16
• Training stopped early or at asymptotic convergence based on validation rollout error over 20-step sequences; complex multi-environment models benefited from slow training for 1–3 weeks.Simple environments typically used fast training for less than a day to a few days; slow training yielded only marginal improvement in those cases.Source: §D.8 · p. 16

E. MLP baseline architectures

• Five ReLU MLP baselines spanned capacities from 3 layers of 128 cells to 9 layers of 512 cells, with the baseline selected by validation rollouts.Each model replaced the GN core while retaining graph-level normalization and update layers; 20 hyperparameter combinations were considered per environment.Source: E. MLP baseline architectures · pp. 16–17

F.1.1. MPC PLANNER WITH LEARNED MODELS / F.1.2. BASELINE MUJOCO-BASED PLANNER

• Learned-model MPC replanned in a receding horizon, while the MuJoCo iLQG baseline alternated forward rollouts and constrained backward control updates.The MPC reused a shifted optimal trajectory and typically ran 3–51 optimization iterations per initial state; iLQG was less effective with complex contacts but worked well for smoother swimmer dynamics.Source: §F.1.1 · p. 17, §F.1.2 · pp. 17–18

F.1.2. BASELINE MUJOCO-BASED PLANNER / F.2. Planning configuration

• Planning configurations covered balance, locomotion, and Jaco pose imitation, with horizons ranging from 20 to 100 steps.The iLQG objective assumed deterministic transitions; listed task rewards included forward velocity for Cheetah and pose loss for Jaco.Source: §F.1.2 · p. 18, §F.2 · p. 18

F.3. Reinforcement learning agents / F.3.1. DDPG

• The continuous-control agents learned off-policy from replay, and DDPG trained a deterministic policy alongside a neural action-value function.DDPG minimized one-step Bellman error for the value function and updated the policy using deterministic policy gradients; target networks were copied every 1000 gradient steps.Source: §F.3 · p. 18, §F.3.1 · p. 19

F.3.2. SVG

• SVG(0) is model-free, whereas SVG(N) uses the Graph Net dynamics model to unroll the system and compute model-based policy gradients.The paper used an action-value variant to compare SVG(0) with one-step SVG(1), and tested SVG(1) on Swimmer 6 using graphs built from observations and actions.Source: §F.3.2 · p. 19

G.1. Dynamic graph

• Dynamic graphs represented each body’s position, orientation, and velocities as nodes, and each joint’s action magnitude as an edge feature.Source: §G.1 · p. 20

G.2. Static graph

• Static graphs included global, body, joint, and actuator properties, but omitted geometry features such as size, density, and shape.The authors note that geometry information may be partly encoded by mass and the inertia tensor.Source: §G.2 · p. 20

H. Supplementary figures

• Supplementary figures show Swimmer 6 prediction on DDPG-generated trajectories and an architecture ablation in which two GNs outperform an equivalent Interaction Network.The ablation also indicates that global updates are important; another figure depicts how Swimmer 6 observations and actions are arranged as a directed graph with duplicated edges indicating direction.Source: Figure H.1 · p. 20, Figure H.2 · p. 21, Figure H.3 · p. 21