DiffTaichi: Differentiable Programming for Physical Simulation
Authors: Yuanming Hu
Luke Anderson
Tzu-Mao Li
Qi Sun
Nathan Carr
Jonathan Ragan-Kelley
Fredo Durand
Venue:  (2020)
Keywords: Differentiable programming
robotics
optimal control
physical simulation
machine learning system

Abstract: We present DiffTaichi, a new differentiable programming language tailored for building high-performance differentiable physical simulators. Based on an imperative programming language, DiffTaichi generates gradients of simulation steps using source code transformations that preserve arithmetic intensity and parallelism. A light-weight tape is used to record the whole simulation program structure and replay the gradient kernels in a reversed order, for end-to-end backpropagation.
We demonstrate the performance and productivity of our language in gradient-based learning and optimization tasks on 10 different physical simulators. For example, a differentiable elastic object simulator written in our language is 4.2x shorter than the hand-engineered CUDA version yet runs as fast, and is 188x faster than the TensorFlow implementation.
Using our differentiable programs, neural network controllers are typically optimized within only tens of iterations.

Published as a conference paper at ICLR 2020
DIFFTAICHI: DIFFERENTIABLE PROGRAMMING FOR PHYSICAL SIMULATION
Yuanming Hu', Luke Anderson', Tzu-Mao Li', Qi Nathan Nathan Carr!,
Jonathan Ragan-Kelley', Frédo Durand'
MIT CSAIL
[yuanming, lukea, fredo] @mi edu {qisu, ncarunarrrrreaaaaaoa com
1 Adobe Research
UC Berkeley
{tzumao, rkk edu
ABSTRACT
We present DiffTaichi, differentiable programming language tailored for building high-performance differentiable physical simulators. Based on an imper- ative programming language, DiffTaichi generates gradients of simulation steps using source code transformations that preserve arithmetic intensity and paral- lelism. A light-weight tape is used record the whole simulation program struc- ture and replay the gradient kernels in reversed order, for end-to-end backprop- agation. We demonstrate the performance and productivity of our language gradient-basee learning and optimization tasks on 10 different physical simula- tors. For example, a differentiable elastic object simulator written in our language is 4.2x shorter than the hand-engineered CUDA version yet runs as fast, and is 188x faster than the TensorFlow implementation. Using our differentiable pro- grams, neural network controllers are typically optimized within only tens of iter- ations.
INTRODUCTION
Network
Policy
Phase
Coal
Weightsbiases
Weights/biases 2
Contcol CONN
diimpmygg
Architecture
tanh
FC,
lquit--p
maHL springirnnggg
NN Controller
NN Controller
NN Controller
RB
Initial State
Parameterizaa
Stite D
Differeniable
Iomulattonn
Stale
Differentablle
Simulation
2045 time steps ...
Statee 2017
Differentable Simulation
Lnve Function
waler rendererpy
volume renderes
smoligg
Time step
Time step
Time step 2047
electricg
billandgg
Figure 1: Left: Our language allows to seamlessly integrate neural network (NN) controller and physical simulation module, and update the weights the controller or the initial state param- eterization (blue). Our simulations typically have 512 2048 time steps, and each time step has up to one thousand parallel operations. Right: 10 differentiable simulators built with DiccTaichi.
Differentiable physical simulators are effective components machine learning systems. For exam- plc, de Avila Belbute-Peres et al. (2018a) and Hu et al. (2019b) have shown that controller optimiza- tion with differentiable simulators converges one to four orders 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 Dichhi a new differentiable programming language for high performance physical simulations on both CPU and GPU. It is based the Taichi programming language (Hu et al.,

---

Published as a conference paper at ICLR 2020
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 "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 Fortran and C++. DiffTaichi likewise adopts imperative approach. The language provides paralle! loops and control flows (such as "if" statements), which are widely used constructs in physical simulations: they simplify common tasks such handling collisions, evaluating boundary conditions, and build- ing iterative solvers. Using an imperative style makes easier to port existing physical simulation code to DiffTaichi.
Flexible Indexing Existing parallel differentiable programming systems provide slement-wiss op- erations on arrays of the same shape, e.g. c[i, j] a(i, il b[i, j]. However, many physical simulation operations, such numerical stencils and particle-gre -grid interactions wise. Common simulation patterns such as y[p[1] * 2, j1 x(q[t j1] 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 easy for the compiler perform access optimizations (Hu et al., 2019a).
not element-
The three requirements motivated us to design a tailored two-scale automatic differentiation sys- tem, which makes DiffTaichi especially suitable for developing complex and perf-peeformance dif- ferentiable physical simulators. possibly with neural network controllers (Fig. 1, left). Using our language, we are able to quickly implement and automatically differentiate 10 physical sinulators!, covering rigid bodies, deformable objects, and fluids (Fig. 1, right). A comprehensive comparison between DiffTaichiand other diffcrentiable programming tools 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 productiv- ity on modern hardware. The key design that distinguishes Taichi from other imperative program- ming languages such C+/CUDA 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, k1), 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 perfor- mance 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 Hu al. (2019a).
The DiffTaichi language frontend is embedded in Python, and Python AST transformer compiles DiffTaichi code to Taichi intermediate representation (IR). Unlike Python, the DiffTaichi language compiled, statically-typer parallel, and differentiable. We extend the Taichi compiler to further compile and automatically differentiate the generated Taichi IR into forward and backward executa- bles.
OOr language, compiler, and simulator code is open-source. All the results in this work can be reproduced by single Python script. Visual results in this work presented in the supplemental video.
2

---

Published as a conference paper at ICLR 2020
We demonstrate the language using 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 casier-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 scalar loss of type float32, 2D tensors x, v, force of size steps Xn_springs and type float32x2, and ID arrays of size n_spring for spring properties: spring_anchor_e (int32), spr spring__nchor__t (fnt32), spring_length (float32).
Defining Kernels mass-spring system modeled by Hooke's law = k[[|x, xalle
1o) Xs -x where is the spring stiffness, F is spring force, X and X₁ are the positions of two mass points, and lo is the rest length. The following kernel loops over all the springs and scatters forces to mass points:
|
6t1. kernel
www.sspprnn__rceeeee t1.132):
# Kerneis can have parameters. Here is parameter with type int32.
for (nseinn_p)) parallel for, preferobly on 6PU a, spring_ spring_ancher_bff] x_a, x_b xit 1, al, x[t 1, b] dist x_a length dist..orm() + 1e-4
(length
Iprinng_ength(ii)
spring_sti stieeesss dist / length
Apply spring impulses to mass paints.
forcest, a] forceft, b] +=
is atomic by defoult
For each particle we semi-implicit Euler time integration with damping: Vt,i c Âta 01-1, + At F, Xt, Xt-1 + Atv,, where Vt,i, Kt,i, my are the velocity, position and mass of particle i m:
at time step f, respectively. is damping factor. The kernel is as follows:
@ti. kernel time_ inttinnnteeettt: ti.132):
for i in range objecct))
damping) compile-tine evaluation since dt and damp ing are constants v[t, 1] v[t 1, i] dt force(t, 1] mass mass 1 in this exomple
math.exp(-dd.
x[t, i] x[t 1, i] v[t, i]
Assembling the Forward Simulator
With these components, define the forward time integra-
tion:
(orward
for in range(1, steps): apply_sprine ing. force(t) time_ (tme__integaateet)
AUTOMATICALLY DIFFERENTIATING PHYSICAL SIMULATORS IN TAICHI
The main goal of DiffTaichi's automatic differentiation (AD) system is to generate gradient simula- tors automatically with minimal code changes to the traditional forward simulators.
Design Decision Source Code Transformation (SCT) (Griewank & Walther, 2008) and Trac- ing (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 flexibil- ity 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):
SCT for differentiating within kernels, and use light-weight tape that only stores function pointers and arguments for end-to-end simulation differentiation. The global tensors natural checkpoints for gradient evaluation.
3

---

Published as a conference paper at ICLR 2020
Ditsin Python Fromenn
Tape System
AST Prgrosssinn
Simpifiation
Forward Program
Tape Contents
Program
conpute force, argsede)
compute. lass.grad:)
Tasta lexsl compute Jarcel
mox partcie argamie
nowe parteses graae..
Teichi Frontend Lowering
Stteerit
Eimulficatoo of
eave par sen cnep force, angant
angami
wove pwr.ppeeeieee Leien spreet
move Lowall
comsute.
Checking
Biniiate Misalle Local Vo
LLV Codeger SIGCUDA
compite lons, argenl
commite forte.g
Lekes gradle)
Figure 2: Left: The DiffTaichi system. We reuse some infrastructure (white boxes) from Taichi, while the blue boxes our extensions for differentiable programming. Right: The tape records kernel launches and replays the gradient kernlls in reverse order during backpropagation.
Assumption Unlike functional programming languages where immutable output buffers are gen- erated, imperative programming allows programmers freely modify global tensors. To make automatic differentiation well-defined under this setting, we make the following assumption im- perative kernels:
Global Data Access Rules:
1) If global tensor element is written than once, then starting from the second write, the write must come in the form of an atomic add ("accumulation"). No read accesses happen to 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, record the whole history of and v, instead of keep- ing only the latest values. The memory consumption issues caused by this can be alleviated via checkpointing, discussed later in Appendix D.
With these assumptions, kernels will not overwrite the outputs of each other, and the goal AD is clear: given primal kernel that takes as input X1, X2,....... Xn and outputs (or accumulates to) Yi, .,,.., Ym. the generated gradient (adjoint) kernel f should take as input X X2,... Xn and Y 22.2.... ..... Y and accumulate gradient contributions Xi, . where each X; is an adjoint of i.e.
i.c 0loss)
Storage Control of Adjoint Tensors Users specify the storage of adjoint tensors using the Taichi data structure description language (Hu al., 2019a), as they are primal tensors. We also provide ti. root. (az)_grad)) to automatically place the adjoint tensors following the layout of their primals.
3.1 LOCAL AD: DIFFERENTIATING TAICHI KERNELS USING SOURCE CODE TRANSFORMS
typical Taichi kernel consists of multiple levels of for loops and body block. To make later AD easier, we introduce two basic code transforms to simplify the loop body, as detailed below.
int 0; if (b > 0)
// flatten branching int a 0; select(b 6, b, 2b);
// eliminate mutable var
{ = b:} [ 2b;]
ssal select(b 0, b, 2b); ssa2 ssal + return ssa2;
return a;
return a;
Figure 3: Simple IR preprocessing before running the AD source code transform (left to right). Demonstrated c++, The actual Taichi IR often more complex. Containing loops are ignored.
Flatten Branching In physical simulation branches are common, e.g., when implementing bound- ary conditions and collisions. To simplify the reverse-mode AD pass, we first replace "if" statements with ternary operators select(cond, value_if_ if_true, salue_if.em ff__ae,)), whose gradients clearly de- fined (Fig. 3, middle). This is a common transformation program vectorization (e.g. Karrenberg & Hack (2011) 1); Pharr & Mark (2012)).
4

---

Published as a conference paper at ICLR 2020
Eliminate Mutable Local Variables After removing branching, we end up with straight-ling loop bodies. To further simplify the IR and make the procedure truly single-assin aassingegg, 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, Dach only has to differentiate the straight- line code without mutable variables, which it achieves with reverse- AD, using 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 not explicitly marked as parallel, we reverse the loop order during AD transforms. We do not support loops that carry mutating local variable since that would require a complex and costly run-time stack maintain the history 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 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 backward kernels are accumulated to the adjoint tensors via atomic adds.
3.2 GLOBAL AD: END-TO-END BACKPROPAGATION USING A WIGGTTWIIGHT TAPE
We construct a tape (Fig. 2, right) of the kernel execution so that gradient kernels can be replayed reversed order. The tape is very liightweight: since the intermediate results are stored in global tensors, during forward simulation the tape only records kernel names and the (scalar) input param- eters, unlike other differentiable functional array systems where all the intermediate buffers have to be recorded by the tape. Whenever DiffTaichi kernel is launched, pointer and parameters to the tape. When evaluating gradients, we traverse the reversed tape, and invoke the gradient kernels with the recorded parameters. Note that Diff'aichi AD is evaluating gradients with respect to input global tensors instead of the input parameters.
append the kernel function
Learning/Optimizattion with Gradients Now we revisit the mass-spring example and make differentiable for optimization. Suppose the goal is optimize the rest lengths of the springs so that the triangle area formed by the three springs becomes 0.2 at the end of the simulation. We first define the loss function:
kernel
def compute_losse loss(t: ti.132): xe1 xft, 01 xft, 11 x02 x[t,
x(t, 2]
# Triangle area from cross product area ab.. (x0[[]]xx02[1] x01[1])x0[001)
Goul: Adjust the spring rest lengths, that this area=0.2 after 1024 time steps (initial area=0)0055)
target_area 0.2
[Non[N ti.sqr(area target_area) # Everything in Toicht is tensor.
g "loss" is a scalar (9-D tensor), thereby indexed with [None].
The programmer uses Tape memorize forward kernel launches. It automatically replays the gradients of these kernels in reverse for backpropagation. Initially the springs have lengths [0.1, 0.14], and after optimization the rest lengths are [0.600, 0.600, 0.529]. This means the springs will expand the triangle according to Hooke's s law and form a larger triangle: [Reproduce: mass mass_spring
wprinngsinplee.pppp.. simple.py]
5

---

Published as a conference paper at ICLR 2020
def main():
Spring Rest Length Optimization
for iter in r20g)
with ti Tape (aapee(osse (loss):
0.000
forward
compute_ Loss_ (steps
0.008
print ('Ieer", iter)
0.004
print (osss d Gradtent descent
0.002
for in range inprinn::))
0.000
ppring._engttlle
0
10
20
Ir * sorring_engthe grad[i]
Gradient descent iberations
Complex Kernels Sometimes the user may want to override the gradients provided by the com- piler. For example, when differentiating 3D singular value decomposition done with an iterative solver, it is better to use manually engineered SVD derivative subroutine for better stability. We provide two more decorators ti. compomleex__ pe_neelel and ti. wtw.llx_keeeeeeeeeeeeeeeeeeeeeee automatic differentiation, as detailed Appendix C. Apart from custom gradients, complex kernels can also be used implement checkpointing, detailed in Appendix D.
overwrite the default
4 EVALUATION
We evaluate DiffTaichi on 10 different physical simulators covering large-scale continuum and small-scale rigid body simulations. All results can reproduced with the provided script. The dynamic/optimizzaion processes are visualized in the supplemental video. In this section we focus our discussions on three simulators. More details on the simulators are Appendix E.
4.1 DIFFERENTABLE CONTINUUM MECHANICS FOR ELASTIC OBJECTS [diffmpm]
First, we build differentiable continuum simulation for soft robotics applications. The physical system is governed by momentum and mass conservation, i.c. p DV V a + pg. Dr + pV We follow ChainQucen's 's implementation (Hu et al., 2019b) and use the moving least squares material point method (Hu et al., 2018) to simulate the system. We were able easily translate the original CUDA simulator into DiffTaichi syntax. Using this simulator and an open-loop controller, can easily train a soft robot to move forward (Fig. 1, diffmmm).
Dp
Performance and Productivity Compared with manual gradient implementations in (Hu et al., 2019b), getting gradients DiffTaichi is effortless. a result, the DiffTaichi implementation is 4.2x shorter in terms of lines of code, and runs almost as fast; compared with TensorFlow, DiffTaichi code 1.7x shorter and 188x faster (Table 1). The Tensorflow implementation verbose due to the heavy use of tf. saither_nd/scatteee and array transposing and broadcasting.
Table 1: di ffmp performance comparison using 6.4K particles. For the lines of code, we only include the essential implementation, excluding boilerplate code. [Reproduce: python3 diffmpm diffmmpm_enchhneeeee benchmark.py)
NVIDIA GTX 1080 Ti GPU. We benchmark in 2D
Approach
Forward Time
Backward Time
Total Time
# Lines of Code
TensorFlow
13.20
35.70 70 ms
48.90 (188.x) x)
190
CUDA
0.10 ms
0.14 ms
0.24 ms (0.92x)
460
DiffTaichi
0.1 I ms
0.15 ms
0.26 ms (1.00x)
110
4.2 DIFFERENTIABLE INCOMPRESSIBLE FLUID SIMULATOR [smoke]
We implemented smoke simulator (Fig. 1, smoke) with semi-Lagrangian advection (Stam, 1999) and implicit pressure projection, following the example in Autograd (Maclaurin et al., 2015). Using gradient descent optimization changes the pattern of the fluid to a target image (Fig. 7a in Appendix). We compare the performance of our system against PyTorch, Autograd, and JAX in Table 2. Note that example from the
the initial velocity field,
able find velocity field that

---

Published as a conference paper at ICLR 2020
Table 2: smoke benchmark against Autograd, PyTorch, and JAX. We used 110 110 grid and 100 time steps, each with 6 Jacobi pressure projections. |Reproduce: python3 smoke_[autogee lauuograd/pytorchh/a//aaa hethhhaaaaaacchiipppaaahhhhi gpu].py]. Note that the Autograd program uses float64 precision, which approximately doubles the run time.
Approach
Forward Time
Backward Time
Total Time
# Essential LoC
PyTorch (CPU, (32)
405 ms
328 ms
733 ms (13.8x)
74
PyTorch (GPU, 132)
254 ms
457 ms
711 ms (13.4x)
74
Autograd (CPU, f64)
307 ms
1197 ms
1504 ms (28.4x)
51
JAX (GPU, f32)
24 ms
75 ms
99 ms (1.9x)
90
DiffTaichi (CPU, f32)
66 ms
132
198 ms (3.7x)
75
DiffTaichi (GPU, f32)
24 ms
29 ms
53 ms (1.0x)
75
Autograd library, this grid-based simulator is intentionally simplified to suit traditional array-based programs. For example, a periodic boundary condition used so that Autograd can represent using numpy.roll without any branching. Still, Taichi delivers higher performance than these array- based systems. The whole program takes 10 seconds to run in DiffTaichi on a GPU, and seconds spent JIT. JAX JIT compilation takes 2 minutes.
4.3 DIFFERENTIARLA RIGID BODY STMULATORS [rigid_body]
We built an impulse-based (Catto, 2009) differentiable rigid body simulator (Fig. 1, riiii_bbddl for optimizing robot controllers. This simulator supports rigid body collision and friction, spring forces, joints, and actuation. The simulation is end-to-end differentiable except for a countable number of discontinuities. Interestingly, although the forward simulator works well, naively differentiating it with DiffTaichi leads to completely misleading gradients, due to the rigid body collisions. We discuss the cause and solution of this issue below.
Improving collision gradients Consider the rigid ball example in Fig. 4 (left), where a rigid ball collides with a friction-less ground. Gravity is ignored, and due to conservation of kinetic energy the ball keeps a constant speed even after this clastic collision.
In the forward simulation, using small At often leads to reasonable result, done in many physics simulators. Lowering the initial ball height will increase the final ball height, since there is less distance to travel before the ball hits the ground and more after (see the loss curves in Fig.4, middle right). However, using a naive time integrator, no matter how small At is, the evaluated gradient of final height w.r.t. initial height will be 1 instead of -1. This counter-intuitive behavior is due to the fact that time discretization itself is not differentiated by the compiler. Fig. 4 explains this effect in greater detail.
Naive Time Integrator
With Time of Impact
Tine of Grodients
Effect Grodients
0.25
Lase Paaved
t=:41t
t=
1=
t=
LEsE
033
t-3
1-AAL
Lese (+1t))
Crediant 700
0.21
058
255
6.0
0.0
6.70
0.75
0.70
3.52
0.20
0.75
1=201
Hight
Iakal Hogg
Figure 4: How gradients can go wrong with naive time integrators. For clarity we use large At here. Left: Since collision detection only happens at multiples of At (2At in this case), lowering the initial position of the ball (light blue) leads to lowered final position. Middle Left: By improving the time integrator to support continuous time of impact (TOI), collisions can be detected at any time, e.g. 1.9A/ (light red). Now the blue ball ends up higher than the green one. Middle Right: Although the two time integration techniques lead to almost identical forward results (in practice At small), the naive time integrator gives an incorrect gradient of 1, but adding TOI yields the correct gradient. Please see supplemental video for better demonstration. [Reproduce: python3 rigid body [tt_otoip]]]] Right: When zooming in, the loss of the naive integrator is decreasing, and the saw-tooth pattern explains the positive gradients. [Reproduce: python3 rigid
zoom]
7

---

Published as a conference paper at ICLR 2020
Mass Spring Robot
Mass Spring Robot
Rigid Body Robot 1
-0.1
-0.2
-0.30
0.3
0.7
-0.35
-0.40
Naluo
-0.4
Na.v
Narrr
-0.45
-0.5
T01
20
60
100
20
80
15
Grodiont Korations
Gradient
Grosion: Descont Rarations
Figure 5: Adding TOI greatly improves gradient and optimization quality. Each experiment is repeated five times. [Reproduce: python3 [mass_ chaas_spring/rriid _ody.pyy [1/2] plot & python3 losses.py
We propose simple solution of adding continuous collision resolution (see, for example, Redon et al. (2002)), which considers precise time of impact (TOI), to the forward program (Fig. 4, middle left). Although barely improves the forward simulation (Fig. 4, middle right), the gradient will be corrected effectively (Fig. 4, right). The details of continuous collision detection are in Appendix F. In real-world simulators, we find the TOI technique leads to significant improvement in gradient quality in controller optimization tasks (Fig. 5). Having TOI or not barely affects forward simula- tion: the supplemental video, we show that a robot controller optimized a simulator with TOI, actually works well in simulator without TOI.
The takeaway is, differentiating physical simulators does not always yield useful gradients of the physical system being simulated, even if the simulator does forward simulation well. In Appendix G, we discuss some additional gradient issues we have encountered.
5 RELATED WORK
Differentiable programming The recent rise of deep learning has motivated the development of differentiable programming libraries for deep NNs, most notably auto-differentiation frameworks such as Theano (Bergstra et al., 2010), TensorFlow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017). However, physical simulation requires complex and customizable operations due to the intrinsic computational irregularity. Using the aforementioned frameworks, programmers have to compose these coarse-gras gaaisee basic operations into desired complex operations. Doing so often leads to unsatisfactory performance.
Earlier work on automatic differentiation focuses on transforming existing scalar code to obtain derivatives (e.g. Utke et al. (2008), Hascoet & Pascual (2013), Pearlmutter & Siskind (2008)). recent trend has emerged for modern programming languages to support differentiable function transformations through annotation (e.g. Innes et al. (2019), Wei et al. (2019)). These frameworks enable differentiating general programming languages, yet they provide limited parallelism.
Differentiable array programming languages such as Halide (Ragan- 2018b), Autograd (Maclaurin et al., 2015), JAX (Bradbury et al., 2018), and Enoki (Jakob, 2019) operate on arrays instead of scalars utilize parallelism. Instead of operating on arrays that are im- mutable, DiffTaichi uses an imperative style with flexible indexing to make porting existing physical simulation algorithms easier.
al., 2013; Li et al.,
Differentiable Physical Simulators Building differentiable simulators for robotics and machine learning has recently increased in popularity. Without differentiable programming, Battaglia et al. (2016), Chang al. (2016) and Mrowca et al. (2018) used NNs to approximate the physical pro- cess and used the NN gradients as the approximate simulation gradients. Degrave et al. (2016) and de Avila Belbute-Peres et (2018b) used Theano and PyTorch respectively build differen- tiable rigid body simulators. Schenck & Fox (2018) differentiates position-based fluid using custom CUDA kernels. Popović et al. (2000) used a differentiable rigid body simulator for manipulating physically based animations. The ChainQueen differentiable elastic object simulator (Hu et al., 2019b) implements forward and gradient versions of continuum mechanics in hand-written CUDA kernels, leading to performance that is two orders of magnitude higher than pure TensorFlow implementation. Liang et al. (2019) built a differentiable cloth simulator for material estimation and motion control. The deep learning community also often incorporates differentiable rendering
8

---

Published as a conference paper at ICLR 2020
operations (OpenDR (Loper & Black, 2014), N3MR (Kato et al., 2018), redner (Li al., 2018a), Mitsuba (Nimier-Davide et al., 2019)) to learn from 3D scenes.
CONCLUSION
We have presented DiffTaichi, a new differentiable programming language designed specifically for building peghoperformance differentiable physical simulators. Motivated by the need for supporting megakernels, imperative programming, and flexible indexing, we developed a tailored two-scale automatic differentiation system. We used Diff'aichi to build 10 simulators and integrated them into deep neural networks, which proved the performance and productivity of DiffTaichi over existing systems, We hope our programming language can greatly lower the barrier of future research on differentiable physical simulation in the machine learning and robotics communities.

---

Published as a conference paper at ICLR 2020
BIBLIOGRAPHY
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, al. Tensorflow: system for large- scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Imple- mentation (OSDI 16), pp. (665-883, 2016.
Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Inter- action networks for learning about objects, relations and physics. 2016.
James Bergstra, Olivier Breuleux, Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, Guillaume Desjardins, Joseph Turian, David Warde-Farley, and Yoshua Bengio. Theano: A epu and gpu math compiler in python. In Proc. 9th Python Science Conf, volume I, pp. 3-10, 2010.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne. JAX: composable transformations of Python+NumPy programs, 2018. URL http: giggtut.bb... com/ gowggoooglleaaa
Erin Catto. Modeling and solving constraints. In Game Developers Conference, pp- 16, 2009.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach to learning physical dynamics. ICLR, 2016.
Filipe de Avila Belbute-Peres Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J Zico Kolter. End-to-end differentiable physics for learning and control. In Advances in Neural Information Processing Systems, pp. 7178-7189, 2018a
Filipe de Avila Belbute- Peres, Kevin A Smith, Kelsey Allen, Joshua B Tenenbaum, and Zico Kolter. End-to-end differentiable physics for learning and control. In Neural Information Pro- cessing Systems, 2018b.
Jonas Degrave, Michiel Hermans, Joni Dambre, et al. A differentiable physics engine for deep learning in robotics. arXiv preprint 1rX1v:1611.01652 2016.
Ronald M Errico. What is an adjoint model? Bulletin of the American Meteorological Society, 78 (11):2577-2592, 1997.
Richard Gordon, Robert Bender, and Gabor T Herman. Algebraic reconstruction techniques (art) for three-dimensional electron microscopy and x-ray photography. Journal of theoretical Biology, 29(3):471-48) 1970.
Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algo- rithmic differentiation, volume 105. Siam, 2008
Laurent Hascoet and Valérie Pascual. The Tapenade automatic differentiation tool: Principles, model, and specification. 39(3):20:1-20:44 2013.
Yuanming Hu, Yu Fang, Ziheng Ge, Ziyin Qu, Yixin Zhu, Andre Pradhana, and Chenfanfu Jiang. A moving least squares material point method with displacement discontinuity and two-way rigid body coupling. 37(4):150, 2018.
Yuanming Hu, Tzu-Mao Li, Luke Anderson, Jonathan Ragan-Kelley, and Frédo Durand. Taichi: A language for high-performance computation on spatially sparse data structures. SIGGRAPH Asia 2019 Technical Papers, pp. 201. ACM, 2019a.
Yuanming Hu, Jiancheng Liu, Andrew Spielberg, Joshua Tenenbaum, William T Freeman, Jiajun Wu, Daniela Rus, and Wojciech Matusik. Chainqueen: A real-time differentiable physical simula- for soft robotics. Proceedings of IEEE International Conference on Robotics and Automation (ICRA), 2019b.
Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckus, Elliot Saba, Viral B Shah, and Will Tebbutt. Zygote: differentiable programming system to bridge machine learning and scientific computing. arXiv preprint arXiv: /00.00/907007777. 2019.
10

---

Published as a conference paper at ICLR 2020
Wenzel Jakob. Enoki: structured vectorization and differentiation modern processor architec- tures, 2019. htwp://www..oomiituubeerreerrree
Chenfanfu Jiang. The material point method for the physics-based simulation of solids and fluids. University of Califomia, Los Angeles, 2015.
Ralf Karrenberg and Sebastian Hack. Whole-function vectorization. In Code Generation and Opti- mization, pp. 141-150, 2011.
Hiroharu Kato, Yoshitaka Ushiku, and Tatsuya Harada. Neural 3d mesh renderer. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp- 3907-3916, 2018.
Stig Larsson and Vidar Thomée. Purtial differential equations with numerical methods, volume 45. Springer Science & Business Media, 2008
Tzu-Mao Li, Miika Aittala, Frédo Durand, and Jaakko Lehtinen. Differentiable monte carlo ray tracing through edge sampling. In SIGGRAPH Asia 2018 Technical Papers, pp. 222. ACM, 2018a.
Tzu-Mao Li, Michaël Gharbi, Andrew Adams, Frédo Durand, and Jonathan Ragan entiable programming for image processing and deep learning in Halide. ACM Transactions on Graphics (TOG), (7(4): 139, 201
Ragan-Kelley, Kelley. Differ-
Junbang Liang, Ming C Lin, and Vladlen Koltun. Differentiable cloth simulation for inverse prob- lems. Advances in Neural Information Processing Systems, 2019.
Matthew M Loper and Michael J Black. Opendr: An approximate differentiable renderer. In Euro- pean Conference on Computer Vision, pp. 154-169. Springer, 2014.
Dougal Maclaurin, David Duvenaud, and Ryan P Adams. Autograd: Effortless gradients in numpy. In ICML 2015 AutoML Workshop, volume 238, 2015.
Aleka McAdams, Andrew Selle, Rasmus Tamstorf, Joseph Teran, and Eftychios Sifakis. Computing the singular value decomposition of 3x3 matrices with minimal branching and elementary floating point operations. Technical report, University of Musconsin-Madisoe Department of Computer Sciences, 2011.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li Fei-Fei, Joshua B Tenenbaum, and Daniel LK Yamins. Flexible neural representation for physics prediction. /806.08047, 2018.
Merlin Nimier-David, Delio Vicini, Tizian Zelmer, and Wenzel Jakob. Mitsuba 2: A retargetable forward and inverse renderer. Transactions on Graphics (Proceedings of SIGGRAPH Asia), 38 (6), November 2019, doi: 10. 10.1145/3355089,335566
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vetito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Barak A. Pearlmutter and Jeffrey Mark Siskind. Reverse-mode AD in functional framework: Lambda the ultimate backpropagator. ACM Transactions on Programming Languages and Sys- tems, 30(2):7:1-7:33 2008
Matt Pharr and William R Mark. ispe: A spmd compiler for high-performance cpu programming.
In Innovative Parallel Computing, pp. 1-13, 2012.
Jovan Popović, Steven M Seitz, Michael Erdmann, Zoran Popović, and Andrew Witkin. Interactive manipulation of rigid body simulations. In Proceedings of the 27th annual conference on Com- puter graphics and interactive techniques, pp. 209-2177 ACM Press/Addison Weenee-Weellyylyy Publishing Co., 2000.
Jonathan Bagan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: A language and compiler for optimizing parallelism, locality, and recom- putation in image processing pipelines. SIGPLAN Not. 48(6):519-530, jun 2013.

---

Published as a conference paper at ICLR 2020
Stéphane Redon, Abderrahmane Kheddar, and Sabine Coquillart. Fast continuous collision detec- tion between rigid bodies. In Computer graphics forum, volume 21, pp. 279-287, Wiley Online Library, 2002.
Connor Schenck and Dieter Fox. Spnets: Differentiable fluid dynamies for deep neural networks. arXiv preprint arXiv: (5Y1:880000094 2018.
Jos Stam. Stable fluids. In Siggraph, volume 99, pp. 121-128, 1999.
Andre Pradhana Tampubolon, Theodore Gast, Gergely Klár, Chuyuan Fu, Joseph Teran, Chenfanfu Jiang, and Ken Museth. Multi-speciies simulation of porous sand and water mixtures. ACM Transactions Graphics (TOG), 36(4):105, 2017.
Jean Utke, Uwe Naumann, Mike Fagan, Nathan Tallent, Michelle Strout, Patrick Heimbach, Chris Hill, and Carl Wunsch. Openad/f: A modular open-source tool for automatic differentiation of fortran codes. 34(4):18, 2008.
Jui-Hsien Wang, Ante Qu, Timothy R Langlois, and Doug James. Toward wave-based sound synthesis for computer animation. ACM Trans. Gruph., 37(4):109-1, 2018.
Richard Wei, Marc Rasi Dan Zheng, and Bart Chrzaszez. Differentiable program- ming prop-poop https: / / thub. com/ com/apple hesee///wiiiiei httlftt/bbbb/aaae /t//f///bb//aatte//6ocss Di DifferenttibleProgr weweeePrrrrrmininnggggo amming. md, 2019. Accessed: 2019-09-25.
R. Wengert, A simple automatic derivative evaluation program. Communications of the ACM, (8):463-464, aug 1964.
12

---

Published as a conference paper at ICLR 2020
A COMPARISON WITH EXISTING SYSTEMS
Workload differences between deep learning and differentiable physical simulation Existing differentiable programming tools for deep learning are typically centered around large data blobs. For example, in AlexNet, the second convolution layer has size 27 27 x 128 x 128. These tools usually provide users with both low-level operations such as tensor add and mul, and high-level op- erations such convolution. The bottleneck of typical deep-learning lap---eeanin- computer vision tasks are convolutions, so the provided high-level operations, with very high arithmetic intensity", can fully exploit hardware capability. However, the provided operations are "atoms" of these differ- entiable programming tools, and cannot be further customized. Users often have to use low-level operations compose their desired high-level operations. This introduces a lot of temporary buffers, and potentially excessive GPU kernel launches. As shown in Hu et al. (2019b), a pure TensorFlow implementation of complex physical simulator 132x slower than CUDA implementation, due to excessive GPU kernel launches and the lack of producer-consungg locality".
The table below compares DiffTaichi with existing tools for build differentiable physical simulators.
Table 3: Comparisons between DiffTaichi and other differentiable programming tools. Note that this table only discusses features related to differentiable physical simulation, and the other tools may not have been designed for this purpose. For example, PyTorch and TensorFlow are designed for classical deep learning tasks and have proven successful in their target domains. Also note that the XLA backend of TensorFlow and JIT feature of PyTorch allow them to fuse operators to some extent, but for simulation we want complete operator fusion within megakernels. "Swift" AD (Wei al., 2019) is partially implemented as of November 2019. "Julia" refers Innes al. (2019).
Feature
DifhTaichi
PyTorch
TensorFlow
Enoki
JAX
Halide
Julia
Swift
GPU Megakernels Imperative Scheme
Parallelism
Flexible Indexing
B DIFFERENTATING STRAIGHT-LINE TAICHI KERNELS USING SOURCE CODE TRANSFORM
Primal and adjoint kernels Recall that in DichTaichi, (primal) kernels are operators that take as input multiple tensors (e.g., X, Y) and output another set of tensors. Mathematically, kernel f has the form
f(Xo, X1, X,,) Yo, Y1, Y
Kernels usually execute uniform operations on these tensors. When comes differentiable pro- gramming, a loss function is defined on the final output tensors. The gradients of the loss function "I," with respect to each tensor are stored in adjoint tensors and computed via adjoint kernels.
The adjoint tensor of (primal) tensor Xisk is denoted as X i5k- Its entries are defined by ijj 0L/0Xijk- At high level, our automatic differentiation (AD) system transforms primal kernll into its adjoint form. Mathematically,
(primal) f(Xo, XI..., Xn)) = Yo, YI....,b. Ym
Reverse-Mode Automatic Differentiation
(adjoint) f"(Xo, X1, X,, Yo: YI ..... Y Xi Xi,,,, X
2FLOPs per byte loaded from/stored main memory.
'The CUDA kemnlss in Hu et al. (2019b) have much higher arithmetic intensity compared to the TensorFlow computational graph system. In other words, when implementing in CUDA immediate results cached in registers, while in TensorFlow they "cached" in main memory.
13

---

Published as a conference paper at ICLR 2020
Differentiating within kernels: The "make adjoint" pass (reverse-mode AD) After the preprocessing passes, which flatten branching and climinate mutable local variables, the "make adjoint" pass transforms forward evaluation (primal) kemnel into its gradient accumulation ("adjoint") kernel. takes straight-line code directly and operates on the hierarchical intermediate representation (IR) of Taichi* Multiple outer for loops allowed for the primal kernel. The Taichi compiler will distribute these parallel iterations onto CPU/GPU threads.
During the "make "make_adjoine adjoint" pass, for each SSA instruction, local adjoint variable will be allocated for gradient contribution accumulation. The compiler will traverse the statements in reverse order, and accumulate the gradients to the corresponding adjoint local variable.
For example, a ID array operation sin 2 has its IR representation as follows:
for i E range(0, 16, step 1) do %1 load x i
%2 = mul %1, %1
%3 (%2)
store y|i] = %3
end for
The above primal kernel will be transformed into the following adjoint kernel:
for i in range(O, 16, step 1) do
// adjoint variables
%ladj = alloca 0.0
%2adj alloca 0.0
%3adj = alloca 0.0
// original forward computation
%1 load x: {
%2 mul %1, %1
%3 sin( (%2)
// reverse accumulation
%4 load y _dj](j
%3adj += %4
%5 cos (%2)
%2adj += %3adj %5
%ladj
%1 %2adj
atomic add x __jddjli [i], %ladj
end for
Note that for clarity the transformed code is not strictly SSA here. The actual IR has more instruc- tions. A following simplification pass will simplify redundant instructions generated by the AD pass.
C COMPLEX KERNELS
Here we demonstrated how to use complex kernels to override the automatic differentiation system. We use singular value decomposition (SVD) of 3 x 3 matrices (M UEV* example. Fast SVD solvers used in physical simulation are often iterative, yet directly evaluate the gradient of this iterative process is likely numerically unstable. Suppose we McAdams al. (2011) the forward SVD solver, and use the method in Jiang (2015) (Section 2.1.1.2) to evalute the gradients, the complex kernels are used as follows:
'Taichi uses hierarchical static single assignment (SSA) intermediate representation (IR) as its internal program representation. The Taichi compiler applies multiple transform passes to lower and simplify the SSA IR in order to get high-performance binary code.
14

---

Published as a conference paper at ICLR 2020
# Do Singular Value Decomposition (SVD) matrices
@ti. kernel def seraaiie_ssdnuumtteeeaaaa
f1.f32):
for 1 in range(n):
input matrix_M(1)
far iter in range range(nuuniterrations))
iteratively solve SVD using McAdams et al. 2011 matrix_u[1] Satrix_Signai]]]
matrix_V[i]
complex kernel that wraps iterative kernel
@ti. cwwplex_keenneeeeeeeeeeeeeeeeeee
svd_
gti kernell
def svd__gadit svd_gradieett():
for range (n)
Implement, for example, section 2.1.1.2 of Jiang (2015)
complex kernel that is registered the svd_ forword complex kernel sti.complex.commmmmmmmmmmmll complex_kernel_eee kerddl forward) def syd_backward (ntarrtttnnntttttitooonnn iiteeaatioon::: i dit different Love SVD
svd_e wvd_gradient ()
D CHECKPOINTING
this section we demonstrate how to use checkpointing via complex kernels. The goal check- pointing to recomputation to save memory space. We demonstrate this using the diffapm example, whose simulation cycle consists of particle to grid transform (p2g), grid boundary con- ditions (grid_op), and grid particle transform (g2p). We assume the simulation has O(n) time steps.
D.1 RECOMPUTATION WITHIN TIME STEPS
naive implementation without checkpointing allocates (n) copied the simulation grid, which can cost a lot of memory space. Actually, if we recompute the grid states during the backward simulation time step by redoing p2g and grid_op, op, we can reused the grid states and allocate only one copy. This checkpointing optimization is demonstrated in the code below:
et cowplex. kernel def advance(s):
clear_gride compute_actuationn p2g(s) grid_op() gsp(s)
@ti heamplee__eenee__rdd(aaooe def advance_grad(S clear_grid()
p2g(s) grid_op() recompute the grid
(2p-grad(s) prid_op.grade
p2g. (ggggrad(s) compute_actee (r)aa)toonngra(((s)
D.2 SEGMENT- WISE RECOMPUTATION
Given simulation with O(n) time steps, if all simulation steps recorded, the space consumption On). This linear space consumption is sometimes too large for high-resolution simulations with long time horizon. Fortunately, reduce the space consumption using a segment-wise
15

---

Published as a conference paper at ICLR 2020
checkpointing trick: We split the simulation into segments of S steps, and in forward simulation store only the first simulation state in each segment. During backpropagation when we need the remaining simulation states in segment, recompute them based on the first state in that segment.
Note that if the segment size O(S). then we only need to store O(n/S) simulation steps for checkpoints and O(S) reusable simulation steps for backpropagation within segments. The total space consumption is O + n/S). Setting 5 = O(vn) reduces memory consumption from ((n) Ovyn). The time complexity remains (n)).
E
DETAILS ON 10 DIFFERENTIABLE SIMULATORS
E.1 DIFFERENTIABLE CONTINUUM MECHANICS FOR ELASTIC OBJECTS [di ffapm]
Figure 6: Controller optimization with our differentiable continuum simulators. Left: the 2D robot with four muscles. Middle: A 3D robot with 16 muscles and 30K particles crawling on the ground. (Reproduce: python3 (diftnppn/ifffmmm3d..pp/// Right: We couple the robot (30K particles) and the liquid simulator (13K particles), and optimize its open-loop controller in this difficult situa- tion. (iepoReeood: python3 liquid.py]
E.2 DIFFERENTIABLE LIQUID SIMULATOR [Liquid]
We follow the weakly compressible fluid model in Tampubolon et al. (2017) and implemented 3D differentiable liquid simulator within the [diffapmad]] framework. Our liquid simulation can be two-way coupled with elastic object simulation (Figure 6, right).
E.3 DIFFERENTIABLE INCOMPRESSIBLE FLUID SIMULATOR [smoke]
(a) smoke
(b) wave
Figure 7: (a): (Left to right) with an optimized initial smoke velocity field, the fluid changes its pattern to "Taichi" symbol. [Reproduce: python3 smoke taichi.py] (b): Unoptimized (top three) and optimized (bottom three) waves at time step 189, and 255. [Reproduce: python3 wave.py]
Backpropagating Through Pressure Projection We followed the baseline implementation in Autograd, and used 10 Jacobi iterations for pressure projection. Technically, 10 Jacobi iterations are not sufficient to make the velocity field fully divergence-free. However, in this example, it does decent job, and we are able successfully backpropagate through the unrolled Jacobi iterations.
In larger-scale simulations, 10 Jacobi iterations are likely not sufficient. Assuming the Poisson solve done by an iterative solver (e.g. multigrid preconditioned conjugate gradients, MGPCG) with multigrid levels and 50 conjugate gradient iterations, then automatic differentiation will likely not be able to provide gradients with sufficient numerical accuracy across this long iterative process. known to numerically
The accuracy is likely worse when conjugate gradients present, they drift as the number of iterations increases. In this case, the user can still use DiffTaichi to implement the forward MGPCG solver, while implementing the backward part of the Poisson solve manually,
16

---

Published as a conference paper at ICLR 2020
likely using adjoint methods (Errico, 1997). DiffTaichi provides "complex kernels" to override the built-in AD system, as shown Appendix
E.4 DIFFERENTIABLE HEIGIT FIELD SHALLOW WATER SIMULATOR [wave]
We adopt the equation in Wang et al. (2018) to model shallow water height field evolution:
where is the height of shallow water, the "speed of sound" and a is a damping coefficient. We use the út and i notations for the first and second order partial derivatives of w.r.t time t respectively.
Wang et al. (2018) used the finite different time-domaii (FDTD) method (Larsson & Thomée, 2008)
discretize Eqn. 1, yielding an update scheme:
Ut,i.j 144 + (e2 ca V u Pt-2,1.5 co.A(VV2))
where
u Augai Wrig+1 + Wr,if-11 1 + Mr 1,7 11,i-1j
Ax
We implemented this wave simulator DiffTaichi to simulate shallow water. We used a grid of resolution 128 x 128 and 256 time steps. The loss function is defined as
L A -
where T' is the final time step, and å is the target height field. 200 gradient descent iterations then used to optimize the initial height field. We set à to be the pattern "Taichi", and Fig. 7b shows the unoptimized and optimized
evolution.
We set the "Taichi" symbol as the target pattern. Fig. 7b shows the unoptimized and optimized final wave patterns. More details discretization is in Appendix E.
E.5 DIFFERENTIABLE PASS-SGRING SYSTEM [mass_spring
We extend the mass-spring system in the main text with ground collision and a NN controller. The time-of-impact fix implemented for improved gradients. The optimization goal is to maximize the distance moved forward with 2048 time steps. We designed three mass-spring robots as shown in Fig. 8 (left).
E.6 DISTERENTIABLE BILLIARD SIMULATOR [billiards]
A differentiable rigid body simulator is built for optimizing billiards strategy (Fig. 8, middle). We used forward Euler for the billiard ball motion and conservation momentum and kinetic energy for collision resolution.
Figure 8: Left: Three mass-spring robots. The and blue springs are actuated. two layer NN is used as controller. [Reproduce: python3 mass paas__sprnggpy [1/2/3] train]. Middle: Optimizing billiards. The optimizer adjusts the initial position and velocity of the white ball, so that the blue ball will reach the target destination (black dot). [Reproduce: python3 billiards.py' Right: Optimizing a robot walking. The rigid robot is controlled with NN controller and learned to walk in gradient descent iterations. [Reproduce: python3 _
17

---

Published as a conference paper at ICLR 2020
E.7 DIFFERENTIABLE RIGID BODY SIMULATOR [rigid_body]
worth noting that discontinuities can happen in rigid body collisions, and at countable number of discontinuities the objective function is non- differentiable. However, apart from these discontinuities, the process still differentiable almost everywhere. The situation of rigid body collision is somewhat similar to the "ReL.U" activation function in neural networks: at point x ReL.U is not differentiable (although continuous), yet is still widely adopted. The rigid body simulation cases more complex than ReLU, not only non-differentiable points, but also discontinuous points. Based on our experiments, in these impulse-based rigid body simulators, still find the gradients useful for optimization despite the discontinuities, especially with our time-of -mmeaacct fix.
Are rigid body collisions differentiable?
have
E.8 DIFFERENTIABE WATER RENDERER [water_renderer]
We implemented differentiable renderers to visualize the refracting water surfaces from We use finite differences to reconstruct the water surface models based on the input height field and refract camera rays to sample the images, using bilinear interpolation for meaningful gradients. To show our system works well with other differentiable programming systems, we use an adversarial optimization goal: fool VGG-16 into thinking that the refracted squirrel image is goldfish (Fig. 9).
[wave] initial water height field
Image Inpult
render_rendere]
Ripple
optimize:
VGG 16
water surface
Figure 9: This three-stage program (simulation, rendering, recognition) is end-to-end differentiable. Our optimized initial water height field evolves to form refraction pattern that perturbs the image into one that fools VGG16 (99.91% % goldfish). [Reproduce: python3 water paater_rendererpy/
E.9 DIFFERENTIABLE VOLUME RENDERER [volune_renderer]
We implemented basic volume renderer that simply uses ray marching (we ignore light, scattering, etc.) to integrate density field over each camera ray. In this task, we render number of target images from different viewpoints, with the camera rotated around the given volume. The goal is then optimize for the density field of the volume that would produce these target images: render candidate images from the same viewpoints and compute an L2 loss between them and the target images, before performing gradient descent the density field (Fig. 10). Essentially, this demonstrates how to use gradients to reconstruct 3D objects out of X-ray photos in a brute-force manner. Other approaches this task include algebraic reconstruction techniques (ART) (Gordon et al., 1970).
Figure 10: Volume rendering of bunny shaped density field. Left: (of the 7) target images. Right: optimized images of the middle bunny after iteration 2, 50, 100. [Reproduce: python3 volume renderer.py]
18

---

Published as a conference paper at ICLR 2020
E.10 DIFFERENTIABLE ELECTRIC FIELD SIMULATOR [electric]
Recall Coulomb's law: F k 21 f. In the right figure, there eight electrodes carrying static charge. The red ball also carries static charge. The controller, which is two-layer neural network, tries to manipulate the electrodes that the red ball follows the path of the blue ball. The bigger the electrode, the more positive charge it carries.
F FIXING GRADIENTS WITH TIME OF IMPACT AND CONTINUOUS COLLISION DETECTION
Here is naive time integrator in the mass-spring system example:
kernel
def advance(t: ti.132):
for range (n__enn_oooeeeeeee :bjoctes):
math.exp damping)
new_v s vft 1, dt gravity ti.Vector (fo.0, 1.01)
old_x x[t
f]
depth old_x [[1_x[1] ground_heigg height
depth
new_v[1]
# assuming a sticky ground (tnfinite coefficient of friction) new_v[o]
new_v[1]
tithout cansidering time of impact, assume the whole uses new_ new_x old_x dt new_y
v[t, i] xft, i] new_x
Implementing TOI in this system is relative straightforward:
gti. kernel
def advance_ toi(t: t1.132):
for i in range (n_ (n_objects)
math.exp
damping)
old_v
v[t 1, 1] + dt gravity * ti.Vector[[8.3 1.9])
old_x xIt 1, i]
new_x old_x dt old_v
toi 0.0
old_v old_v
if new_x[1] ground_hewee height and old_v[1] < -1e-4;
le-d safe guard ieportant for qumerical stability
toi -[old_x[1] ground_heiieeeee height) / old_v[1] # Compute the time of impact new_v ti, (i. 0.0])
Note that with of impact, de is divided into two parts, the first using old_v, and second part using new_v new_x old_x toi old_y (dt toi) A new_y
v[t, i] xft, i] new_x
In rigid body simulation, the implementation follows the same idea yet is slightly more complex. Please refer rigid_body py for more details.
G ADDITIONAL TIPS ON GRADIENT BEHAVIORS
Initialzzaton matters: flat lands and local minima in physical processes A trivial example of objective flat land in billiards. Without proper initialization, gradient descent will make no progress since gradients are (Fig. 11). Also note the local minimum near (-5, 0.03).
19

---

Published as a conference paper at ICLR 2020
Billiard Scene Objective
Billiard Scene Objective
0.06
0.06
0.05
0.05
0.04
0.01
0.02
0.22
-75
-30
-23
0
23
50
75
-20
-10
Angle velority
Ang vebocity
Most initial angles yield flat objective (final distance between the blue ball and black destination) of 0.065, since the white ball does not collide with any other balls and imposes no effect the pink ball via the chain reaction. Right: A zoomed-in view of the middle figure. The complex collisions lead to lot of local minimums. [Reproduce: python3 billiards.py 1.0/0.23]
Figure 11: Left:
Scanning initial velocity in the billiard example. Middle:
In mass_spaa jass__siing and rigid_boci once the robot falls down, gradient descent will quickly become trapped. A robot the ground will make further progress, no matter how it changes its controller. This leads to more non-trivial local minimum and zero gradient case.
Ideal physical models are only "ideal": discontinuities and singularities Real-world macro- scopic physical processes are usually continuous. However, building upon ideal physical models, even the forward physical simulation results can contain discontinuities. For example, body model with friction, changing the initial rotation of the box can lead to different corners hitting the ground first, and result in a discontinuity (Fig. 12). In electric and mass_spring, due to the / and + terms, when r + 0, gradients can be very inaccurate due to numerical precision issues. Note that d(1/r)/dr -1/r², and the gradient is more numerically problematic than the primal for small Safeguarding is critically important for gradient stability.
rigid
Rigid Body Simulation Discontinuity
coeff of friction=0 coeff of friction=1
0.5
Loss (m)
0.4
B
0.3
0.2
-1.0
-0.5
0.0
0.5
1.0
Initial Rotation Angle (degrees)
Friction in rigid body with collision is a common source of discontinuity. In this scene rigid body hits the ground. Slightly rotating the rigid body changes which corner (A/B) hits the ground first, and different normal/friction impulses will be applied to the rigid body. This leads to discontinuity in its final position (loss=final coordinate). [Reproduce: python3 rigid body discontinuity.py/ Please see our supplemental video for more details.
Figure 12:
20


Reviews:
[{"review_id": "ryeQ4y6htH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a programming language for building differentiable physics simulators. This is a very interesting goal, as differentiable systems are a crucial building block for many deep learning methods and similar optimization techniques. \n\nThe system presented by the authors is certainly impressive. Unfortunately, the paper itself covers a wide range of topics, and consists of an overview of the language with a programming tutorial, a collection of ten results, and a brief discussion of problems when computing gradients. \n\nThe core of the proposed work, the programming language seems to be quite powerful. However, it seems to be built on an existing system, which was published as a programming language for simulation in this years siggraph asia conference (Taichi: A language for high-performance computation on spatially sparse data structures. In SIGGRAPH Asia 2019 Technical Papers, pp. 201. ACM, 2019a). This ICLR submission seems to extend this system to build and provide gradient information automatically along with the simulation itself. There seem to be few technical challenges here, and many aspect discussed in section 2 are shared with the original simulation language.\n\nThe examples cover a nice range of cases, from simple mass spring systems and a rendering case to complex 3d simulations. Here, I was a bit surprised that the paper only compares to autograd, which has been succeeded by jax. The latter also provides a compiler backend to produce GPU code with gradients, and as such seems very closely related to the proposed language. From the submission, it's hard to say which version has advantages. The examples seem to be a sequence of demos of the language, rather than illustrating different technical challenges or improvements for a scientific conference. Or at least a discussion of these differences is currently missing in the text.\n\nSection four also mostly gives the impression of a loose discussion. The gradients for rigid body impacts are interesting, but seem relevant only for a subset of 2D examples shown in the paper. The discussion of gradient explosions is quite ad-hoc, and would be stronger with a more detailed analysis.\n\nThe submission as a whole aims for a very interesting direction, but I think the paper would benefit from focusing on a certain range of problems, such as the rigid body control cases, in conjunction with topics such as the improved gradients. Instead, the current version tries to combine this topic with a systems overview, a tutorial and loosely related discussions. Combined with the length of 10 pages, I think the work could use a revision rather than being accepted in its current form.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJgrncv9KB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nDifferentiable physics simulation is an important research area, especially for optimal control and reinforcement learning. While I am impressed by the large variety of examples demonstrated in the paper, I am leaning towards rejecting the paper because of its poor presentation. The paper only gives a simple and high-level example of the language (optimizing the rest length of springs that form a triangle), very brief descriptions of 10 examples and some discussions about the difficulty of computing useful gradients, but without any in-depth discussion how everything is implemented. This is not enough for an ICLR paper. For example, the paper does not answer some of the fundamental problems of differentiable physics. For example, collision and contact are inherently non-differentiable. How does the paper handle it in the examples of locomotion and billiards (Figure 4)? In addition, how does the paper back-propagate the gradient through the incompressibility conditions (Poisson solve) of fluid simulation? \n\nHere is my suggestions how to improve the writing. There are several ways to write the paper, with different emphasis. If this paper is more about introducing a new programming language, Appendix B Compiler Design and Implementation would be important and should be moved to main text. If the paper want to emphasize how to handle the non-differentiable cases of the simulation, then detailed derivations of contact, collision, and linear/nonlinear solving (due to incompressibility conditions or implicit integrators) should be presented. If the paper would like to demonstrate how differentiable physics simulation can help with controller optimization, then two to three examples, such as the locomotion control for soft bodies or rigid bodies, should be analyzed in far more details, and compared with traditional method without differentiable simulation. It is good to focus on one of the above points, based on the venue that this paper is submitted to. Currently, the paper is trying to touch all three. But due to the page limit, it is not thorough, or detailed in any one of them.\n\n-------------------------Update after rebuttal------------------------------\nThank you for the revision of the paper and the additional comparisons with Jax. The revised version reads much better. The response and the revision addressed most of my concerns. Thus, I raised my rating to weak accept.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "H1loP67jqH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "*Summary*\nThis 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.\n\n*Rating*\nThe paper is interesting and easy to read. While some part of the underlying functionality of DiffSim is directly derived from previous work (Taichi), the paper does describe a non-trivial contribution.\n\nI lack the background to comment constructively about expectations for these simulations or the fidelity of the methods in this paper. What evidence can you offer regarding the physical fidelity achievable and how that relates to issues of scalability, gradient behavior, size of time steps, code complexity, etc.? For a sense of context, what might be needed to simulate a 7 DoF robotic arm and learn a controller that would reasonably transfer to a real robot?\n\nOverall, I'm optimistic about this paper, and would tend to vote for acceptance.\n\n*Notes*\npg3: define k (spring stiffness?)\npg4: what is the value of 'mass' for this simulation?\nFig 8: what is the x-axis in the two right plots? initial height?\nFig 10: right plot title should probably be \"Gradient Explosion with Damping\"", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a language for optimizing through physical simulations. The reviewers had a number of concerns related to paper organization and insufficient comparisons to related work (jax). During the discussion phase, the authors significantly updated their paper and ran additional experiments, leading to a much stronger paper.
Decision: Accept (Poster)