Propagation Networks for Model-Based Control Under Partial Observation

I Introduction

Physics engines are critical for planning and control in robotics. To plan for a task, a robot may use a physics engine to simulate the effects of different actions on the environment and then select a sequence of actions to reach a desired goal configuration. The utility of the resulting action sequence depends on the accurate prediction of the physics engine; so a high-fidelity physics engine plays a critical role in robot planning. Most physics engines used in robotics, such as Mujoco , Bullet , and Drake , use approximate contact models, and recent studies , , have demonstrated discrepancies between their predictions and real-world data. These mismatches prevent the above physics engines from solving contact-rich tasks.

Recently, researchers have started building general-purpose neural physics simulators, aiming to approximate complex physical interactions with neural networks . They have succeeded to model the dynamics of both rigid bodies and deformable objects (e.g., ropes). More recent work has used interaction networks for discrete and continuous control .

Interaction networks, however, have two major limitations. First, interaction nets only consider pairwise interactions between objects, restricting its use in real-world scenarios, where simultaneous multi-body interactions often occur. Typical examples include Newton’s cradle (Fig. 1(a)) or rope manipulation (Fig. 1(b)). Second, they need to observe the full states of a environment; however, many real-world control tasks involve dealing with partial observable states. Fig. 1(c) shows an example, where a robot wants to push a set of blocks into a target configuration; however, only the red blocks in the top layer are visible to the camera.

In this paper, we introduce Propagation Networks (PropNet), a differentiable, learnable engine that simulates multi-body object interactions. PropNet handles partially observable situations by operating on a latent dynamics representation; it also enables instantaneous propagation of signals beyond pairwise interactions using multi-step effect propagation. Specifically, by representing a scene as a graph, where objects are the vertices and object interactions are the directed edges, we initialize and propagate the signals through the directed paths in the interaction graph at each time step.

Experiments demonstrate that PropNet consistently outperforms interaction networks in forward simulation. PropNet’s ability to accurately handle partially observable states brings significant benefits for control. Compared with interaction nets and state-of-the-art model-free deep reinforcement learning algorithms, model-based control using propagation networks is more sample-efficient, accurate, and generalizes better to new, partially observable scenarios.Our project page: http://propnet.csail.mit.edu

II Related Work

In recent years, researchers have been building differentiable physics simulators in various forms . For example, approximate, analytical differentiable rigid body simulators have been deployed for tool manipulation and tool-use planning .

Among them, two notable efforts on learning differentiable simulators include interaction networks and neural physics engines . These methods restrict themselves to pairwise interactions for generalizability. However, this simplification limits their ability to handle simultaneous, multi-body interactions. In this work, we tackle this problem by learning to propagate the signals multiple steps on the interaction graph. Gilmer et al. have recently explored message passing networks, but with a focus on quantum chemistry.

II-B Model-Predictive Control with a Learned Simulator

Recent work on model-predictive control with deep networks often learns an abstract-state transition function, instead of an explicit account of environments . Subsequently, they use the learned model or value function to guide the training of the policy network. Instead, PropNet learns a general physics simulator that takes raw object observations (e.g., positions, velocities) as input. We then integrate it into classic trajectory optimization algorithms for control.

A few recent papers exploit the power of interaction networks for planning and control. Many of them use interaction networks to imagine—rolling out approximate predictions—to facilitate training a policy network . In contrast, we use propagation networks as a learned dynamics simulator and directly optimize trajectories for continuous control. By separating model learning and control, our model generalizes better to novel scenarios. Recently, Sanchez-Gonzalez et al. also explored applying interaction networks for control. Compared with them, our propagation networks can handle simultaneous multi-body interactions and deal with partially observable scenarios.

III Learning the Dynamics

We assume that the interactions within a physical system can be represented as a directed graph, $G=\langle O,R\rangle$, where vertices $O$ represent the objects, and edges $R$ correspond to their relations (Fig. 3). Graph $G$ can be represented as

Specifically, $o_{i}=\langle x_{i},a^{o}_{i},p_{i}\rangle$, where $x_{i}=\langle q_{i},\dot{q}_{i}\rangle$ is the state of object $i$, containing its position $q_{i}$ and velocity $\dot{q}_{i}$. $a^{o}_{i}$ denote its attributes (e.g., mass, radius), and $p_{i}$ is the external force on object $i$. For the relations, we have

where $u_{k}$ is the receiver, $v_{k}$ is the sender, and $a^{r}_{k}$ is the type and attributes of relation $k$ (e.g., collision, spring connection).

Our goal is to build a learnable physics engine to approximate the underlying physical interactions. We can then use it to infer the system dynamics and predict the future from the observed interaction graph $G$:

where $G_{t}$ denotes the scene states at time $t$. We aim to learn $\phi(\cdot)$, a learnable dynamics model, to minimize $\|G_{t+1}-\phi(G_{t})\|_{2}$.

Below we review our baseline model Interaction Networks (IN) . IN is a general-purpose, learnable physics engine, performing object- and relation-centric reasoning about physics. IN defines an object function $f_{O}$ and a relation function $f_{R}$ to model objects and their relations in a compositional way. The future state at time $t+1$ is predicted as

where $o_{i,t}=\langle x_{i,t},a^{o}_{i},p_{i,t}\rangle$ denotes object $i$ at time $t$, $u_{k}$ and $v_{k}$ are the receiver and sender of relation $r_{k}$, and $\mathcal{N}_{i}$ denotes the relations where object $i$ is the receiver.

III-B Propagation Networks

IN defines a flexible and efficient model for explicit reasoning of objects and their relations in a complex system. It can handle a variable number of objects and relations and has performed well in domains like n-body systems, bouncing balls, and falling strings. However, one fundamental limitation of IN is that at every time step $t$, it only considers local information in the graph $G$ and cannot handle instantaneous propagation of forces, such as Newton’s cradle shown in Fig. 2, where ball A’s impact produces a compression wave that propagates through the balls immediately . As force propagation is a common phenomenon in rigid-body dynamics, this shortcoming has limited IN’s practical applicability.

To address the above issues, we propose Propagation Networks (PropNet) to handle the instantaneous propagation of forces efficiently. Our method is inspired by message passing, a classic algorithm in graphical models .

Effect propagation requires multi-step message passing along the directed edges in graph $G$. Forces ejected from ball A (Fig. 2) should be propagated through the connected balls to ball B within a single time step. Force propagation is hard to analyze analytically for complex scenes. Therefore, we let PropNet learn to decide whether an effect should be propagated further or withheld.

At time $t$, we denote the propagating effect from relation $k$ at propagation step $l$ as $e^{l}_{k,t}$, and the propagating effect from object $i$ as $h^{l}_{i,t}$. Here, we have $1\leq l\leq L$, where $L$ is the maximum propagation steps within each step of the simulation. Propagation can be described as

where $f^{l}_{O}(\cdot)$ denotes the object propagator at propagation step $l$， and $f^{l}_{R}(\cdot)$ denotes the relation propagator. Depending on the complexity of the task, the network weights can be shared among propagators at different propagation steps.

We name this model Vanilla PropNet. Experimental results show that the selection of $L$ is task-specific, and usually a small $L$ (e.g., $L=3$) can achieve a good trade-off between the performance and efficiency.

III-B2 Object- and relation-encoding with residual connections

We notice that Vanilla PropNet is not efficient for fast online control. As information such as states $o_{i,t}$ and attributes $a^{r}_{k}$ are fixed at a specific time step, they can be shared without re-computation between each sequential propagation step. Hence, inspired by the ideas on fast RNNs training , we propose to encode the shared information beforehand and reuse them along the propagation steps. We denote the encoder for objects as $f^{\text{enc}}_{O}(\cdot)$ and the encoder for relations as $f^{\text{enc}}_{R}(\cdot)$. Then,

In practice, we add residual links between adjacent propagation steps that connect $h^{l}_{i,t}$ and $h^{l-1}_{i,t}$. This helps address gradient vanishing and exploding problem, and provides access to historical effects. The update rules become

where propagators $f^{l}_{O}(\cdot)$ and $f^{l}_{R}(\cdot)$ now take a new sets of inputs, which is different from Vanilla PropNet.

Based on the assumption that the effects between propagation steps can be represented as simple transformations (e.g., identity-mapping in Newton’s cradle), we can use small networks as function approximators for the propagators $f^{l}_{O}(\cdot)$ and $f^{l}_{R}(\cdot)$ for better efficiency. We name this updated model Propagation Networks (PropNet).

III-C Partially Observable Scenarios

For many real-world situations, however, it is often hard or impossible to estimate the full state of environments. We extend Eqn. 3 using PropNets to handle such partially observable cases by operating on a latent dynamics model:

where $\tau(\cdot)$ is an encoding function that maps the current observation to a latent representation. As shown in Figure 4b, $\tau(\cdot)$ consists of two parts: first, PropNets $g(\cdot)$ that map the current observation to object-centric representations; second, $\omega(\cdot)$ that aggregates the object-centric representations into a fixed-dimensional global representation. We use a global representation for partially observable cases, because the number and set of observable objects vary over time, making it hard to define object-centric dynamics. In fully observable environments, $\tau(\cdot)$ reduces to an identity mapping and the dynamics is defined on the object level over the state space (Eqn. 3 and Fig. 4a). To train such a latent dynamics model, we seek to minimize the loss function: $\mathcal{L}_{\text{forward}}=\|\tau(G_{t+1})-\phi(\tau(G_{t}))\|_{2}$.

In practice, we use a small history window of length $T_{\text{history}}$ for the state representation, i.e., the input to $\phi(\cdot)$ is the concatenation of $\tau(G_{t}),\tau(G_{t-1}),...,\tau(G_{t-T_{\text{history}}+1})$.

Using the above loss alone leads to trivial solutions such as $\phi(x)=\tau(x)=0$ for any valid $x$. We tackle this based on an intuitive idea: an ideal encoding function $\tau(\cdot)$ should be able to reserve information about the scene observation. Hence, we use an aggregation function $\omega(\cdot)$ that has no learnable parameters like summation or average and introduce a decoding function $\psi(\cdot)$ to ensure a nontrivial $\tau(\cdot)$ by minimizing an additional auto-encoder reconstruction loss : $\mathcal{L}_{\text{encode}}=\|G-\psi(g(G))\|_{2}$, where $\psi(\cdot)$ is realized as PropNets. The full model is shown in Figure 4b.

IV Control Using Learned Dynamics

Compared to model-free approaches, model-based methods offer many advantages, such as generalization and sample efficiency, as it can approximate the policy gradient or value estimation without exhausted trials and errors.

However, an accurate model of the environment is often hard to specify and brings significant computational costs for even a single-step forward simulation. It would be desirable to learn to approximate the underlying dynamics from data.

A learned dynamics model is naturally differentiable. Given the model and a desired goal, we can perform forward simulation, optimizing the control inputs by minimizing a loss between simulated results and a goal. The model can also estimate the uncertain attributes online by minimizing the difference between predicted future and actual outcome. Alg. 1 outlines our control algorithm, which provides a natural testbed for evaluating the dynamics models.

Let $\mathcal{G}_{g}$ be our goal and $\hat{u}_{1:T}$ be the control inputs (decision variables), where $T$ is the time horizon. These task-specific control inputs are part of the dynamics graph. Typical choices include observable objects’ initial velocity/position and external forces/attributes on objects/relations. We denote the graph encoding as $G^{\tau}=\tau(G)$, and the resulting trajectory after applying the control inputs as $\mathcal{G}=\{G_{i}^{\tau}\}_{i=1:T}$. The task here is to determine the control inputs by minimizing the gap between the actual outcome and the specified goal $\mathcal{L}_{\text{goal}}(\mathcal{G},\mathcal{G}_{g})$.

Our propagation networks can do forward simulation by taking the dynamics graph at time $t$ as input, and produce the graph at next time step, $\hat{G}_{t+1}^{\tau}=\phi(G_{t}^{\tau})$. Let’s denote the forward simulation from time step $t$ as $\hat{\mathcal{G}}=\{\hat{G}_{i}^{\tau}\}_{i=t+1\dots T}$ and the history until time $t$ as $\bar{\mathcal{G}}=\{G_{i}^{\tau}\}_{i=1\dots t}$. We can back-propagate from the loss $\mathcal{L}_{g}(\bar{\mathcal{G}}\cup\hat{\mathcal{G}},\mathcal{G}_{g})$ and use stochastic gradient descent (SGD) to update the control inputs. This is known as the shooting method in trajectory optimization .

If the time horizon $T$ is too long, the learned model might deviate from the ground truth due to accumulated prediction errors. Hence, we use Model-Predictive Control (MPC) to stabilize the trajectory by doing forward simulation at every time step as a way to compensate the simulation error.

In many situations, inherent attributes such as masses, friction, and damping are not directly observable. Instead, we can interact with the objects and use PropNet to estimate these attributes online (denoted as $A$) with SGD updates by minimizing the difference between the predicted future states and the actual future states $\mathcal{L}_{\text{state}}(\hat{G}_{t}^{\tau},G_{t}^{\tau})$.

V Experiments

In this section, we evaluate the performance of our model on both simulation and control in three scenarios: Newton’s Cradle, Rope Manipulation, and Box Pushing. We also test how the model generalizes to new scenarios and how it learns to adapt online.

We aim to predict the future states of physical systems. We first describe the network used across tasks and then present the setup of each task as well as the experimental results.

For the IN baseline, we use the same network as described in the original work . For Vanilla PropNet, we adopt similar network structure where the relation propagator $f_{R}^{l}(\cdot)(1\leq l\leq L)$ is an MLP with four 150-dim hidden layers and the object propagator $f_{O}^{l}(\cdot)(1\leq l\leq L-1)$ has one 100-dim hidden layer. Both output a 100-dim propagation vector. For fully observable scenarios, $f^{L}_{O}(\cdot)$ has one 100-dim hidden layer and outputs a 2-dim vector representing the velocity at the next time step. For partially observable cases, $f^{L}_{O}(\cdot)$ outputs one 100-dim vector as the latent representation.

For PropNet, we use an MLP with three 150-dim hidden layers as the relation encoder $f_{R}^{\text{enc}}(\cdot)$ and one 100-dim hidden layer MLP as the object encoder $f_{O}^{\text{enc}}(\cdot)$. Light-weight neural networks are used for the propagators $f_{O}^{l}(\cdot)$ and $f_{R}^{l}(\cdot)$, both of which only contain one 100-dim hidden layer.

A typical Newton’s cradle consists of a series of identically sized rigid balls suspended from a frame. When one ball at the end is lifted and released, it strikes the stationary balls. Forces will transmit through the stationary balls and push the last ball upward immediately. In our fully observable setup, the graph $G$ of $n$ balls has $2n$ objects representing the balls and the corresponding fixed pinpoints above the balls, as shown in Fig. 2a, where $n=5$. There will be $2n$ directed relations describing the rigid connections between the fixed points and the balls. Collisions between adjacent balls introduce another $2(n-1)$ relations.

We generated 2,000 rollouts over 1,000 time steps, of which 85% of the rollouts are randomly chosen as the training set, while the rest are held as the validation set. The model was trained with a mini-batch of 32 using Adam optimizer with an initial learning rate of 1e-3. We reduce the learning rate by 0.8 each time the validation error stops decreasing for over 20 epochs.

Fig. 2a-c show some qualitative results, where we compare IN and PropNet. IN cannot propagate the forces properly: the rightmost ball starts to swing up before the first collision happens. Quantitative results also show that our method significantly outperforms IN in tracking object positions. For 1,000 forward steps, IN results in an MSE of 336.46, whereas PropNet achieves an MSE of 7.85.

We then manipulate a particle-based rope in a 2D plane using a spring-mass model, where one end of the rope is fixed to a random point near the center and the rest of the rope is free to move. Two circular obstacles are placed at random positions near the rope and are fixed to the ground. Random forces are applied to the masses on the rope and the rope is moving in compliant with the forces. More specifically, for a rope containing $n$ particles, there will be a total of $n+2$ objects. Each pair of adjacent masses will have spring relations connecting each other, resulting in $2(n-1)$ directed edges in the dynamics graph $G$. Each mass will have a collision relation with each fixed obstacle, which adds to the graph another $4n$ edges. Frictional force applied to each mass is modeled as a directed edge connecting the mass itself.

We use the same network as described above and generate 5,000 rollouts over 100 time steps. Fig. 5(a) and Fig. 6(a) show qualitative and quantitative results, respectively. We train the models with a 15-dim rope and evaluated in situations where the rope length can vary between 10 and 20. As can be seen from the figures, although the length of the underlying force propagation is fewer than Newton’s Cradle’s, our proposed method can still track the ground truth much more accurately and outperform IN by a large margin.

In this case, we are pushing a pile of boxes forward (Fig. 5(c)). We place a camera at the top of the scene, and only red boxes are observable. More challengingly, the observable boxes are not tracked. Therefore, the visibility of a specific box might change over time. The vertices in the graph are then defined as the state of the observable boxes and edges are defined as directional relations connecting every pair of observable boxes. Specifically, if there are $n$ observable boxes, $n(n-1)$ edges are automatically generated. The dynamics function $\phi(\cdot)$ then takes both the scene representation and the action (i.e., position and velocity of the pusher) as input to perform an implicit forward simulation. As it is hard to explicitly evaluate a latent dynamics model, we evaluate the downstream control tasks instead.

We also provide ablation studies on how the number of propagation steps $L$ influences the final performance. Empirically, a larger $L$ can model a longer propagation path. They are however harder to train and more likely to overfit the training set, often leading to poor generalization. Fig. 6(a) and 6(b) show the ablation studies regarding the choice of $L$. PropNet achieves a high accuracy at $L=3$, with a good trade-off between speed and accuracy. Vanilla PropNet achieves its best accuracy at $L=2$ but generalizes less well as $L$ increases further. This shows the benefits of using the shared encoding and residual connections used in PropNet, as described in Section III-B2.

V-B Control

We now evaluate the applicability of the learned model on control tasks. We first describe the three tasks: Newton’s Cradle, Rope Manipulation, and Box Pushing, which include both open-loop and feedback continuous control tasks, as well as fully and partially observable environments. We evaluate the performance against various baselines and test its ability on generalization and online adaptation.

In this scenario, we assume full-state observation and a control task would be to determine the initial angle of the left-most ball, so as to let the right-most ball achieve a specific height, which can be solved with an accurate forward simulation model.

This is an open-loop control task where we only have control over the initial condition. We thus use a simplified version of Alg. 1. Given the initial physics graph and a learned dynamics model, we iteratively do forward simulation and update the control inputs by minimizing the loss function $\mathcal{L}_{\text{goal}}(\mathcal{G},\mathcal{G}_{g})$. In this specific task, the loss $\mathcal{L}_{\text{goal}}$ is the $\mathcal{L}_{2}$ distance between the target height of the right-most ball and the highest height that has been achieved in $\mathcal{G}$.

We initialize the swing up angle as $45^{\circ}$ and then optimize the angle with a learning rate of $0.1$ for $50$ iterations using Adam optimizer. We compare our model with IN. Qualitative results are shown in Fig. 2e. Quantitatively, PropNet’s output angle has an MSE of 3.08 from the ground truth initial angle, while the MSE for interaction nets is 296.66.

Here we define the task as to move the rope to a target configuration, where the only controls are the top two masses at the moving end of the rope (Fig. 5(b)). The controller tries to match the target configuration by “swinging” the rope, which requires to leverage the dynamics of the rope. The loss $\mathcal{L}_{\text{goal}}$ here is the $\mathcal{L}_{2}$ distance between the resulting configuration and the goal configuration.

We first assume the attributes of the physics graph is known (e.g., mass, friction, damping) and compare the performance between Proportional-Derivative controller (PD) , Model-free Deep Reinforcement Learning (Actor-Critic method optimized with PPO - DRL), as well as Interaction Networks (IN) and Propagation Networks (PropNet) with Alg. 1. Fig. 7 shows quantitative results, where bars marked as “Normal” are the results in this task (a hand-tuned PD controller has an MSE of 2.50). PropNet outperforms the competing baselines. Fig. 5(b) shows a qualitative sample. Compared with the PD controller, our method leverages the dynamics and manages to match the target, instead of naively matching the free end of the rope.

We then consider situations where some of the attributes are unknown and can only be guessed before actually interacting with the objects. We randomly add noise of 15% of the original scale to the attributes as the initial guesses. The “Bias” bars in Fig. 7 show that models trained with ground-truth attributes will encounter performance drop when the supplied attributes are not accurate. However, model-based methods can do online adaptation using the actual output from the environment as feedback to correct the attribute estimation. By updating the estimated attributes over the first $20$ steps of the time horizon with standard SGD, we can improve the manipulation performance so as to catch up with the situations where attributes are accurate (bars marked as “Adapt” in Fig. 7).

We further test whether our model generalizes to new scenarios, where the length of the rope is varied between 10 to 20. As can be seen in Fig. 7, our proposed method can still achieve a good performance, even though the original PropNet is only trained in situations with a fixed length 15 (PD has an MSE of 2.72 for generalization).

In this case, we aim to push a pile of boxes to a target configuration within a predefined time horizon (Fig. 5(c)). We assume partial observation where a camera is placed at the top of the scene, and we can only observe the states of the boxes marked in red. The model trained with partial observation is compared with two baselines: DRL and IN. The loss function $\mathcal{L}_{\text{goal}}$ used for MPC is the $\mathcal{L}_{2}$ distance between the resulting scene encoding and the target scene encoding.

We evaluate the performance by the Chamfer Distance (CD) between the observable boxes at the end of the episode and the target configurations, where for each box in each set, CD finds the nearest box in the other set, and sums the distance up. The negative of the distance is used as the reward for DRL. Fig. 5(c) and Fig. 7(b) show qualitative and quantitative results, respectively. Our method outperforms the baselines due to its explicit modeling of the dynamics and its ability to handle multi-object interactions.

VI Conclusion

We have presented propagation networks (PropNet), a general learnable physics engine that outperforms the previous state-of-the-art with a large margin. We have also demonstrated PropNet’s applicability in model-based control under both fully and partially observable environments. With propagation steps, PropNet can propagate the effects along relations and model the dynamics of long-range interactions within a single time step. We have also proposed to improve PropNet’s efficiency by adding residual connections and shared encoding.

Acknowledgement

This work was supported by: Draper Laboratory Incorporated, Sponsor Award No. SC001-0000001002; NASA - Johnson Space Center, Sponsor Award No. NNX16AC49A; NSF #1524817; DARPA XAI program FA8750-18-C000; ONR MURI N00014-16-1-2007; and Facebook.

References