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 describes DiffSim, a differentiable programming system for learning with physical simulation. The system (built on the Taichi system) allows users to specify a forward simulation in a Python-like syntax, after which the program is compiled and iteratively run in both forward-mode and gradients computed for system parameters and controllers, as desired. A variety of simple simulations are included, demonstrating that the automatically generated CUDA code runs as fast as hand-written CUDA code (and noticeably faster than TensorFlow or PyTorch implementations), while requiring far fewer lines of code. The final section details two issues--time of  impact errors due to discrete time intervals and gradient explosions with long time horizons--and some potential solutions.