Differentiable Implicit Soft-Body Physics

Physics engines have been extensively used in robotics and computer graphics, but the increased adoption of neural networks to solve inverse problems in these domains has revealed the need for additional features in such engines (Degrave et al., 2019; Hu et al., 2019c, b, a). For example, using neural networks to solve inverse problems where physical constraints must be satisfied, such as motor control, often requires model-free reinforcement learning algorithms that approximate the derivatives of the simulator by sampling multiple trajectories using stochastic control policies. This is due to the non-differentiable nature of most existing simulators, which have traditionally focused on efficiency and stability of forward computations only. Differentiable simulators are useful when we want to incorporate physics into neural networks as a differentiable layer, which allows us to run backpropagation as commonly done in supervised learning.

In this paper, we focus on soft-body physics and muscle-like actuation mechanisms, motivated by humans and other animals that control their bodies via storage and release of energy in their muscles to generate a wide range of motions, and we demonstrate that our differentiable simulator is effective in optimizing control policies for locomotion (Figure 1) and can achieve better sample efficiency than model-free reinforcement learning. In contrast to other differentiable soft-body physics approaches that use explicit forward models to define state transitions, we focus on implicit state transitions defined via function minimization. Implicit state transitions appear in implicit numerical integration methods, which offer the benefits of large time steps and excellent numerical stability, but require a special treatment to achieve differentiability due to the absence of an explicit differentiable forward pass.

We first consider quasistatic motions, which are primarily driven by potential energy minimization and ignore inertial effects, which we will consider in Section 5. While typical elastic objects are passive and their motion is determined by initial conditions, in this work we focus on active objects, which combine passive elasticity with an internal actuation mechanism such as muscles or robotic actuators.

The model we use for our experiments is based on deformable meshes simulated using the finite element method (FEM) and contractile fibers modeled as springs. As seen in Figure 2, the output of the simulator depends on the action signal $\mathbf{a}$. More specifically, it is computed by minimizing the potential energy $E$ of the system. For soft-body physics, the potential energy is conventionally defined as a function of vertex positions $\mathbf{x}$, but to support actuation mechanisms we consider energy functions $E$ that also take action signals $\mathbf{a}$ as input:

The energy function we use for each contractile fiber in the model shown in Figure 2 is:

where $l(\mathbf{x})$ is the current length of the fiber, $l_{0}$ is its original rest length and $k$ is its stiffness. However, we want to emphasize that the approach we propose to achieve differentiability is not tied to this particular model. The important characteristic of the model presented in this section is that it defines the output of the simulator via energy minimization (Equation 1), but many possible energy functions exist for this purpose and our approach can achieve differentiability for arbitrary energy functions.

The energy function we use for our experiments also includes contributions from triangle deformations due to Neo-Hookean elasticity. Although there are many variants of Neo-Hookean elasticity (Smith et al., 2018), the specific model we use is based on the energy density function $\Psi$ presented in (Sifakis & Barbič, 2012), which can be expressed in terms of the isotropic invariants $I_{1}$ and $J$ that can be computed from vertex positions $\mathbf{x}$:

The only modification we make to this model is that we use a quadratic approximation of the logarithmic function to avoid issues with non-positive values, which may appear when triangles are inverted (Irving et al., 2004). We refer to (Sifakis & Barbič, 2012) for additional details on how to compute $I_{1}$ and $J$ from vertex positions $\mathbf{x}$ and additional background in FEM simulation. However, we want to emphasize again that our goal in this paper is not to advocate any specific elasticity model or actuation mechanism. We provide some details of the energy functions we use for our experiments for reference, but our goal is to provide a general framework to achieve differentiability in soft-body simulations where the forward pass requires a minimization procedure. The specific function that the simulator minimizes and the minimization algorithm implemented in the simulator may vary depending on application requirements.

Although the specifics may have some differences, there is a high-level interpretation of the energy functions commonly used in soft-body physics that can reveal some of their similarities with neural networks, which are not always evident in more conventional presentations of this topic in the context of continuum mechanics and partial differential equations. Many energy models used in FEM simulation such as Neo-Hookean, Saint Venant-Kirchhoff and corotated elasticity are defined as a non-linear function of the deformation gradient $\mathbf{F}$, a square matrix (2x2 for triangles and 3x3 for tetrahedra) that can be defined as a linear function of the vertex positions $\mathbf{x}$. In that sense, all these models can be understood as shallow neural networks that use different non-linearities. We can also see a similar pattern in spring energy functions. Figure 3 shows a diagram that makes this interpretation more evident.

One difference with more conventional neural network architectures is that the non-linearity in FEM models is not element-wise. For example, an FEM model on 2D triangle meshes computes the energy of every triangle by applying a non-linear function on 2x2 matrices. This approach has some similarities with capsule networks (Sabour et al., 2017), since the non-linearity acts on groups of neurons whose values represent properties of the same entity (one triangle of the mesh), and a similar observation can be made about spring energy functions.

In fact, the energy function used in our simulator is implemented using differentiable operations commonly available in deep learning packages (we used PyTorch (Paszke et al., 2019), without any additional libraries). This also allows us to automatically differentiate the energy function which is useful to implement the forward pass that requires running a minimization procedure (Equation 1).

Note, however, that the fact that we can implement the energy function using differentiable operations does not immediately provide us with a backward pass for the simulator. Although explicitly unrolling a minimization procedure (for example, gradient descent with a fixed step size) could in principle allow us to run the backward pass, this is not necessarily the best option, since it requires storing in memory all the intermediate steps of the optimizer and limits its applicability to optimization methods that only use differentiable operations in their implementation.

In the next section, we present an implicit differentation approach that allows us to differentiate through the simulator regardless of how the optimization procedure is implemented, it does not require storing intermediate steps and is applicable even if the optimizer runs non-differentiable operations such as line search, which is common in implementations of soft-body physics (Wang & Yang, 2016; Gast & Schroeder, 2015; Liu et al., 2017; Smith et al., 2018).

For machine learning problems, we are interested in incorporating physics into neural networks as a layer, as shown in Figure 4. The difficulty in running backpropagation on the computation graph presented in Figure 4 is that $\frac{dL}{d\mathbf{a}}$ is not readily available because $\mathbf{x_{1}}$ is not defined explicitly as a differentiable function of $\mathbf{a}$. In this section, we present an implicit differentiation approach to compute this derivative, and we point out some aspects that make this approach feasible to implement in reverse-mode automatic differentiation systems commonly used for neural network training. To present an initial derivation using conventional matrix operations, we assume vector representations using one-dimensional arrays for the vertex positions $\mathbf{x}\in\mathbb{R}^{2n}$ (where $n$ is the number of vertices) and the action $\mathbf{a}\in\mathbb{R}^{a}$ (where $a$ is the number of actuators), but we show how to generalize this derivation to support other representations for $\mathbf{x}$ and $\mathbf{a}$ in Section 4.1.

The output of the simulation layer is defined implicitly via energy minimization, which implies that the force $\mathbf{f}\in\mathbb{R}^{2n}$ (Equation 4) at the solution is zero (Equation 5).

We observe that the fact that $\mathbf{f}$ is zero for any $\mathbf{a}$ (Equation 5) implies that $\frac{d\mathbf{f}}{d{\mathbf{a}}}$ is also zero for any $\mathbf{a}$. Applying the chain rule, we get:

In the next equations, we will omit the explicit evaluations at $(\mathbf{x_{1}}(\mathbf{a}),\mathbf{a})$ for notational convenience. From Equation 6 we can obtain a formula for $\frac{d{\mathbf{x_{1}}}}{d{\mathbf{a}}}$:

Considering the loss $L$ as a function of $\mathbf{x}$, we can use Equation 7 to obtain a formula for $\frac{d{L}}{d{\mathbf{a}}}$:

We then reorder the matrix multiplications of Equation 8 to get our final formula for the backward pass:

The derivation of the last equation leveraged the fact that $-\frac{\partial\mathbf{f}}{\partial\mathbf{x}}$ is symmetric, since it is equal to the Hessian of the energy as can be seen via Equation 4. This Hessian is positive definite when evaluated at a quasistatic configuration $\mathbf{x_{1}}(\mathbf{a})$, since it corresponds to a local minimum of the energy $E$. Thus, the quantity inside the parenthesis in Equation 9, which requires solving a linear system, can be conveniently computed via a conjugate gradient solver in a matrix-free fashion. The fact that we can implement Equation 9 without explicitly constructing the Hessian matrix is particularly important in the context of neural networks. Neural networks are often implemented using reverse-mode automatic differentiation systems which do not provide an efficient way to compute Hessian matrices, but can compute Hessian-vector products in a matrix-free fashion, which is sufficient to implement Equation 9 (we used PyTorch). We provide more details of this matrix-free implementation in Section 4.1.

We now consider the effect of inertia, which we previously ignored in quasistatic motions. For this purpose, we consider sequences of states $\mathbf{s_{t}}$ that evolve over time. The state of our simulation is defined by vertex positions and velocities $\mathbf{s_{t}}=(\mathbf{x_{t}},\mathbf{v_{t}})$. A good starting point to characterize the dynamics of our system is to consider the equations of motion in a continuous-time setting:

where $\mathbf{f}$ is the force and $\mathbf{M}$ is the mass. For simulation purposes, we assume a particular time discretization scheme given by a numerical integration method, defining a function $\mathrm{sim}$ that updates the state of the system after $h$ units of time:

Explicit methods such as symplectic Euler have been used before for differentiable soft-body physics (Hu et al., 2019c, a). However, in this paper we focus on implicit numerical integration methods that do not provide an explicit differentiable forward pass. Although implicit numerical integration methods are sometimes formulated as root finding problems using the force function, many implicit numerical integration methods can be reformulated as minimization problems using the energy function. The latter formulation is preferred in practical implementations due to its better numerical properties (Martin et al., 2011; Gast & Schroeder, 2015; Liu et al., 2017; Rojas et al., 2018). We adopt the optimization-based formulation of implicit numerical integration for our differentiable simulator. The general form that these methods have is the following:

where $g$ is a scalar-valued function that includes the potential energy $E$, and $G$ is an explicit velocity update rule. For example, the optimization-based formulation of backward Euler integration commonly used in simulation (Gast & Schroeder, 2015; Liu et al., 2017) can be written as:

The implicit update rule for $\mathbf{x_{1}}$ can be differentiated using the same implicit differentiation strategy presented in Section 4 and the explicit update rule for $\mathbf{v_{1}}$ can be differentiated using standard backpropagation. This provides us with a fully-differentiable implicit numerical integration layer (Figure 5) that we can use for policy optimization.

We presented a soft-body physics layer that can be differentiated even when an explicit differentiable forward pass is not available. By adopting an optimization-based approach to define state transitions, we derived a differentiation rule that is agnostic to the specific implementation of the forward pass, which allows us to achieve differentiability for implicit numerical integration methods. We showed how our implicit differentiation approach can be implemented using reverse-mode automatic differentiation, which allows us to implement and compose our simulator with neural networks using the same automatic differentiation system.

As with recent work in differentiable physics for explicit methods, such as DiffTaichi (Hu et al., 2019a), one interesting direction of future work for implicit methods is to devise performance optimizations for the simulator while keeping the same flexibility and productivity provided by automatic differentiation. Our current implementation does not take full advantage of spatial sparsity, which is one of the main reasons why simulators implemented in DiffTaichi are often faster than implementations in conventional deep learning packages. The forward pass of our simulator can also be further improved by adopting different numerical optimization strategies. We adopted a descent method with line search that only uses information of first-order derivatives, but this could be improved in the future. Our goal in this paper is not to provide a specific implementation of the forward pass, but rather to present a general approach that leverages reverse-mode automatic differentiation and allows us to use different implementations of the forward pass while keeping the simulator differentiable.

We demonstrated that our differentiable simulator is compatible with backpropagation through time and can be used to optimize control policies taking into account the effect of multiple simulation steps. However, as with other recurrent models, vanishing and exploding gradients could become an issue for very long time horizons. We ran some preliminary experiments that indicate that combining our differentiable simulator with a trainable value function can mitigate this issue, which is also common in infinite-horizon reinforcement learning, but further research is needed to properly assess the benefits of this approach. Although a differentiable physics engine allows us to optimize deterministic policies, it is possible that adding some stochasticity, for example to improve exploration during training, could mitigate potential issues with local optima, which has been reported in prior work in differentiable physics. We believe it is worth evaluating the benefits of incorporating these modifications into the training loop in future work.

We would like to thank Alexander Winkler for the useful discussions. This material is based upon work supported by the National Science Foundation under Grant Numbers IIS-2008915, IIS-2008584, IIS-1763638 and IIS-1764071. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.