Source-linked AI summary
Graph networks as learnable physics engines for inference and control
Alvaro Sanchez-Gonzalez, Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, Peter Battaglia
TL;DR
The paper asks how agents can understand and control complex physical scenes whose structure is difficult to represent. It introduces graph-network models with object- and relation-centric inductive bias, finding accurate, generalizable prediction across eight physical systems, system identification, and differentiable planning and policy optimization.
Problem
Understanding and controlling complex physical scenes requires rich structural knowledge, but their combinatorial complexity makes this difficult.
Method
The paper uses graph networks to represent bodies as nodes and joints as edges, learning forward, inference, and control models over these structured representations.
Results
Across eight physical systems, GN-based models supported accurate and generalizable prediction, system identification, and competitive control performance.
Takeaways & Limitations
The framework supports learnable, object- and relation-centric modeling for prediction, inference, planning, and policy learning in challenging physical systems.
Takeaways & Limitations
The models were less successful when systems offered little repeated structure, and they do not address error compounding over long trajectory predictions.
Abstract
from arXiv · showhide
Understanding and interacting with everyday physical scenes requires rich knowledge about the structure of the world, represented either implicitly in a value or policy function, or explicitly in a transition model. Here we introduce a new class of learnable models--based on graph networks--which implement an inductive bias for object- and relation-centric representations of complex, dynamical systems. Our results show that as a forward model, our approach supports accurate predictions from real and simulated data, and surprisingly strong and efficient generalization, across eight distinct physical systems which we varied parametrically and structurally. We also found that our inference model can perform system identification. Our models are also differentiable, and support online planning via gradient-based trajectory optimization, as well as offline policy optimization. Our framework offers new opportunities for harnessing and exploiting rich knowledge about the world, and takes a key step toward building machines with more human-like representations of the world.
1. Introduction
The paper addresses how intelligent agents can understand and control combinatorially complex physical systems by learning object- and relation-centric graph representations. Its GN-based models share learned dynamics across system elements and support prediction, inference, and control across varied physical systems.
- 1. Introduction: Physical systems become combinatorially complex as the number of elements, possible configurations, or interactions increases.The paper gives multi-link chains and bouncing-ball systems as examples of rapidly growing possibilities.
- 1. Introduction: Graph representations apply shared objectwise and relation-wise computations, enabling combinatorial generalization to new scenarios built from understood components and rules.Particle-based physics engines illustrate this principle by reusing common body dynamics and interaction rules.
- 1. Introduction: Graph networks represent physical bodies as nodes and joints as edges, with learned node, edge, and global updates encoding body, interaction, and system dynamics.Knowledge is shared across system elements, supporting generalization to new systems composed of the same body and joint building blocks.
- 1. Introduction: Across seven simulated systems and one real robotic system, GN-based models supported accurate generalizable prediction, system identification, and competitive control.The model learns to approximate physical laws from supervised current-state/next-state pairs without specific a priori physical-law knowledge.
- 1. Introduction: The work contributes GN-based forward models, inference models, and control algorithms for graph-based physical modeling, planning, and policy learning.The control algorithm uses the learned forward and inference models for planning and policy learning.
2. Related Work
The paper situates its approach among research on generative models, graph neural networks, learned physical dynamics, system identification, robust policies, exploration, and model-based planning.
- 2. Related Work: The work draws on cognitive-science accounts that place rich generative models at the center of perception, reasoning, and decision-making.
- 2. Related Work: Its core implementation belongs to the broader class of graph neural networks, while its physical-dynamics contribution builds on earlier learned-dynamics approaches.
- 2. Related Work: The inference model relates to work on explicit system identification, policies robust to hidden-property variation, and exploration under uncertainty.
- 2. Related Work: The learned models are used for model-based planning in a manner related to classic approaches using predefined models.
3. Model
The model represents physical scenes with graphs and processes them through graph-network blocks to predict dynamics, infer hidden properties, and support differentiable control. Separate static and dynamic representations allow recurrent inference to encode unobserved system properties for forward prediction.
- Graph representation of a physical system: A physical system graph contains global features, node features for bodies, and directed edge features with sender and receiver indices for joints.
- Graph representation of a physical system: Static and dynamic properties are represented separately: static graphs encode system parameters, while dynamic graphs encode instantaneous body states.Static parameters include gravity, mass, inertia, joint properties, and motor properties; dynamic states include positions and quaternion orientations.
- Graph networks: A GN block maps an input graph to an output graph with the same structure but updated edge, node, and global features.Its edge-wise, node-wise, and global functions perform message passing in sequence.
- Forward models: The feed-forward forward model uses sequential GNs to predict each body’s state one time-step ahead from the current graph.The first GN produces a latent graph that is concatenated with the input before the second GN returns predictions.
- Forward models: The recurrent forward model applies GRU sub-modules across edges, nodes, and global features before a GN block.
- Inference models: The recurrent inference model observes dynamic-state trajectories and constructs latent representations of unobserved static properties for more accurate forward prediction.It graph-concatenates inferred information with a dynamic input and updates a hidden graph through recurrent processing.
- Control algorithms: Because the GN is differentiable, the learned forward and inference models support gradient-based trajectory optimization and model-based control, while SVG supports joint model and policy learning.
4. Methods
The experiments trained GN forward models on simulated and real robotic trajectories, testing prediction, generalization, and model-based control across varied physical systems.
- Experimental setup: Training data came from seven actuated Mujoco environments and recorded JACO robot trajectories under human control.The simulated data used random controls, while the real data came from a stacking task.
- Experimental setup: The generalization dataset varied system parameters and structure, including link lengths, masses, motor gears, and Swimmer link counts from 3-15.
- Control: GN-based MPC optimized action sequences by backpropagating reward gradients through predicted N-step trajectories.
- Control: Within an SVG(1) regime, the GN forward model served as a differentiable environment simulator while the stochastic policy was trained jointly with it.
- Evaluation: Models were compared with constant-prediction and learned MLP baselines, plus a ground-truth-model DDPG-based or MPC control baseline where applicable.
- Evaluation: Prediction was evaluated using normalized squared one-step state-difference and trajectory rollout errors, averaged across position, orientation, and velocity components.Reported errors used 1000 100-step test sequences unless otherwise specified.
5. Results: Prediction
GN forward models predicted dynamics accurately across simulated and real systems, generalized across parameter and structural variation, and often outperformed MLP baselines. Generalization was strongest for familiar or closely related structures, while recurrent GNs improved real-robot rollout accuracy.
- Single-system prediction: GN models had lower prediction error than MLP baselines in 6 of 7 simulated control systems.The difference was especially pronounced for repeatedly structured Swimmers and negligible for Pendulums.
- Generalization: GN models showed smaller error increases than MLPs from training data to test data and DDPG-agent trajectories in Swimmer6.
- Parametric and structural generalization: A single GN accurately predicted systems with continuously varying masses, lengths, joint angles, and other static parameters.These properties were supplied through node and edge attributes.
- Parametric and structural generalization: GN models trained on mixed Swimmer datasets also predicted systems with different numbers of links and completely different system structures.
- Parametric and structural generalization: Zero-shot prediction was very accurate for held-out 7- and 10-link Swimmers, then degraded gradually for 11-15-link Swimmers while remaining visually close to ground truth.
- Real robot prediction: For real JACO trajectories, the recurrent GN tracked orientations and angular velocities more accurately than the feed-forward GN over 100-step rollouts.The feed-forward model nevertheless performed below the constant-prediction baseline.
6. Results: Inference
The GN inference model identifies hidden system properties from an initial observation phase and supports accurate subsequent trajectory predictions when that phase is informative. Performance degrades when the identification phase does not make the hidden system identifiable, and inferred representations can be reused across initial states.
- The inference model predicted accurately after observing 20 steps of randomly controlled dynamics while only dynamic state variables were observed.The observed variables were position, orientation, linear velocity, and angular velocity; hidden properties were inferred during the identification phase.
- Informative control inputs during identification produced very accurate predictions for pendulums with unobserved lengths, whereas no-input identification produced very poor predictions.
- The feed-forward model may fail to capture long-range temporal dependencies caused by lag or hysteresis.
- The system-identification model generally approached the performance of a model given true static parameters, while mismatched identification parameters could perform as poorly as having no identification information.
- Once identification is complete, the inferred representation can be stored and reused to predict trajectories from different initial states.This avoids re-identifying the same system for every new trajectory.
7. Results: Control
The GN models support gradient-based model-predictive control and model-based reinforcement learning across several physical systems. They achieve competitive planning, generalize across system structures, and improve learning efficiency in the tested SVG setting.
- Model-predictive control: GN-based MPC optimized action sequences by backpropagating predicted trajectory rewards through the differentiable forward model.The method repeatedly minimized negative predicted reward by gradient descent over the planning horizon.
- Model-predictive control: GN-based MPC generated target-reaching JACO trajectories and reasonable Cheetah movements under 20-step planning horizons.
- Model-predictive control: Swimmer6 MPC performance was close to a planner using true Mujoco physics, and matched strong baselines when trained on random and DDPG trajectories.
- Model-predictive control: A single GN model trained on Swimmers with 3-6 and 8-9 links generalized nearly as well as the Mujoco baseline to Swimmers with 7 and 10-15 links.
- Model-based reinforcement learning: The GN-based SVG(1) agent reached higher performance after fewer episodes than the model-free SVG(0) agent.
- Model-based reinforcement learning: Using more than one forward step in SVG did not significantly improve performance and could worsen it for SVG(5+).
8. Discussion
The work presents graph networks as accurate, robust, and generalizable learnable models for prediction, system identification, and planning in challenging physical systems. Its benefits are strongest when bodies and joints share structure, while long-horizon error compounding remains unresolved.
- Graph networks provide an object- and relation-centric inductive bias for learnable forward and inference models.
- Across experiments, the models were accurate, robust, and generalizable for prediction, system identification, and planning.
- Performance was strongest in systems with common body and joint structure, such as Swimmers, and weaker where opportunities for parameter sharing were limited, such as Cheetah.
- The approach does not address the compounding of errors over long trajectory predictions, a common model-based planning problem.
- Future work includes real-world control, simulation-to-real transfer, stochastic environments, and system identification of system structure as well as parameters.
- Preliminary experiments found little benefit from pre-training the model, although further exploration was warranted.
Supplementary Material: Graph Networks as Learnable Physics Engines for
The supplementary material documents experimental configurations, parameter ranges, random-data generation, and representative prediction and control trajectories. These details span varied physical systems and their action, state, and structural parameters.
- Representative trajectory tables encode whether videos show prediction or control, fixed or parameterized settings, system identification, and the system abbreviation.
- The systems included pendulum, cartpole, and multi-link configurations with specified state and action dimensions.
- The supplementary configurations varied system structure and physical parameters, including link counts, lengths, thicknesses, masses, and motor gears.
- Training data were generally generated with random control inputs sampled from spline interpolations of randomly generated values.
- The supplementary figures provide sample random sequences and video references for the generated trajectories.
C.2. Datasets
The experiments use simulated and real trajectories to train, validate, and test graph-network models across fixed and parametrized physical systems. Forward prediction combines static system information with dynamic states and actions, while system identification infers latent properties from observed sequences.
- Datasets: Fixed systems used 10,000 100-step training sequences, plus 1,000 validation and 1,000 test sequences.These corresponded to about 10^6 supervised training examples.
- Datasets: Parametrized environments used 20,000 100-step training sequences, plus 5,000 validation and 5,000 test sequences.The training data represented about 2 · 10^6 supervised examples.
- Datasets: Real JACO data contained 2,000 100-step trajectories, split into 1,800 training, 100 validation, and 100 test trajectories.Each timestep was 40 ms, and observations used joint-angle sine/cosine values and joint velocities.
- Model inputs and outputs: The forward model constructs static and dynamic graphs from system parameters, current state, and actions, then predicts a delta state used to update the system.Input and output features are normalized, while inferred latent static graphs are not normalized.
- System identification: System identification uses a 20-step observed sequence to infer an abstract static graph before producing the next-state prediction.During rollouts, identification is performed once and the same inferred graph is reused for all one-step predictions.
- Training procedure: Training adds noise to dynamic input nodes so predictions remain within joint-compatible configurations.The recurrent inference GN and dynamics GNs are trained end-to-end using a shared prediction loss.
D.8. Training details
Training used Adam with validation rollout error for early stopping or convergence, with learning-rate decay schedules selected according to model complexity. Simpler environments trained faster, while complex models required longer training.
- Optimization: Models used batch size 200, Adam optimization, an initial learning rate of 1−4, and decay factor 0.975.Decay was updated every 50,000 steps for fast training or 200,000 steps for slow training.
- Stopping and duration: Training stopped early or at asymptotic convergence according to 20-step validation rollout error.Simple environments typically trained in less than a day to a few days, while complex models used longer schedules.
E. MLP baseline architectures
The MLP baseline spans five ReLU architectures and replaces the graph-network core with layers that flatten graph inputs and reconstruct graph outputs. Planning and reinforcement-learning comparisons use explicit trajectory objectives and model-free or model-based baselines.
- MLP architectures: The MLP baseline includes five ReLU networks ranging from 3 to 9 hidden layers and 128 to 512 hidden cells per layer.The listed configurations include 3×128, 3×512, 5×256, 9×128, and 9×512 hidden-layer/cell combinations.
- Planning baselines: MPC optimizes action sequences by rolling out a learned dynamics model and applying gradients of total reward with respect to the actions.The planner operates over a specified horizon and iteratively updates the action sequence.
- MLP implementation: The MLP replaces the two-layer GN core by flattening input graphs and reconstructing output graphs, while retaining graph-level normalization and update layers.The best validation-rollout model for each environment was selected from 20 hyperparameter combinations.
- Planning objective: The total planning objective sums per-timestep costs and a terminal cost under deterministic transitions xt+1 = f(xt, ut).The objective is evaluated over a finite horizon from an initial state.
- Planning limitations: The iLQG baseline is not optimal for complex contacts, although DDP-style methods work well for relatively smooth swimmer dynamics.A ground-truth Mujoco model was also used for a differential dynamic programming baseline.
- RL baselines: RL comparisons include model-free DDPG and SVG(0), alongside model-based SVG(N) using the graph-network model to unroll dynamics for N steps.The model-based variant directly compares a one-step GN horizon with the model-free SVG(0) baseline.