Dynamic Inference on Graphs using Structured
Transition Models
Saumya Saxena1 and Oliver Kroemer1

Abstract—Common everyday tasks such as picking up anobject in one smooth motion or pushing off a wall to quickly
turn a corner involve complex dynamic interactions between thehuman and the environment. These dynamic interactions are
critical in successful execution of these tasks. Thus, in orderto enable robots to perform such dynamic tasks effectively,
we need to consider the dynamics of the robot, the individualobjects, as well as the interactions between them. In this work,
we present a method that enables efficient learning of thedynamics of interacting systems by simultaneously learning a
dynamic graph structure and a stable and locally linear forward
dynamic model of the system. The dynamic graph structure
encodes evolving contact modes along a trajectory by making
probabilistic predictions over the edge activations. The learned
stable and locally linear dynamics enable the use of optimal
control algorithms such as iLQR for long-horizon planning andcontrol for complex interactive tasks. Through experiments insimulation and in the real world, we evaluate the performance ofour method by using the learned interaction dynamics for control
and demonstrate generalization to more objects and interactions
not seen during training.

I. I NTRODUCTION
In this work, we develop a method for efficiently learning the dynamics of interacting systems by simultaneously
learning a relational dynamic graph structure and a stablelocally linear forward dynamic model of the system. We use
the spring-mass-damper model as a structural prior over thelocal object-centric dynamics of our system. This provides us
with two benefits: linearity and stability. For positive valuesof mass, stiffness and damping parameters, the linear springmass-damper model is always stable about the equilibrium
point (which is also learned) [2].
We explore an exciting application of this work in the fieldof apprenticeship learning wherein the learned locally linear
model can be used for learning the quadratic cost function
underlying expert demonstrations using differentiable LQR as
the policy class [1, 4]. The learned behavior can then be
generalized to unseen goal conditions.
The key contributions of this work are two-fold: 1) a method
for learning stable locally linear dynamics for non-linear
interactive systems using graph neural networks by encodingchanging dynamics and contacts as part of the graph structure,enabling strong generalization properties to more objects inthe scene, and 2) using the learned locally linear dynamics todevise a robust control scheme that utilizes the recurrent nature
1 Robotics

MellonInstitute, School
University, PittsburghofPAokroemer}@andrew.cmu.edu

Computer Science, Carnegie
15123, USA {saumyas,of the learned graph structure to adapt the model predictions
and the policy to observed contact events.
II. M ETHOD
Our training pipeline is composed of two main modules thatare trained together, 1) a graph inference module that learns thedynamic graph structure, and 2) a forward dynamics modulethat learns the stable locally linear forward dynamics given
the graph structure. An overview of the method is shown inFig. 1.
A. Graph Inference Module
The state of the system at time t is represented as a graph
Gt = (Nt , Et ) where nodes Nt represent objects in the scene
and edges Et connect interacting nodes. Node nit ∈ Nt isrepresented using a feature vector xit where xit = [qit , q̇it ].
qit and q̇it are the position and velocity of the ith object
respectively. The edge eijt ∈ Et connecting nodes i and j isij
represented with edge features dij
t where dt is the distanceji
between nodes nt and nt .
Starting with a set of trajectories of fully connected graphs(G0 , u0 , . . . , GT , uT ) (where ut is the control applied at timet, and T is the length of the trajectory), we first pass themthrough the inference graph network that performs message
passing between the nodes and outputs embedded graphs.
Next, to aggregate information temporally, we pass the edge
embeddings mij
0,...,T through a GRU that outputs a discrete
probability distribution pijt over the edge types (active orinactive) for each edge. We sample from this distribution to
get the edge activations aij
t . Differentiability of this sampling
procedure can be ensured by employing techniques such as
Gumbel softmax [3].
B. Forward Dynamics Module
Using the edge activations output from the Graph Inference
Module, we remove the inactive edges from the fully
connected input graphs and then pass them through theforward dynamics graph network that outputs another setof embedded graphs. The node embeddings z̃ij
t are passed
though a fully connected neural network that outputs springmass-damper parameters [ m1i , kti , cit , x̂it ] independently fort
each node where mit is the mass, kti is the stiffness of thespring, cit is the damping, x̂it is the equilibrium point.
The loss function for training the model is written as,XX∗predij
Ldyn =||xi1:T − xi1:T ||2 +KL[pij0:T −1 || q0:T −1 ]ii,jGraph inference module

Spatial aggregation{xt }i{dt }
ij

Inferencegraphnetwork

Forward dynamics module

Temporal aggregation

{zt }

ij
ij
{m0 , …, mT }

i

ij
{mt }

GRU

{p0 , …, pT }ij

ij

Spatial aggregation

Edgeactivationsij
ij{a0 , …, aT }

{dt }
ij

i
{xt }

Forwarddynamics
graphnetwork

Locally linear dynamics prediction{m̃t }
ij

{z̃t}
i

i*
i
{xt+1} ⟺ {xt+1}
Loss
pred

SampleFully
connected
NN{At , Bt }ii

{xt+1 = At xt + Bt ut}
iiiiFig. 1. Method overview. Our method has two main parts 1) Graph inference module: we start with a set of fully connected graphs and pass them throughthe inference graph network that performs message passing between the nodes and outputs embedded graphs. Next, to aggregate information temporally, we
ijpass the edge embeddings mij
0,...,T through a GRU that outputs a discrete probability distribution pt over the edge types (active or inactive) for each edge.
ij
We sample from this distribution to get the edge activations at . 2) Forward dynamics module: using the edge activations, we remove the inactive edges fromthe fully connected input graphs and then pass them through the forward dynamics graph network that outputs another set of embedded graphs. The nodeembeddings z̃ijt are passed though a fully connected neural network that outputs stable locally linear transition dynamics. We forward propagate the learned
dynamics using the input control and get the next state which is used for loss calculation during training

pred

∗

where xi∗t+1 is the observed next state and xit+1 = Ait xit +Bit ut + oit is the predicted next state for node ni . pij
0:T −1 isthe probability distribution over predicted edge activations andqij
0:T −1 is a prior on the edge activations.

objects (orange) in the scene, not seen during training. Row 3shows Box2D Task 2 with same number of objects (two) as
seen during training. Row 4 shows Box2D Task 2 generalization scenario for picking up many more objects than during
training. Row 5 shows real world experiments with the 7DOF
Franka-Emika Panda arm performing a dynamic pickup taskusing the learned model.
We observe that, when tested with the same number of object interactions as seen during training, the N-step prediction
error is 0.16cm for task 1 and 0.14cm for task 2. When tested
for generalization (rows 2 and 4 in Fig. 2) the N-step prediction
error is 0.34cm for task 1 and 0.39cm for task 2 – less thana centimeter. Hence, we show strong generalization to moreobjects and interactions in the scene.
IV. C ONCLUSION AND F UTURE W ORK
In this work, we showed that by learning the dynamics
of interactive systems by simultaneously learning a dynamic
interaction graph and a stable locally linear forward dynamic
model given the graph, we can use the learned dynamics
for long-horizon control using iLQR-MPC. This work also
demonstrates on the generalization benefits of learning adynamic graph structure.
R EFERENCES

Fig. 2. Box2D and real world experiments

III. EXPERIMENTS
The proposed method is evaluated in simulation using theBox2D environment as shown in Fig. 2. We consider two tasks
in this environment: Task 1: Dynamically (without stopping)
picking up a 2D object using a 2D gripper and taking it toa goal region (Fig. 2 Rows 1 & 2), and Task 2: Dynamic
pickup of multiple objects (Fig. 2 Rows 3 & 4). For each ofthe tasks, we use the learned dynamic models for control using
iLQR-MPC and calculate the N-step prediction error. Row 1in Fig. 2 shows the executed trajectory for Box2D Task 1 withsame number of objects as seen during training. Row 2 shows
Box2D Task 1 generalization scenario with many distractor[1] Brandon Amos, Ivan Jimenez, Jacob Sacks, Byron Boots,and J Zico Kolter. Differentiable mpc for end-to-end
planning and control. Advances in neural information
processing systems, 31, 2018.
[2] Karl Johan Åström and Richard M Murray. Feedbacksystems, Chapter 8. Princeton university press, 2010.
[3] Eric Jang, Shixiang Gu, and Ben Poole. Categoricalreparameterization with gumbel-softmax. arXiv preprint
arXiv:1611.01144, 2016.
[4] Saumya Saxena, Alex LaGrassa, and Oliver Kroemer.
Learning reactive and predictive differentiable controllers
for switching linear dynamical models. In 2021 IEEEInternational Conference on Robotics and Automation(ICRA), pages 7563–7569. IEEE, 2021.