Year: 2021

Abstract:
In the rapidly advancing area of learned physical simulators, nearly all methods train a forward model that directly predicts future states from input states. However, many traditional simulation engines use a constraint-based approach instead of direct prediction. Here we present a framework for constraint-based learned simulation, where a scalar constraint function is implemented as a trainable function approximator, and future predictions are computed as the solutions to a constraint satisfaction problem. We implement our method using a graph neural network as the constraint function and gradient descent as the constraint solver. The architecture can be trained by standard backpropagation. We test the model on a variety of challenging physical domains, including simulated ropes, bouncing balls, colliding irregular shapes and splashing fluids. Our model achieves better or comparable performance to top learned simulators. A key advantage of our model is the ability to generalize to more solver iterations at test time to improve the simulation accuracy. We also show how hand-designed constraints can be added at test time to satisfy objectives which were not present in the training data, which is not possible with forward approaches. Our constraint-based framework is applicable to any setting in which forward learned simulators are used, and more generally demonstrates key ways that learned models can leverage popular methods in numerical methods.

text: In the rapidly advancing area of learned physical simulators, nearly all methods train a forward model that directly predicts future states from input states. However, many traditional simulation engines use a constraint-based approach instead of direct prediction. Here we present a framework for constraint-based learned simulation, where a scalar constraint function is implemented as a trainable function approximator, and future predictions are computed as the solutions to a constraint satisfaction problem. We implement our method using a graph neural network as the constraint function and gradient descent as the constraint solver. The architecture can be trained by standard backpropagation. We test the model on a variety of challenging physical domains, including simulated ropes, bouncing balls, colliding irregular shapes and splashing fluids. Our model achieves better or comparable performance to top learned simulators. A key advantage of our model is the ability to generalize to more solver iterations at test time to improve the simulation accuracy. We also show how hand-designed constraints can be added at test time to satisfy objectives which were not present in the training data, which is not possible with forward approaches. Our constraint-based framework is applicable to any setting in which forward learned simulators are used, and more generally demonstrates key ways that learned models can leverage popular methods in numerical methods.

heading: 1 INTRODUCTION
text: Consider a bowling ball colliding with a bowling pin. You might explain this event as involving a pair of forces being generated, one which causes the pin to move, and the other which causes the ball to careen away with a different direction and speed. This kind of intuitive cause-and-effect approach is analogous to physical simulators that apply an explicit forward model to calculate a future state directly from the current one, such as when numerically integrating discretized equations of motion.
An alternative, but equally valid, way to explain the collision is in terms of constraint satisfaction: the ball and pin cannot occupy the same location at the same time, and their combined energies and momenta must be conserved, so the post-collision trajectories are the only way the future can unfold without violating these constraints. This constraint-based approach is analogous to physical simulators that use an implicit function to model a system of constraints over the current and future states, and which generate a prediction by searching for a future state that respects all constraints.
Both families of simulators—those based on explicit, forward functions versus those which define the dynamics implicitly, via constraints—are widely used in physics, engineering, and graphics. In principle they can model the same types of dynamics, however they differ in how their respective predictions are computed and in practice strike different trade-offs that determine why one or the other is preferred in different domains. For example, explicit methods are popular for large systems with (mostly) independent local effects whose space and time derivatives are relatively smooth, and their accuracy can often be increased by discretizing space and time more finely. Implicit approaches are often preferred for systems with strong interactions, such rigid and stiff dynamics, and more accurate solutions can often be found by using more sophisticated constraint solvers or by increasing the computational budget (e.g., solver iterations) allocated to searching for solutions. In machine learning (ML), there have been rapid advances recently in methods for learning to simulate complex dynamic processes, however almost all (e.g., Sanchez-Gonzalez et al. (2020); Pfaff et al. (2021)) have focused on explicit forward model approaches, with few exceptions (Yang et al., 2020).
Here we present a framework for learning to simulate complex dynamics via constraint satisfaction. Our “Constraint-based Graph Network Simulator” (C-GNS) defines a single scalar-valued constraint function that represents whether a future state satisfies the physical constraints, conditioned on the current and previous states. The constraint function is implemented as a Graph Neural Network (GNN) (Bronstein et al., 2017; Battaglia et al., 2018), which can model systems with rich compositional structure—multiple bodies, complex meshes, etc. To predict the next state via the constraint function’s implicit representation of the dynamics, a gradient-based solver finds a proposed state which satisfies the constraints. We train it through the solver by backpropagation. We also introduce a hybrid approach that proposes and refines the future state using an explicit iterative predictor, rather than solving for learned constraints.
We tested the C-GNS on a variety of challenging physical simulation domains generated by several different simulation engines: simulated rope, bouncing balls, and bouncing irregular rigid shapes (MuJoCo, Todorov et al. (2012)) and splashing fluids (Flex, Macklin et al. (2014)). We found that the C-GNS’s simulated rollouts were more accurate than a state-of-the-art Graph Net Simulator (GNS) (Sanchez-Gonzalez et al., 2020) with comparable number of parameters. At test time, the C-GNS could use additional solver iterations to improve its predictive accuracy, striking desired speed-accuracy trade-offs. It could also satisfy new, hand-designed constraints jointly alongside its learned constraints. Neither of these capabilities are possible in explicit forward-style approaches.

heading: 2 BACKGROUND AND RELATED WORK
text: Constraint solvers are central to many physics simulators. Most rigid-body and game engines use constraints to model joints, collision and contact (Baraff, 1994). They are used for limiting strain in realistic cloth simulation (Thomaszewski et al., 2009), and are a core component in Eulerian incompressible fluid solvers to solve for pressure (Chorin, 1967). Recently, position-based (Müller et al., 2007) and projective dynamics methods (Bouaziz et al., 2014) have become very popular for interactive simulation. These methods express dynamics purely as constraints, and can simulate a wide range of physical systems from rigids over soft-bodies to fluids (Macklin et al., 2014).
Machine learning methods for accelerating scientific simulation of complex systems, such as turbulence (Kochkov et al., 2021; Wang et al., 2020) and aerodynamics (Thuerey et al., 2020; Zhang et al., 2018), have grown rapidly in recent years. GNN-based learned simulators, in particular, are a very flexible approach which can model a wide range of systems, from articulated dynamics (Sanchez-Gonzalez et al., 2018) to particle-based physics (Mrowca et al., 2018; Li et al., 2019; Sanchez-Gonzalez et al., 2020) and mesh-based continuum systems (Pfaff et al., 2021; De Avila Belbute-Peres et al., 2020), and generalize well to unseen scenarios. Combining learning algorithms with principles from physics and numerical methods, such as auxiliary loss terms and rich inductive biases, can improve sample complexity, computational efficiency, and generalization (Wu et al., 2018; Karniadakis et al., 2021; Chen et al., 2018; Rubanova et al., 2019). Imposing Hamiltonian (Greydanus et al., 2019; Sanchez-Gonzalez et al., 2019; Chen et al., 2019) and Lagrangian (Lutter et al., 2019; Cranmer et al., 2020; Finzi et al., 2020) mechanics in learned simulators offers unique speed/accuracy tradeoffs and can preserve symmetries more effectively.
Recent methods have been proposed for learning constraint functions and solving them in a model’s forward pass (Duvenaud et al. (2020)’s “Deep Implicit Layers” tutorial is an excellent hands-on survey). Such models can play games (Amos & Kolter, 2017; Wang et al., 2019), optimize power flow (Donti et al., 2021), support robotic planning (Loula et al., 2020), and perform combinatorial optimization (Bartunov et al., 2020). Solvers such as gradient descent and Newton’s method are differentiable, and support training by backpropagation, but this can be computationally expensive, so approaches such as Deep Equilibrium Models (DEM) (Bai et al., 2019; 2020) use implicit differentiation to compute gradients only at the solution point.
Despite the popularity of constraint-based traditional simulators, only a single simulator which uses learned constraints has been reported (Yang et al., 2020). Their “Neural Projections” method, based on Goldenthal et al. (2007), iteratively proposes a future state with an Euler step, then projects the proposal onto a learned constraint manifold, implemented as a multilayer perceptron (MLP). Crucially, their constraint function only measures how much an individual state violates the learned constraints, and thus is not an implicit representation of the dynamics. It is suitable for quasi-static regimes, but not scenarios such as the elastic collisions in the bowling ball example described above.
Under review as a conference paper at ICLR 2021
<latexit sha1_base64="pfoxs8rsC/moP79HBUVBdU/MhtI=">AAAB6nicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkUI9FLx4r2g9oQ9lsJ+3SzSbsboQS+hO8eFDEq7/Im//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ7dzvPKHSPJaPZpqgH9GR5CFn1FjpoTvwBuWKW3UXIOvEy0kFcjQH5a/+MGZphNIwQbXueW5i/Iwqw5nAWamfakwom9AR9iyVNELtZ4tTZ+TCKkMSxsqWNGSh/p7IaKT1NApsZ0TNWK96c/E/r5ea8NrPuExSg5ItF4WpICYm87/JkCtkRkwtoUxxeythY6ooMzadkg3BW315nbSvql6tWruvVRo3eRxFOINzuAQP6tCAO2hCCxiM4Ble4c0Rzovz7nwsWwtOPnMKf+B8/gDdJ42H</latexit>
<latexit sha1_base64="t5Z8j6uw1dDrul9BiPoeucvIxm0=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipqQflilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4A4OuM/g==</latexit>
<latexit sha1_base64="HJ/awZ5XFr1cDfruVAeJQcOIixk=">AAAB8HicbVDLSgNBEOz1GeMr6tHLYBA8hV0J6DHoxWME85BkCbOTSTJkZnaZ6RXCkq/w4kERr36ON//GSbIHTSxoKKq66e6KEiks+v63t7a+sbm1Xdgp7u7tHxyWjo6bNk4N4w0Wy9i0I2q5FJo3UKDk7cRwqiLJW9H4dua3nrixItYPOEl4qOhQi4FgFJ302B1RzNrTHvZKZb/iz0FWSZCTMuSo90pf3X7MUsU1Mkmt7QR+gmFGDQom+bTYTS1PKBvTIe84qqniNszmB0/JuVP6ZBAbVxrJXP09kVFl7URFrlNRHNllbyb+53VSHFyHmdBJilyzxaJBKgnGZPY96QvDGcqJI5QZ4W4lbEQNZegyKroQguWXV0nzshJUK9X7arl2k8dRgFM4gwsI4ApqcAd1aAADBc/wCm+e8V68d+9j0brm5TMn8Afe5w8R35CX</latexit>
<latexit sha1_base64="ajLqfCZov3o6FTnfb54h4xEyGI0=">AAAB9HicbVBNS8NAEJ3Ur1q/qh69LBZBEEoiBT0WvXisYGuhDWWz3bZLN5u4OymUkN/hxYMiXv0x3vw3btsctPXBwOO9GWbmBbEUBl332ymsrW9sbhW3Szu7e/sH5cOjlokSzXiTRTLS7YAaLoXiTRQoeTvWnIaB5I/B+HbmP064NiJSDziNuR/SoRIDwShaye+OKKbtrJfihZf1yhW36s5BVomXkwrkaPTKX91+xJKQK2SSGtPx3Bj9lGoUTPKs1E0Mjykb0yHvWKpoyI2fzo/OyJlV+mQQaVsKyVz9PZHS0JhpGNjOkOLILHsz8T+vk+Dg2k+FihPkii0WDRJJMCKzBEhfaM5QTi2hTAt7K2EjqilDm1PJhuAtv7xKWpdVr1at3dcq9Zs8jiKcwCmcgwdXUIc7aEATGDzBM7zCmzNxXpx352PRWnDymWP4A+fzB7oYkhM=</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="pmdXaQtx/RkbjEEj0JjG94undSA=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2FZpQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMTqPqAaBZfYMtwIvE8U0igQ2AnGNzO/84RK81jemUmCfkSHkoecUWOl5kO/XHGr7hxklXg5qUCORr/81RvELI1QGiao1l3PTYyfUWU4Ezgt9VKNCWVjOsSupZJGqP1sfuiUnFllQMJY2ZKGzNXfExmNtJ5Ege2MqBnpZW8m/ud1UxNe+RmXSWpQssWiMBXExGT2NRlwhcyIiSWUKW5vJWxEFWXGZlOyIXjLL6+S9kXVq1VrzVqlfp3HUYQTOIVz8OAS6nALDWgBA4RneIU359F5cd6dj0VrwclnjuEPnM8fuYOM5A==</latexit>
<latexit sha1_base64="Jxyp7Y07VUyiyNEiRyqcT+XYQDk=">AAAB9HicbVDLSgNBEJz1GeMr6tHLYBA8hV0J6DGoB71FMA9IljA76U2GzD6c6Q2GZb/DiwdFvPox3vwbJ8keNLGgoajqprvLi6XQaNvf1srq2vrGZmGruL2zu7dfOjhs6ihRHBo8kpFqe0yDFCE0UKCEdqyABZ6Elje6nvqtMSgtovABJzG4ARuEwhecoZFcv5d2EZ4wvbnLsl6pbFfsGegycXJSJjnqvdJXtx/xJIAQuWRadxw7RjdlCgWXkBW7iYaY8REbQMfQkAWg3XR2dEZPjdKnfqRMhUhn6u+JlAVaTwLPdAYMh3rRm4r/eZ0E/Us3FWGcIIR8vshPJMWIThOgfaGAo5wYwrgS5lbKh0wxjianognBWXx5mTTPK061Ur2vlmtXeRwFckxOyBlxyAWpkVtSJw3CySN5Jq/kzRpbL9a79TFvXbHymSPyB9bnDziakmY=</latexit>
<latexit sha1_base64="rQRE9g+W9jgLwJeS3fY3jHXSBJo=">AAAB83icbVDLSgNBEJyNrxhfUY9eBoPgKexKQI9BPXiMYB6QXcLspDcZMvtgplcMy/6GFw+KePVnvPk3TpI9aGJBQ1HVTXeXn0ih0ba/rdLa+sbmVnm7srO7t39QPTzq6DhVHNo8lrHq+UyDFBG0UaCEXqKAhb6Erj+5mfndR1BaxNEDThPwQjaKRCA4QyO5wSBzEZ4wu83zQbVm1+056CpxClIjBVqD6pc7jHkaQoRcMq37jp2glzGFgkvIK26qIWF8wkbQNzRiIWgvm9+c0zOjDGkQK1MR0rn6eyJjodbT0DedIcOxXvZm4n9eP8XgystElKQIEV8sClJJMaazAOhQKOAop4YwroS5lfIxU4yjialiQnCWX14lnYu606g37hu15nURR5mckFNyThxySZrkjrRIm3CSkGfySt6s1Hqx3q2PRWvJKmaOyR9Ynz+gO5IT</latexit>
<latexit sha1_base64="kZkYF8749t+4eNJpcL8A/9X3kSM=">AAAB83icbVBNS8NAEJ3Ur1q/qh69BIvgqSRS0GOxF48V7Ac0oWy2m3bpZhN2J2IJ+RtePCji1T/jzX/jts1BWx8MPN6bYWZekAiu0XG+rdLG5tb2Tnm3srd/cHhUPT7p6jhVlHVoLGLVD4hmgkvWQY6C9RPFSBQI1gumrbnfe2RK81g+4CxhfkTGkoecEjSSFw4zD9kTZq08H1ZrTt1ZwF4nbkFqUKA9rH55o5imEZNIBdF64DoJ+hlRyKlgecVLNUsInZIxGxgqScS0ny1uzu0Lo4zsMFamJNoL9fdERiKtZ1FgOiOCE73qzcX/vEGK4Y2fcZmkyCRdLgpTYWNszwOwR1wximJmCKGKm1ttOiGKUDQxVUwI7urL66R7VXcb9cZ9o9a8LeIowxmcwyW4cA1NuIM2dIBCAs/wCm9War1Y79bHsrVkFTOn8AfW5w+etZIS</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="FgDtDkQujOEF3ZM1C6wBi95/Q5Y=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPspJMMmZ1dZmaFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glhwbVz328mtrW9sbuW3Czu7e/sHxcOjpo4SxbDBIhGpdkA1Ci6xYbgR2I4V0jAQ2ArGNzO/9YRK80jem0mMfkiHkg84o8ZKrYfHtMzPp71iya24c5BV4mWkBBnqveJXtx+xJERpmKBadzw3Nn5KleFM4LTQTTTGlI3pEDuWShqi9tP5uVNyZpU+GUTKljRkrv6eSGmo9SQMbGdIzUgvezPxP6+TmMGVn3IZJwYlWywaJIKYiMx+J32ukBkxsYQyxe2thI2ooszYhAo2BG/55VXSvKh41Ur1rlqqXWdx5OEETqEMHlxCDW6hDg1gMIZneIU3J3ZenHfnY9Gac7KZY/gD5/MHvsCPMA==</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="FgDtDkQujOEF3ZM1C6wBi95/Q5Y=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPspJMMmZ1dZmaFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glhwbVz328mtrW9sbuW3Czu7e/sHxcOjpo4SxbDBIhGpdkA1Ci6xYbgR2I4V0jAQ2ArGNzO/9YRK80jem0mMfkiHkg84o8ZKrYfHtMzPp71iya24c5BV4mWkBBnqveJXtx+xJERpmKBadzw3Nn5KleFM4LTQTTTGlI3pEDuWShqi9tP5uVNyZpU+GUTKljRkrv6eSGmo9SQMbGdIzUgvezPxP6+TmMGVn3IZJwYlWywaJIKYiMx+J32ukBkxsYQyxe2thI2ooszYhAo2BG/55VXSvKh41Ur1rlqqXWdx5OEETqEMHlxCDW6hDg1gMIZneIU3J3ZenHfnY9Gac7KZY/gD5/MHvsCPMA==</latexit>
<latexit sha1_base64="nk/O4KvoHG7nWgCKsKo+m2fcKfM=">AAAB/nicbVBNS8NAEN3Ur1q/ouLJy2IRPJVECnos9uKxgm2VJoTNdtMu3WzC7kQsoeBf8eJBEa/+Dm/+G7dtDtr6YODx3gwz88JUcA2O822VVlbX1jfKm5Wt7Z3dPXv/oKOTTFHWpolI1F1INBNcsjZwEOwuVYzEoWDdcNSc+t0HpjRP5C2MU+bHZCB5xCkBIwX2kSdJKEhwj6Mg94A9Qt6cTAK76tScGfAycQtSRQVagf3l9ROaxUwCFUTrnuuk4OdEAaeCTSpepllK6IgMWM9QSWKm/Xx2/gSfGqWPo0SZkoBn6u+JnMRaj+PQdMYEhnrRm4r/eb0Moks/5zLNgEk6XxRlAkOCp1ngPleMghgbQqji5lZMh0QRCiaxignBXXx5mXTOa269Vr+pVxtXRRxldIxO0Bly0QVqoGvUQm1EUY6e0St6s56sF+vd+pi3lqxi5hD9gfX5A3C7lc8=</latexit>
<latexit sha1_base64="vQ3/3W0LApHcMPwu9rwFnhCUcHI=">AAAB73icbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ie0oWw2k3bpZhN3N0Ip/RNePCji1b/jzX/jts1BWx8MPN6bYWZekAqujet+O4W19Y3NreJ2aWd3b/+gfHjU0kmmGDZZIhLVCahGwSU2DTcCO6lCGgcC28HoZua3n1Bpnsh7M07Rj+lA8ogzaqzU6YUoDCUP/XLFrbpzkFXi5aQCORr98lcvTFgWozRMUK27npsaf0KV4UzgtNTLNKaUjegAu5ZKGqP2J/N7p+TMKiGJEmVLGjJXf09MaKz1OA5sZ0zNUC97M/E/r5uZ6MqfcJlmBiVbLIoyQUxCZs+TkCtkRowtoUxxeythQ6ooMzaikg3BW355lbQuql6tWrurVerXeRxFOIFTOAcPLqEOt9CAJjAQ8Ayv8OY8Oi/Ou/OxaC04+cwx/IHz+QOaD4+w</latexit>
<latexit sha1_base64="fAxixLzqPP0ibRTKMP9FMpa2phA=">AAACB3icbZDLSsNAFIZPvNZ6i7oUZLAIlUJJpKAboejGZQV7o41lMpm0QycXZiZCCd258VXcuFDEra/gzrdxmnahrT8MfPznHM6c3405k8qyvo2l5ZXVtfXcRn5za3tn19zbb8goEYTWScQj0XKxpJyFtK6Y4rQVC4oDl9OmO7ye1JsPVEgWhXdqFFMnwP2Q+Yxgpa2eedS+T4usZJ+O0SXKWFMJdT3KFUbtnlmwylYmtAj2DAowU61nfnW9iCQBDRXhWMqObcXKSbFQjHA6zncTSWNMhrhPOxpDHFDppNkdY3SiHQ/5kdAvVChzf0+kOJByFLi6M8BqIOdrE/O/WidR/oWTsjBOFA3JdJGfcKQiNAkFeUxQovhIAyaC6b8iMsACE6Wjy+sQ7PmTF6FxVrYr5cptpVC9msWRg0M4hiLYcA5VuIEa1IHAIzzDK7wZT8aL8W58TFuXjNnMAfyR8fkD8P+W0w==</latexit>
<latexit sha1_base64="fAxixLzqPP0ibRTKMP9FMpa2phA=">AAACB3icbZDLSsNAFIZPvNZ6i7oUZLAIlUJJpKAboejGZQV7o41lMpm0QycXZiZCCd258VXcuFDEra/gzrdxmnahrT8MfPznHM6c3405k8qyvo2l5ZXVtfXcRn5za3tn19zbb8goEYTWScQj0XKxpJyFtK6Y4rQVC4oDl9OmO7ye1JsPVEgWhXdqFFMnwP2Q+Yxgpa2eedS+T4usZJ+O0SXKWFMJdT3KFUbtnlmwylYmtAj2DAowU61nfnW9iCQBDRXhWMqObcXKSbFQjHA6zncTSWNMhrhPOxpDHFDppNkdY3SiHQ/5kdAvVChzf0+kOJByFLi6M8BqIOdrE/O/WidR/oWTsjBOFA3JdJGfcKQiNAkFeUxQovhIAyaC6b8iMsACE6Wjy+sQ7PmTF6FxVrYr5cptpVC9msWRg0M4hiLYcA5VuIEa1IHAIzzDK7wZT8aL8W58TFuXjNnMAfyR8fkD8P+W0w==</latexit>
<latexit sha1_base64="9UY9Tzo0xSguDEfPuYbWHGsRnh0=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPsZJIMmZ1dZnqFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glgKg6777eTW1jc2t/LbhZ3dvf2D4uFR00SJZrzBIhnpdkANl0LxBgqUvB1rTsNA8lYwvpn5rSeujYjUPU5i7od0qMRAMIpWaj08pmX3fNorltyKOwdZJV5GSpCh3it+dfsRS0KukElqTMdzY/RTqlEwyaeFbmJ4TNmYDnnHUkVDbvx0fu6UnFmlTwaRtqWQzNXfEykNjZmEge0MKY7MsjcT//M6CQ6u/FSoOEGu2GLRIJEEIzL7nfSF5gzlxBLKtLC3EjaimjK0CRVsCN7yy6ukeVHxqpXqXbVUu87iyMMJnEIZPLiEGtxCHRrAYAzP8ApvTuy8OO/Ox6I152Qzx/AHzucPZ+qO9w==</latexit>
<latexit sha1_base64="4l4tyu1wQNISUpBmeK1HI2A9PJA=">AAAB7nicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkUI9FLx4r2A9pQ9lsN+3SzSbsToQS+iO8eFDEq7/Hm//GbZuDtj4YeLw3w8y8IJHCoOt+O4WNza3tneJuaW//4PCofHzSNnGqGW+xWMa6G1DDpVC8hQIl7yaa0yiQvBNMbud+54lrI2L1gNOE+xEdKREKRtFKnf6YYvY4G5QrbtVdgKwTLycVyNEclL/6w5ilEVfIJDWm57kJ+hnVKJjks1I/NTyhbEJHvGepohE3frY4d0YurDIkYaxtKSQL9fdERiNjplFgOyOKY7PqzcX/vF6K4bWfCZWkyBVbLgpTSTAm89/JUGjOUE4toUwLeythY6opQ5tQyYbgrb68TtpXVa9Wrd3XKo2bPI4inME5XIIHdWjAHTShBQwm8Ayv8OYkzovz7nwsWwtOPnMKf+B8/gCDZ4+x</latexit>
<latexit sha1_base64="ErhthiE4wcyOR5LhKX4LnvmTkNM=">AAAB+HicbVBNS8NAEN34WetHox69LBahXkoiBb0IRS+epIL9oo1ls922SzebsDsRasgv8eJBEa/+FG/+G7dtDtr6YODx3gwz8/xIcA2O822trK6tb2zmtvLbO7t7BXv/oKHDWFFWp6EIVcsnmgkuWR04CNaKFCOBL1jTH19P/eYjU5qH8h4mEfMCMpR8wCkBI/XsQndEIGmnl+2HpHR7mvbsolN2ZsDLxM1IEWWo9eyvbj+kccAkUEG07rhOBF5CFHAqWJrvxppFhI7JkHUMlSRg2ktmh6f4xCh9PAiVKQl4pv6eSEig9STwTWdAYKQXvan4n9eJYXDhJVxGMTBJ54sGscAQ4mkKuM8VoyAmhhCquLkV0xFRhILJKm9CcBdfXiaNs7JbKVfuKsXqVRZHDh2hY1RCLjpHVXSDaqiOKIrRM3pFb9aT9WK9Wx/z1hUrmzlEf2B9/gAd6pK9</latexit>
<latexit sha1_base64="ErhthiE4wcyOR5LhKX4LnvmTkNM=">AAAB+HicbVBNS8NAEN34WetHox69LBahXkoiBb0IRS+epIL9oo1ls922SzebsDsRasgv8eJBEa/+FG/+G7dtDtr6YODx3gwz8/xIcA2O822trK6tb2zmtvLbO7t7BXv/oKHDWFFWp6EIVcsnmgkuWR04CNaKFCOBL1jTH19P/eYjU5qH8h4mEfMCMpR8wCkBI/XsQndEIGmnl+2HpHR7mvbsolN2ZsDLxM1IEWWo9eyvbj+kccAkUEG07rhOBF5CFHAqWJrvxppFhI7JkHUMlSRg2ktmh6f4xCh9PAiVKQl4pv6eSEig9STwTWdAYKQXvan4n9eJYXDhJVxGMTBJ54sGscAQ4mkKuM8VoyAmhhCquLkV0xFRhILJKm9CcBdfXiaNs7JbKVfuKsXqVRZHDh2hY1RCLjpHVXSDaqiOKIrRM3pFb9aT9WK9Wx/z1hUrmzlEf2B9/gAd6pK9</latexit>
<latexit sha1_base64="9UY9Tzo0xSguDEfPuYbWHGsRnh0=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXAnoMevEYwTwkWcPsZJIMmZ1dZnqFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7glgKg6777eTW1jc2t/LbhZ3dvf2D4uFR00SJZrzBIhnpdkANl0LxBgqUvB1rTsNA8lYwvpn5rSeujYjUPU5i7od0qMRAMIpWaj08pmX3fNorltyKOwdZJV5GSpCh3it+dfsRS0KukElqTMdzY/RTqlEwyaeFbmJ4TNmYDnnHUkVDbvx0fu6UnFmlTwaRtqWQzNXfEykNjZmEge0MKY7MsjcT//M6CQ6u/FSoOEGu2GLRIJEEIzL7nfSF5gzlxBLKtLC3EjaimjK0CRVsCN7yy6ukeVHxqpXqXbVUu87iyMMJnEIZPLiEGtxCHRrAYAzP8ApvTuy8OO/Ox6I152Qzx/AHzucPZ+qO9w==</latexit>
<latexit sha1_base64="pmdXaQtx/RkbjEEj0JjG94undSA=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2FZpQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMTqPqAaBZfYMtwIvE8U0igQ2AnGNzO/84RK81jemUmCfkSHkoecUWOl5kO/XHGr7hxklXg5qUCORr/81RvELI1QGiao1l3PTYyfUWU4Ezgt9VKNCWVjOsSupZJGqP1sfuiUnFllQMJY2ZKGzNXfExmNtJ5Ege2MqBnpZW8m/ud1UxNe+RmXSWpQssWiMBXExGT2NRlwhcyIiSWUKW5vJWxEFWXGZlOyIXjLL6+S9kXVq1VrzVqlfp3HUYQTOIVz8OAS6nALDWgBA4RneIU359F5cd6dj0VrwclnjuEPnM8fuYOM5A==</latexit>
<latexit sha1_base64="YET3kkOS2zY8mMv2K4bjQDXERhw=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipyQblilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4AyKuM7g==</latexit>
<latexit sha1_base64="e4/cShLnpc3G1u2go1/xMgSgs0s=">AAAB8XicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48V7Ae2oWy2m3bpZhN3J0IJ/RdePCji1X/jzX/jts1BWx8MPN6bYWZekEhh0HW/ncLa+sbmVnG7tLO7t39QPjxqmTjVjDdZLGPdCajhUijeRIGSdxLNaRRI3g7GNzO//cS1EbG6x0nC/YgOlQgFo2ilh04/60n+SHDaL1fcqjsHWSVeTiqQo9Evf/UGMUsjrpBJakzXcxP0M6pRMMmnpV5qeELZmA5511JFI278bH7xlJxZZUDCWNtSSObq74mMRsZMosB2RhRHZtmbif953RTDKz8TKkmRK7ZYFKaSYExm75OB0JyhnFhCmRb2VsJGVFOGNqSSDcFbfnmVtC6qXq1au6tV6td5HEU4gVM4Bw8uoQ630IAmMFDwDK/w5hjnxXl3PhatBSefOYY/cD5/AHGkkMY=</latexit>
<latexit sha1_base64="t5Z8j6uw1dDrul9BiPoeucvIxm0=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkoMeiF48t2A9oQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgobm1vbO8Xd0t7+weFR+fikreNUMWyxWMSqG1CNgktsGW4EdhOFNAoEdoLJ3dzvPKHSPJYPZpqgH9GR5CFn1FipqQflilt1FyDrxMtJBXI0BuWv/jBmaYTSMEG17nluYvyMKsOZwFmpn2pMKJvQEfYslTRC7WeLQ2fkwipDEsbKljRkof6eyGik9TQKbGdEzVivenPxP6+XmvDGz7hMUoOSLReFqSAmJvOvyZArZEZMLaFMcXsrYWOqKDM2m5INwVt9eZ20r6perVpr1ir12zyOIpzBOVyCB9dQh3toQAsYIDzDK7w5j86L8+58LFsLTj5zCn/gfP4A4OuM/g==</latexit>
<latexit sha1_base64="NkuDH4GMVAhCg8VuQG4AALMi2Sw=">AAACNHicbVBNSxxBEO3RxJhNNKsec2myBFaIy0xYMBdB9BLIxUBWd9jZDDU9NdrY0zN214jLMD/Kiz8kFwnkEJFc8xvSu+4hfhQ0vH6vXnXXS0olLfn+T29h8dnzpRfLL1uvXq+svmmvrR/aojICB6JQhRkmYFFJjQOSpHBYGoQ8UXiUnO5P9aNzNFYW+htNShzncKxlJgWQo+L2lyhFRcBDvsO3IuWMKfBIQ6IgDnkW1xHhBdX7TdMduovCM07NBx5u8si4uRTX4U74ve7KzaaJ2x2/58+KPwbBHHTYvA7i9o8oLUSVoyahwNpR4Jc0rsGQFAqbVlRZLEGcwjGOHNSQox3Xs6Ub/t4xKc8K444mPmP/d9SQWzvJE9eZA53Yh9qUfEobVZR9GtdSlxWhFncPZZXiVPBpgjyVBgWpiQMgjHR/5eIEDAhyObdcCMHDlR+Dw4+9oN/rf+13dvfmcSyzt+wd67KAbbNd9pkdsAET7JJds9/sxrvyfnm33p+71gVv7tlg98r7+w94eaot</latexit>

heading: 3 MODEL FRAMEWORK
text: Simulation basics A physical trajectory, measured at discrete time intervals, is a sequence of states, (X1, . . . , XT ), where Xt represents properties such as the positions, velocities, masses, etc, of elements of the system. A physical simulator, s, is a function that maps current and/or previous state(s), which we term the context, X≤t, to a predicted future state, X̂t+1 = s(X≤t) (see Figure 1a)1. A simulated physical trajectory, termed a rollout, (Xt, X̂t+1, X̂t+2, . . . ), can be generated by repeatedly applying s to its own predicted state, X̂t+1 = s(X̂≤t).
Simulators are often comprised of a PREDICTOR mechanism which maps the context X≤t to an update value Ŷ , that represents information about the system’s temporal evolution at the current time. Then Ŷ is used by an UPDATER mechanism to update the current state to the next state: X̂t+1 = UPDATER(X≤t, Ŷ ), e.g., updating current positions and velocities represented by Xt with new velocities and accelerations represented by Ŷ , to predict the next state.
Explicit simulators Across science, engineering, and graphics, a popular class of simulators are defined explicitly: the state update Ŷ is predicted directly fromX≤t using an explicit forward function, Ŷ = fD(X≤t), as illustrated in Figure 1b. Among the rapidly growing family of learned simulators, the forward function fD is typically implemented using a neural network (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021).
Constraint-based implicit simulators Here we explore learned simulators based on implicit formulations of the dynamics. Rather than predicting the desired state directly, as in explicit formulations, our implicit simulator uses a differentiable constraint function, c = fC(X≤t, Ŷ ), where c is a scalar that quantifies how well a proposed state update Ŷ agrees withX≤t. A future prediction is generated by applying a solver, such as an optimization or zero-finding algorithm, to find a Ŷ that satisfies the constraint function, and applying the UPDATER to update Xt to X̂t+1. The fC can represent all the physical constraints in the system, including the time dynamics.
1Despite that physics is Markovian, we use X≤t as input because our framework can also apply to dynamic processes which are non-Markovian. Providing previous states can also often be helpful when there are hidden properties of the system which are only identifiable over a sequence of observed states, and when a state does not represent velocity or momentum information.
As illustrated in Figure 1d, we formulate our constraint-solving procedure via an iterative method that starts with an initial proposal, Y (0). On the i-th iteration, the solver uses the gradient of fC w.r.t. Y at the current proposal to compute a change to the proposal, δY = −λ ∇Y fC(X≤t, Y )|Y=Y (i) . This δY is then used to revise the proposal to, Y (i+1) = Y (i) + δY . This process repeats for N steps, and the final proposal value is treated as the PREDICTOR’s output, Ŷ = Y (N).
Our constraint-based model’s fC is defined as a trainable function approximator which is real-valued and lower bounded at zero, and uses gradient descent to find Ŷ that minimizes it, where λ is a fixed step size. This induces the semantics that the desired Ŷ = arg minY fC(X≤t, Y ).
We also explore a second constraint-solving procedure, inspired by Yang et al. (2020)’s Neural Projections’ use of “fast projection” Goldenthal et al. (2007). Specifically, λ = − fC(X≤t,Y (i))
‖∇Y fC(X≤t,Y )| Y =Y (i)
‖2 . Unlike gradient descent, fast projection is a zero-finding algorithm, so in
this case fC is not lower bounded. This induces the semantics that fC(X≤t, Ŷ ) = 0.
This general formulation of constraint-based learned simulation can be trained by backpropagating loss gradients through the solver loop2. The computational budget of the forward pass can be varied via the number of solver iterations N .
Explicit iterative simulators As a hybrid between forward and constraint-based simulators, we introduced a model which iteratively refines a proposed state update, like in the constraint-based approach described above, but using an explicit function to directly output a δY at each iteration, rather than solving a constraint function (see Figure 1c). See Section 4.3 for details.

heading: 4 EXPERIMENTS

heading: 4.1 EXPERIMENTAL TASK DOMAINS
text: We test our framework on a variety of physical environments, shown in Figure 2: ROPE, BOUNCING BALLS and BOUNCING RIGIDS, whose ground truth training and test data were generated by the MuJoCo physics simulator, as well as BOXBATH from Li et al. (2019). These environments demonstrate a diverse set of physical constraints: ‘hard’ constraints (preserving the shape of the rigid object and resolving collisions), and ‘soft’ constraints on fluid movement, handling gravity and preserving the momentum of the rope and bouncing balls. See the Supplementary Materials for details.

heading: 4.2 MODEL IMPLEMENTATIONS
text: Representing the physical system Our experimental domains are physical systems comprised of sets of interacting point-like elements, e.g., objects, particles, mesh vertices, etc. We represent the state as Xt = (p j t )
j=1...|Xt|, where |Xt| is the number of elements, and pjt is the j-th element’s position at time t. There are also other static properties of the physical elements, e.g., masses, material types, etc., which we represent with Z to keep it distinct from the dynamic state information represented by Xt. The input context is X≤t = (Z,Xt−3, Xt−2, Xt−1, Xt).
2Implicit differentiation at the solution point should be applicable as well, and potentially offer computational benefits as mentioned in the Section 2, though we do not explore that here.
In our implementation, Ŷ , represents the predicted changes in position (i.e., the “average velocity” across the time step)3, ŷj = ∆p̂jt+1 = p̂ j t+1 − pjt . The UPDATER then computes X̂t+1 using p̂jt+1 = p j t + ∆p̂ j t+1, where p j t is provided in the input X≤t.
Constructing the input graph Our implementations of the fD, fDI, and fC use GNNs as the function approximators, so we need to pack the context, X≤t, and (for the fDI and fC) the proposed state update information, Y (i), into an input graph, Gt = (Vt, Et). The edges Et represent possible interactions among the elements, such as fully connected edges to represent collisions and rigid attachments in BOUNCING BALLS and BOUNCING RIGIDS, spring constraints in ROPE, and interactions among particles within a fixed connectivity radius in BOXBATH.
We enforced translation-invariance by construction, by never providing absolute positions as input to the models. Instead, the j-th input node’s features are the static properties, and a sequence of the three most recent position changes (i.e. average velocities), vjt = [z j ,∆pjt−2,∆p j t−1,∆p j t ], where, ∆pjt = p j t − pjt−1. For fDI and fC, which also take the solver’s current proposed Y (i), we also concatenate the proposed average velocity from the i-th solver iteration, yj,(i) − pjt , as input. For the input edge feature for an edge that connects from node j to k, we also provide the relative displacement vector between the nodes’ positions, ejkt = p k t − pjt .
GNN-based Encode-Process-Decode core We implemented fD, fDI, and fC using Graph Networks (GN) (Battaglia et al., 2018), arranged in the Encode-Process-Decode architecture, similar to previous work on GN-based learned simulators (Sanchez-Gonzalez et al., 2018; 2020; Pfaff et al., 2021). The Encoder uses two MLPs to encode node and edge features into high-dimensional latent vectors. The Processor applies multiple GNs, with unshared weights, in sequence, with node and edge residual connections at each step. We do not use global updates for the GNs. The Decoder uses an MLP to produce an output for each node.
The fD directly returns Ŷ . The fDI returns a change to the proposed update δY for the current iteration. The fC’s Decoder returns a scalar for each node to produce a constraint value per node {cj |j = 1 . . . |V |}. These node-wise constraint values are averaged to compute a single scalar c constraint for the entire system, c = fC(X≤t, Ŷ ) = 1|V | ∑|V | j=1 c j .
Solving the constraint For fDI and fC we initialize Y (0) = ∆pjt to the most recent average velocity4. We used auto-differentiation in JAX to compute the gradient function,∇Y fC, and the step size λwas specific to the model variant, as described below. During training we used N = 5 solver iterations.

heading: 4.3 MODEL VARIANTS
text: The key questions in this work are whether constraint-based learned simulators can compete with explicit, forward learned simulators, whether implementing the constraint function with GNNs is more effective than with MLPs, and how minima-based constraint functions solved by gradient descent compare to constraints defined as the zeros of a function which are solved by fast projection (Goldenthal et al., 2007). The following model variants allow us to answer these questions.
Forward GNN This is an explicit, forward GNN-based learned simulator based on the GNS models from Sanchez-Gonzalez et al. (2020); Pfaff et al. (2021). It directly predicts the state update Ŷ from the past time points X≤t.
C-GNS Gradient Descent (C-GNS-GD) and C-GNS-Fast Projections (C-GNS-FP) These are our proposed constraint-based GNN models. For the C-GNS-GD, the scalar per-node output cj was squared, to force the overall fC to be non-negative, and a gradient descent solver with a fixed step size, λ = 0.001, was used to minimize it. For C-GNS-FP, the λ was based on “fast projection” (Goldenthal et al., 2007; Yang et al., 2020), as described in Section 3. Supplementary Figure B.5(c-d) shows ablations.
3For BOXBATH we vary a number of modelling choices to best match those in Sanchez-Gonzalez et al. (2020). The major difference is that we set Ŷ to be the average acceleration rather than average velocity. See Supplementary Materials for other differences.
4To ensure analogous information is provided downstream of fD, the update rule also includes the previous average velocity: p̂jt+1 = p j t + ∆p j t + ŷ j
Iterative GNN We implemented a hybrid between the Forward GNN and C-GNS, as shown in Figure 1c. It was identical to the C-GNS models, except its fDI directly predicted proposed state updates as in fD, rather than being computed via the gradients as was done with fC.
ConstraintMLP Gradient Descent (ConstraintMLP-GD) and ConstraintMLP-Fast Projections (ConstraintMLP-FP) These were MLP-based constraint models, which, rather than using GNNs to implement fC, instead concatenated the embeddings of all the input nodes into a single vector and passed them to an MLP implementation of fC. By default, these models cannot handle variable-length inputs, so we padded smaller states with zeros up to the maximum state size. The ConstraintMLP-FP was the MLP analog to our C-GNS-FP, and was similar to Neural Projections (Yang et al., 2020). The ConstraintMLP-GD used gradient descent, and was the MLP analog to our C-GNS-GD. We omit the results for the ConstraintMLP models on BOXBATH (1024 nodes), as MLPs do not generally work well on physical systems with more than a few particles (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018).

heading: 4.4 TRAINING AND EVALUATION
text: We trained the models to make next-step predictions, by computing the L2 loss between the predicted X̂t+1 and the corresponding ground truth Xt+1, averaged over nodes. All model weights and biases were trained using standard backpropagation with the Adam optimizer.
At test time, we compute 1-step metrics by evaluating the 1-step errors along each point of the ground truth trajectory. We also evaluate rollout errors by iteratively applying the learned model starting from an initial state, over 160 rollout steps, and computing the error between the predicted and ground truth trajectories.

heading: 5 RESULTS
text: Predictive accuracy5 Our experimental results show that our C-GNS-GD’s performance was generally better than the other model variants. Figure 3 compares the different models on 1-step and rollout position MSE (see Supplementary Table B.1 for numerical results). For each dataset, we
5Videos of the model rollouts are available at sites.google.com/view/constraint-based-simulator
used the same number of message-passing steps (MP) for all GN-based models. We used 2 MPs for the ROPE dataset, and 1 MP for all other tasks.
The C-GNS-GD has lower 1-step MSE between the ground truth and predicted positions than other models across all datasets. Qualitatively, we observed that for Forward GNN with a single messagepassing step, the box in BOXBATH “melts” over time, as the forward model cannot preserve its rigid shape (see Videos). The comparable C-GNS-GD, by contrast, maintains the rigidity more effectively. These quantitative results suggest that constraint-based learned simulators are competitive alternative to explicit, forward learned simulators. We generally found that the Iterative GNN was fairly competitive with the C-GNS-GD in overall performance and better than the Forward GNN.
We also found that the C-GNS-FP was generally less stable across seeds, and not as accurate as the C-GNS-GD. The same conclusion holds for ConstraintMLP-FP versus ConstraintMLP-GD. We speculate that the fast projection algorithm may make training challenging because the step size λ is proportional to fC, which may cause poor zero-finding early in training when the fC is not yet informative. Additionally, we find that C-GNS-FP algorithm becomes unstable in the areas with shallow constraint gradients, perhaps because its λ depends on the inverse of the gradient’s norm.
We explored how varying the message-passing steps and solver iterations (N ) influenced the relative performance among the models in our ROPE dataset. Figure 4 shows that the C-GNS-GD generally required fewer parameters and message-passing steps to achieve comparable 1-step MSE to the other models. Supplementary Figure B.3 shows similar results for the rollout MSE. For most combinations of message-passing steps and number of solver iterations, C-GNS-GD (green) outperforms the Iterative GNN (yellow), C-GNS-FP (purple) as well as the Forward GNN (blue) with the same number of MPs (the Forward GNN is not iterative model, so we plot it as a single bar). We hypothesize that the solver iterations in the C-GNS and Iterative GNN may play a similar role to message passing with shared weights .
Interpreting the learned constraints To better understand the learned fC functions in the C-GNSGD, Figure 5 visualizes the node-wise constraint values as a function of Y (proposed average velocity) for different nodes in the ROPE dataset while holding the other nodes’ proposed update Y fixed.
We also overlay the sequence of five points that represent the proposed Y (i) steps from the solver where all nodes were jointly optimized. The figure shows the learned fC has a minimum near the ground truth Y, which the gradient descent steps are able to reach.
Incorporating novel constraints at test time We next explored a unique advantage of the constraint-based model: because the fC measures the degree the physical constraints are violated, we can incorporate additional, hand-designed constraints at test time, and use the model to potentially satisfy them. For the ROPE dataset, we designed three constraint functions that return positive values which increase quadratically as the rope enters different “forbidden” regions of the space: a vertical wall, a horizontal floor, and a disk-shaped region. We weighted these constraint terms by a coefficient hyperparameter and added each of the hand-designed constraints to the learned fC term of C-GNS-GD and ran the forward evaluation of the model.
As shown in Figure 6, the model was able to simulate the dynamics in a way that the corresponding forbidden region was avoided. In some cases, satisfying the joint constraint resulted in unintuitive behaviors, such as the rope links changing in length to adapt to the obstacle (Videos). However, this is to be expected, as the minimum of the joint constraint may not overlap with the minimum of the learned constraint, which is the one that would otherwise guarantee length preservation. For this example we added a further hand-designed constraint that incentivizes maintaining relative distances between nodes. In general this is a powerful example of how constraint-based models can generalize outside their training data, and solve both for the learned dynamics and arbitrary desired constraints.
Generalizing to larger systems via increased solver iterations In principle, iterative and constraint-based simulators should find more accurate solutions by increasing the number of solver iterations, N . We investigated whether the C-GNS-GD and Iterative GNN trained on ROPE could generalize from Ntrain = 5 on which they were trained, to Ntest ∈ [0, 15]. We also analyzed whether increased solver iterations could improve generalize performance from training on ropes with 5−10 nodes, to test ropes with 20 nodes.
Figure 7a (top row) shows that for test ropes that match the 5−10 nodes experienced during training, the Iterative GNN (light blue) overfits very heavily to Ntest = Ntrain = 5: error increases abruptly for N ≤ 4 and N ≥ 6. By contrast, the C-GNS-GD (light red) generalizes much better to different Ntest. Figure 7a (bottom row) shows that for test ropes with 20 nodes, the Iterative GNN again overfits, while the C-GNS-GD can generalize well to longer ropes if Ntest is increased.
We also trained the Iterative GNN and C-GNS-GD with additional loss terms that were applied to the Y (i) on each solver iteration, not only the final one, Ŷ = Y (N). We used an exponential decay factor, α = 0.25, which downweighted this additional loss term more heavily for earlier solver proposals. The dark blue and red curves in Figure 7a show how this additional loss further improves generalization to more solver iterations and larger systems as test time for the Iterative GNN, but especially the C-GNS-GD. Figure 7b visualizes how increasing the solver iterations systematically improves the quality of the long-term rollout accuracy in the ROPE dataset.
Together these results show the C-GNS-GD is effective in making use of additional resources at test time. This opens the exciting possibility of training on small, simple systems, and testing on large, complex systems. See Supplementary Figure B.2 for further details.

heading: 6 DISCUSSION
text: We presented a general-purpose framework for constraint-based learned simulation, where a learned constraint function implicitly represents the dynamics, and future predictions are generated via a constraint solver. We implemented our framework using GNNs as the constraint function and gradient descent as the constraint solver, and tested it in a variety of challenging physical simulation problems. Our results showed that our C-GNS has competitive or better performance compared to previous learned simulators. We demonstrated unique abilities to generalize to novel, hand-designed constraints, and use more solver iterations than experienced during training to improve the accuracy on larger systems.
We can hypothesize about the relationship between explicit, forward learned simulators and implicit, constraint-based ones in terms of the sharing schemes of these architectures. The C-GNS has a stronger inductive bias than the Forward GNN. The transformation of fC in C-GNS effectively ties the parameters in the resulting ∇Y fC function, and the solver iterations are analogous to how a recurrent neural network’s parameters are shared over iterations. In contrast, the message-passing steps in the Forward GNN used in our work are unshared. In principle, the fD of the Forward GNN is more expressive because if given enough depth, after training it could learn to take parameter values that are equivalent to the shared parameters of C-GNS. Our results shown in Figure 4 supports this possibility: the Forward GNN with many more message-passing steps eventually approaches the C-GNS’s performance. Moreover, we speculate the C-GNS’s inductive biases contribute to its advantages in terms of incorporating novel hand-designed constraints and generalizing to more solver iterations and larger systems.
More broadly, the performance, generality and unique advantages of constraint-based learned simulation make it an important new direction in the advancement of machine learning methods for complex simulation problems in science and engineering.

heading: 7 REPRODUCIBILITY STATEMENT
text: We are committed to open-source the model code after the paper is accepted. Also, we are going to open-source the MuJoCo datasets that we generated for this paper. We provide more details on the model implementation as well as the hyperparameters used for each model in the Supplementary Material.

Review summaries:
This paper aims to add explicit/human-defined constraints to learning-based simulation frameworks, where a learned constraint function implicitly regularizes the dynamics, and future predictions are generated via a constraint solver. The authors built the framework on top of graph neural networks (GNNs) to capture the compositionality of the underlying system and enforce the constraint using an implicit constraint function optimized via gradient descent.  The authors tested the proposed method in four physical simulation environments, including rope, bouncing balls, bouncing rigids, and BoxBath. Experiment results show that the proposed C-GNS has a competitive or better performance compared to prior learned simulators. In addition, they have also demonstrated that C-GNS can generalize to unseen, hand-designed constraints by applying more solver iterations than experienced during training to improve the accuracy on larger systems.

This paper presents a neural-network simulator that learns to solve constraints inspired by classic physics-based simulation. The proposed simulator uses a graph neural network to encode a constraint solver and shows results in a number of simulation environments.  The main contribution of the paper is its idea of using graph networks to encode constraint-based simulation. 

This paper presents to simulate physics via a constraint-based approach instead of direct prediction. In particular, the authors first employ GNN taking as input the history positions and dynamics, to capture the interaction between different particles within the system, whose output is considered as the constraint satisfaction scalar. Then, the gradient with respect to the constraint function is applied as the update of the dynamics over a certain number of solver iterations. The experiments are conducted on a variety of challenging physical domains, including simulated ropes, bouncing balls, colliding irregular shapes and splashing fluids.

This manuscript proposes to learn the numerical solutions to Lagrangian physical simulation with a constraint-based inference method on graph neural networks. This method involves an iterative update during inference and thus enables test-time dynamical correction. The contributions are: (1) this manuscript builds a scalar predictor to indicate how well the constraint is agreed. (2) this manuscript proposes using the graph neural networks as the backbone to deal with a variable length of the physical domain. They also examine the effectiveness with a bunch of experiments including the state prediction experiment on four different environments and multiple ablation studies towards different hyper-parameters.