DiffTaichi: Differentiable Programming for Physical Simulation

1 INTRODUCTION . Differentiable physical simulators are effective components in machine learning systems . For example , de Avila Belbute-Peres et al . ( 2018a ) and Hu et al . ( 2019b ) have shown that controller optimization with differentiable simulators converges one to four orders of magnitude faster than model-free reinforcement learning algorithms . The presence of differentiable physical simulators in the inner loop of these applications makes their performance vitally important . Unfortunately , using existing tools it is difficult to implement these simulators with high performance . We present DiffTaichi , a new differentiable programming language for high performance physical simulations on both CPU and GPU . It is based on the Taichi programming language ( Hu et al. , 2019a ) . The DiffTaichi automatic differentiation system is designed to suit key language features required by physical simulation , yet often missing in existing differentiable programming tools , as detailed below : Megakernels Our language uses a “ megakernel ” approach , allowing the programmer to naturally fuse multiple stages of computation into a single kernel , which is later differentiated using source code transformations and just-in-time compilation . Compared to the linear algebra operators in TensorFlow ( Abadi et al. , 2016 ) and PyTorch ( Paszke et al. , 2017 ) , DiffTaichi kernels have higher arithmetic intensity and are therefore more efficient for physical simulation tasks . Imperative Parallel Programming In contrast to functional array programming languages that are popular in modern deep learning ( Bergstra et al. , 2010 ; Abadi et al. , 2016 ; Li et al. , 2018b ) , most traditional physical simulation programs are written in imperative languages such as Fortran and C++ . DiffTaichi likewise adopts an imperative approach . The language provides parallel loops and control flows ( such as “ if ” statements ) , which are widely used constructs in physical simulations : they simplify common tasks such as handling collisions , evaluating boundary conditions , and building iterative solvers . Using an imperative style makes it easier to port existing physical simulation code to DiffTaichi . Flexible Indexing Existing parallel differentiable programming systems provide element-wise operations on arrays of the same shape , e.g . c [ i , j ] = a [ i , j ] + b [ i , j ] . However , many physical simulation operations , such as numerical stencils and particle-grid interactions are not elementwise . Common simulation patterns such as y [ p [ i ] * 2 , j ] = x [ q [ i + j ] ] can only be expressed with unintuitive scatter/gather operations in these existing systems , which are not only inefficient but also hard to develop and maintain . On the other hand , in DiffTaichi , the programmer directly manipulates array elements via arbitrary indexing , thus allowing partial updates of global arrays and making these common simulation patterns naturally expressible . The explicit indexing syntax also makes it easy for the compiler to perform access optimizations ( Hu et al. , 2019a ) . The three requirements motivated us to design a tailored two-scale automatic differentiation system , which makes DiffTaichi especially suitable for developing complex and high-performance differentiable physical simulators , possibly with neural network controllers ( Fig . 1 , left ) . Using our language , we are able to quickly implement and automatically differentiate 10 physical simulators1 , covering rigid bodies , deformable objects , and fluids ( Fig . 1 , right ) . A comprehensive comparison between DiffTaichiand other differentiable programming tools is in Appendix A . 2 BACKGROUND : THE TAICHI PROGRAMMING LANGUAGE . DiffTaichi is based on the Taichi programming language ( Hu et al. , 2019a ) . Taichi is an imperative programming language embedded in C++14 . It delivers both high performance and high productivity on modern hardware . The key design that distinguishes Taichi from other imperative programming languages such as C++/CUDA is the decoupling of computation from data structures . This allows programmers to easily switch between different data layouts and access data structures with indices ( i.e . x [ i , j , k ] ) , as if they are normal dense arrays , regardless of the underlying layout . The Taichi compiler then takes both the data structure and algorithm information to apply performance optimizations . Taichi provides “ parallel-for '' loops as a first-class construct . These designs make Taichi especially suitable for writing high-performance physical simulators . For more details , readers are referred to Hu et al . ( 2019a ) . The DiffTaichi language frontend is embedded in Python , and a Python AST transformer compiles DiffTaichi code to Taichi intermediate representation ( IR ) . Unlike Python , the DiffTaichi language is compiled , statically-typed , parallel , and differentiable . We extend the Taichi compiler to further compile and automatically differentiate the generated Taichi IR into forward and backward executables . 1Our language , compiler , and simulator code is open-source . All the results in this work can be reproduced by a single Python script . Visual results in this work are presented in the supplemental video . We demonstrate the language using a mass-spring simulator , with three springs and three mass points , as shown right . In this section we introduce the forward simulator using the DiffTaichi frontend of Taichi , which is an easier-to-use wrapper of the Taichi C++14 frontend . Allocating Global Variables Firstly we allocate a set of global tensors to store the simulation state . These tensors include a scalar loss of type float32 , 2D tensors x , v , force of size steps ×n_springs and type float32x2 , and 1D arrays of size n_spring for spring properties : spring_anchor_a ( int32 ) , spring_anchor_b ( int32 ) , spring_length ( float32 ) . Defining Kernels A mass-spring system is modeled by Hooke ’ s law F = k ( ‖xa − xb‖2 − l0 ) xa−xb ‖xa−xb‖2 where k is the spring stiffness , F is spring force , xa and xb are the positions of two mass points , and l0 is the rest length . The following kernel loops over all the springs and scatters forces to mass points : @ ti.kernel def apply_spring_force ( t : ti.i32 ) : # Kernels can have parameters . Here t is a parameter with type int32 . for i in range ( n_springs ) : # A parallel for , preferably on GPU a , b = spring_anchor_a [ i ] , spring_anchor_b [ i ] x_a , x_b = x [ t - 1 , a ] , x [ t - 1 , b ] dist = x_a - x_b length = dist.norm ( ) + 1e-4 F = ( length - spring_length [ i ] ) * spring_stiffness * dist / length # Apply spring impulses to mass points . force [ t , a ] += -F # `` += '' is atomic by default force [ t , b ] += F For each particle i , we use semi-implicit Euler time integration with damping : vt , i = e−∆tαvt−1 , i+ ∆t mi Ft , i , xt , i = xt−1 , i + ∆tvt , i , where vt , i , xt , i , mi are the velocity , position and mass of particle i at time step t , respectively . α is a damping factor . The kernel is as follows : @ ti.kernel def time_integrate ( t : ti.i32 ) : for i in range ( n_objects ) : s = math.exp ( -dt * damping ) # Compile-time evaluation since dt and damping are constants v [ t , i ] = s * v [ t - 1 , i ] + dt * force [ t , i ] / mass # mass = 1 in this example x [ t , i ] = x [ t - 1 , i ] + dt * v [ t , i ] Assembling the Forward Simulator With these components , we define the forward time integration : def forward ( ) : for t in range ( 1 , steps ) : apply_spring_force ( t ) time_integrate ( t ) 3 AUTOMATICALLY DIFFERENTIATING PHYSICAL SIMULATORS IN TAICHI . The main goal of DiffTaichi ’ s automatic differentiation ( AD ) system is to generate gradient simulators automatically with minimal code changes to the traditional forward simulators . Design Decision Source Code Transformation ( SCT ) ( Griewank & Walther , 2008 ) and Tracing ( Wengert , 1964 ) are common choices when designing AD systems . In our setting , using SCT to differentiate a whole simulator with thousands of time steps , results in high performance yet poor flexibility and long compilation time . On the other hand , naively adopting tracing provides flexibility yet poor performance , since the “ megakernel '' structure is not preserved during backpropagation . To get both performance and flexibility , we developed a two-scale automatic differentiation system ( Figure 2 ) : we use SCT for differentiating within kernels , and use a light-weight tape that only stores function pointers and arguments for end-to-end simulation differentiation . The global tensors are natural checkpoints for gradient evaluation . Assumption Unlike functional programming languages where immutable output buffers are generated , imperative programming allows programmers to freely modify global tensors . To make automatic differentiation well-defined under this setting , we make the following assumption on imperative kernels : Global Data Access Rules : 1 ) If a global tensor element is written more than once , then starting from the second write , the write must come in the form of an atomic add ( “ accumulation ” ) . 2 ) No read accesses happen to a global tensor element , until its accumulation is done . In forward simulators , programmers may make subtle changes to satisfy the rules . For instance , in the mass-spring simulation example , we record the whole history of x and v , instead of keeping only the latest values . The memory consumption issues caused by this can be alleviated via checkpointing , as discussed later in Appendix D. With these assumptions , kernels will not overwrite the outputs of each other , and the goal of AD is clear : given a primal kernel f that takes as input X1 , X2 , . . . , Xn and outputs ( or accumulates to ) Y1 , Y2 , . . . , Ym , the generated gradient ( adjoint ) kernel f∗ should take as input X1 , X2 , . . . , Xn and Y ∗1 , Y ∗ 2 , . . . , Y ∗ m and accumulate gradient contributions to X ∗ 1 , X ∗ 2 , . . . , X ∗ m , where each X ∗ i is an adjoint of Xi , i.e . ∂ ( loss ) /∂Xi . Storage Control of Adjoint Tensors Users can specify the storage of adjoint tensors using the Taichi data structure description language ( Hu et al. , 2019a ) , as if they are primal tensors . We also provide ti.root.lazy_grad ( ) to automatically place the adjoint tensors following the layout of their primals . 3.1 LOCAL AD : DIFFERENTIATING TAICHI KERNELS USING SOURCE CODE TRANSFORMS . A typical Taichi kernel consists of multiple levels of for loops and a body block . To make later AD easier , we introduce two basic code transforms to simplify the loop body , as detailed below . Flatten Branching In physical simulation branches are common , e.g. , when implementing boundary conditions and collisions . To simplify the reverse-mode AD pass , we first replace “ if ” statements with ternary operators select ( cond , value_if_true , value_if_false ) , whose gradients are clearly defined ( Fig . 3 , middle ) . This is a common transformation in program vectorization ( e.g . Karrenberg & Hack ( 2011 ) ; Pharr & Mark ( 2012 ) ) . Eliminate Mutable Local Variables After removing branching , we end up with straight-line loop bodies . To further simplify the IR and make the procedure truly single-assignment , we apply a series of local variable store forwarding transforms , until the mutable local variables can be fully eliminated ( Fig . 3 , right ) . After these two custom IR simplification transforms , DiffTaichi only has to differentiate the straightline code without mutable variables , which it achieves with reverse-mode AD , using a standard source code transformation ( Griewank & Walther , 2008 ) . More details on this transform are in Appendix B. Loops Most loops in physical simulation are parallel loops , and during AD we preserve the parallel loop structures . For loops that are not explicitly marked as parallel , we reverse the loop order during AD transforms . We do not support loops that carry a mutating local variable since that would require a complex and costly run-time stack to maintain the history of local variables . Instead , users are instructed to employ global variables that satisfy the global data access rules . Parallelism and Thread Safety For forward simulation , we inherit the “ parallel-for '' construct from Taichi to map each loop iteration onto CPU/GPU threads . Programmers use atomic operations for thread safety . Our system can automatically differentiate these atomic operations . Gradient contributions in backward kernels are accumulated to the adjoint tensors via atomic adds .

Summary: This paper introduces DiffSim, a programming language for high-performance differentiable physics simulations. The paper demonstrates 10 different simulations with controller optimization. It shows that the proposed language is easier to use and faster than the other alternatives, such as CUDA and TensorFlow. At the end, the paper provides insightful discussions why the gradient of the simulation could be wrong.