PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffTaichi: Differentiable Programming for Physical Simulation

Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Ragan-Kelley, Frédo Durand

arXiv:1910.00935v3cs.LGcs.GRphysics.comp-phstat.ML

TL;DR

Existing tools make it difficult to implement high-performance differentiable physical simulators, despite their importance in machine learning systems. DiffTaichi addresses this with a tailored differentiable programming language and automatic differentiation system, supporting 10 physical simulators and showing shorter, faster implementations than existing systems.

  • Problem

    Existing tools make it difficult to implement differentiable physical simulators with high performance, although such simulators are important components in machine learning systems.

  • Method

    DiffTaichi combines an imperative physical-simulation language with megakernels, source-code transformations, and a tailored two-scale automatic differentiation system.

  • Results

    DiffTaichi was used to quickly implement and automatically differentiate 10 physical simulators, while an elastic-object implementation was 4.2× shorter than manual gradients, almost as fast, and 188× faster than TensorFlow.

  • Takeaways & Limitations

    DiffTaichi lowers the barrier to developing high-performance differentiable physical simulators for machine learning and robotics research.

  • Takeaways & Limitations

    Differentiating a simulator does not always yield useful gradients of the physical system, even when forward simulation is accurate.

Abstract

from arXiv · show

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.

1 INTRODUCTION

DiffTaichi targets the difficulty of implementing high-performance differentiable physical simulators by combining tailored language features with automatic differentiation. The system supports complex simulators and neural-network controllers across 10 physical-simulation domains.

  • Differentiable physical simulators matter because they can make controller optimization converge one to four orders of magnitude faster than model-free reinforcement learning.Their inner-loop performance is therefore important, while existing tools make high-performance implementation difficult.
  • DiffTaichi is a differentiable programming language for high-performance physical simulations on CPU and GPU.
  • Megakernels: Megakernels fuse multiple computation stages, then apply source-code transformation and just-in-time compilation to produce efficient differentiated kernels.The fused kernels have higher arithmetic intensity than linear-algebra operators in TensorFlow and PyTorch for physical simulation tasks.
  • Imperative Parallel Programming: Imperative parallel programming provides loops and control flow suited to collisions, boundary conditions, and iterative solvers, while simplifying ports of existing simulation code.
  • Flexible Indexing: Flexible indexing directly expresses numerical stencils, particle-grid interactions, and partial global-array updates without unintuitive scatter/gather operations.Explicit indexing also enables compiler access optimizations.
  • DiffTaichi enabled rapid automatic differentiation of 10 simulators spanning rigid bodies, deformable objects, and fluids, including systems with neural-network controllers.

2 BACKGROUND: THE TAICHI PROGRAMMING LANGUAGE

Taichi provides an imperative, embedded programming foundation in which computation is separated from data structures and parallel loops are first-class constructs. The DiffTaichi frontend compiles statically typed, parallel, differentiable code into Taichi intermediate representation and forward and backward executables for simulation programs.

  • Taichi separates computation from data structures, allowing data layouts to change while array-like indexing remains available to programmers.The compiler uses both data-structure and algorithm information for performance optimization.
  • Taichi's first-class parallel-for loops and data-layout flexibility support high-performance physical-simulation programming.
  • The DiffTaichi frontend is embedded in Python and compiles code into Taichi intermediate representation and automatically differentiated forward and backward executables.Unlike Python, DiffTaichi is compiled, statically typed, parallel, and differentiable.
  • The tutorial demonstrates a three-spring, three-mass-point simulator using global tensors for loss, state, forces, and spring properties.
  • The mass-spring model computes spring forces with Hooke's law and scatters those forces to mass points through a parallel kernel.
  • Forward simulation repeatedly applies spring forces and semi-implicit Euler integration across time steps, updating velocity and position with damping.

3 AUTOMATICALLY DIFFERENTIATING PHYSICAL SIMULATORS IN TAICHI

DiffTaichi differentiates physical simulators through a two-scale system that transforms kernels locally and replays recorded kernel launches globally. Its design preserves parallel execution while imposing data-access rules and supporting optimization workflows.

  • Two-scale automatic differentiation: DiffTaichi combines source code transformation within kernels with a lightweight tape for end-to-end simulation differentiation.The tape records kernel names, function pointers, and scalar parameters, then replays gradient kernels in reverse order.
  • Global data access: Automatic differentiation assumes repeated global-tensor writes use atomic accumulation and that reads occur only after accumulation is complete.These rules prevent kernels from overwriting one another’s outputs and make adjoint computation well-defined.
  • Practical constraints: Forward simulators may need to record complete histories of state variables, increasing memory use; checkpointing can alleviate this cost.The mass-spring example records the whole history of position and velocity rather than only their latest values.
  • Local automatic differentiation: Branches are converted to differentiable select operations, mutable local variables are eliminated, and straight-line code is differentiated with reverse-mode AD.These preprocessing steps simplify the intermediate representation before the standard source code transformation.
  • Parallelism and loop handling: Parallel loop structures are preserved during differentiation, while backward atomic operations accumulate gradient contributions into adjoint tensors.Non-parallel loops are reversed during AD; loops carrying mutating local variables are unsupported.
  • Optimization example: DiffTaichi supports gradient-based optimization workflows, including spring rest-length optimization targeting a triangle area of 0.2.In the example, 200 gradient-descent iterations change the spring lengths from [0.1, 0.1, 0.14] to [0.600, 0.600, 0.529].

4 EVALUATION

DiffTaichi was evaluated on 10 physical simulators spanning continuum, fluid, and rigid-body settings. The evaluations report substantially shorter implementations, faster execution than array-based systems, and improved rigid-body gradient quality with continuous time of impact.

  • Evaluation: 10 physical simulators spanning large-scale continuum and small-scale rigid-body simulations were evaluated with reproducible scripts.The paper focuses discussion on three simulators and provides additional simulator details in Appendix E.
  • Performance and Productivity: 4.2× shorter and almost as fast as the manual-gradient CUDA implementation, while 1.7× shorter and 188× faster than TensorFlow for diffmpm.The benchmark uses 2D simulations with 6.4K particles on an NVIDIA GTX 1080 Ti GPU.
  • Smoke Simulator: 10 seconds total runtime, including 2 seconds of JIT compilation, versus 2 minutes for JAX JIT compilation in the smoke benchmark.The grid-based simulator was intentionally simplified to suit traditional array-based programs.
  • Rigid-Body Gradients: A naive rigid-body time integrator produces a gradient of 1 instead of −1 for final height with respect to initial height, regardless of how small ∆t is.The forward result can remain reasonable because time discretization itself is not differentiated by the compiler.
  • Rigid-Body Gradients: Adding continuous time of impact barely changes forward simulation but corrects gradients and significantly improves controller-optimization gradient quality.Figure 5 reports improved gradient and optimization quality with experiments repeated five times.
  • Evaluation: Differentiating a simulator can yield unusable physical gradients even when its forward simulation works well.The paper identifies this as a broader gradient issue and discusses additional cases in Appendix G.

5 RELATED WORK

Related work spans differentiable programming for deep neural networks, general-purpose code transformation, array-based parallel systems, and differentiable physical simulators. DiffTaichi differs by targeting imperative physical simulation with flexible indexing and preserved parallelism.

  • Differentiable Programming: Deep-learning autodiff libraries such as Theano, TensorFlow, and PyTorch primarily target differentiable neural-network operations.Physical simulation instead requires complex and customizable operations because of computational irregularity.
  • Differentiable Programming: General scalar-code transformation and annotated differentiable languages support general programming languages but provide limited parallelism.These approaches represent an earlier and emerging line of automatic-differentiation research, respectively.
  • Differentiable Programming: Array-based systems such as Halide, Autograd, JAX, and Enoki exploit parallelism, whereas DiffTaichi uses imperative programming with flexible indexing.The design is intended to ease porting existing physical simulation algorithms.
  • Differentiable Physical Simulators: Prior differentiable physical simulators used neural-network approximations, Theano, PyTorch, or custom CUDA kernels for robotics, fluids, and rigid bodies.The cited systems cover several simulation modalities and implementation strategies.

6 CONCLUSION

DiffTaichi is presented as a language for high-performance differentiable physical simulation, combining tailored automatic differentiation with implementations across 10 simulators. The authors report performance and productivity gains and aim to lower barriers for future research.

  • Conclusion: DiffTaichi combines support for megakernels, imperative programming, and flexible indexing in a tailored two-scale automatic-differentiation system.These design choices motivated the language and its simulator implementations.
  • Conclusion: 10 simulators were built and integrated into deep neural networks, demonstrating reported performance and productivity advantages over existing systems.The conclusion frames these results as evidence for the language’s utility in differentiable physical simulation.
  • Conclusion: The authors hope DiffTaichi lowers the barrier to future differentiable physical-simulation research in machine learning and robotics.

A COMPARISON WITH EXISTING SYSTEMS

The comparison frames differentiable physical simulation as a workload distinct from conventional deep learning, requiring support for irregular computation and efficient fusion. Table 3 compares systems specifically on features relevant to that workload.

  • Workload Differences: Deep-learning tools process large data blobs with low-level tensor operations and high-level operators, while physical simulation has different computational bottlenecks.The passage uses convolution-heavy computer vision as an example of the deep-learning workload.
  • System Comparison: Table 3 compares DiffTaichi with existing differentiable programming tools using features related specifically to differentiable physical simulation.The comparison does not evaluate whether other tools succeed in their original target domains.
  • System Comparison: TensorFlow XLA and PyTorch JIT can fuse operators partially, but the comparison emphasizes complete operator fusion within megakernels for simulation.The table notes that PyTorch and TensorFlow were designed for classical deep-learning tasks.

B DIFFERENTATING STRAIGHT-LINE TAICHI KERNELS USING SOURCE CODE TRANSFORM

DiffTaichi differentiates imperative parallel simulation kernels by transforming primal code into adjoint kernels and replaying gradients through the simulation structure. Its checkpointing and complex-kernel mechanisms address memory use and numerically delicate operations while supporting coupled simulators.

  • Primal kernels map input tensors to output tensors, while adjoint tensors store loss gradients and adjoint kernels compute them.
  • Reverse-mode source transformation traverses kernel statements backward, accumulating local adjoints and atomically adding gradients to input arrays.
  • Complex kernels let users override built-in differentiation, including manually specified gradients for iterative SVD procedures whose direct differentiation may be numerically unstable.
  • Checkpointing recomputes intermediate grid states during backward simulation instead of storing every state, reducing memory for long simulations.
  • O(S + n/S) space becomes O(√n) when S = O(√n), while time complexity remains O(n).
  • DiffTaichi supports differentiable continuum simulators, including a liquid simulator that can be two-way coupled with elastic-object simulation.

E.3 DIFFERENTIABLE INCOMPRESSIBLE FLUID SIMULATOR [smoke]

The wave and smoke examples use differentiable fluid formulations to optimize initial fields toward target patterns. The section also identifies accuracy limits for differentiating long iterative pressure solves and discontinuities in rigid-body cases.

  • Smoke simulator: Gradient descent on an initial smoke velocity field changes the fluid pattern toward a target image using semi-Lagrangian advection and implicit pressure projection.
  • Pressure projection: Ten Jacobi iterations are sufficient in this example to backpropagate through pressure projection, but they do not make the velocity field fully divergence-free.
  • Pressure projection: In larger simulations, automatic differentiation through long iterative Poisson solves may produce insufficiently accurate gradients, motivating manually implemented backward solves via complex kernels.
  • Wave simulator: The shallow-water simulator uses a wave equation with height, speed-of-sound, and damping terms, discretized by the finite difference time-domain method.
  • Wave simulator: A 128 × 128 grid and 256 time steps model shallow-water evolution, while 200 gradient-descent iterations optimize the initial height field toward the Taichi pattern.
  • Rigid-body simulator: Rigid-body collisions are non-differentiable at countably many discontinuities, yet gradients remain useful for optimization in experiments, especially with the time-of-impact fix.

E.8 DIFFERENTIABLE WATER RENDERER [water_renderer]

The water renderer makes wave simulation, surface rendering, and image recognition end-to-end differentiable. Optimizing the initial water height field produces a refraction pattern that fools VGG16.

  • Finite differences reconstruct the water surface from a height field, while bilinear interpolation provides meaningful gradients for refracted camera-ray sampling.
  • The three-stage pipeline—simulation, rendering, and recognition—is end-to-end differentiable.
  • The optimized initial water height field creates a refraction pattern that makes VGG16 classify the image as goldfish with 99.91% confidence.

E.9 DIFFERENTIABLE VOLUME RENDERER [volume_renderer]

The volume-rendering example optimizes a density field from multi-view target images using ray marching and an L2 loss. The accompanying simulator examples illustrate differentiable control and collision-handling implementations.

  • Volume renderer: A basic ray-marching volume renderer integrates density along camera rays while ignoring light and scattering.
  • Volume renderer: The density field is optimized by matching candidate renders from multiple viewpoints to target images with an L2 loss.
  • Volume renderer: Figure 10 compares target views with optimized bunny images after iterations 2, 50, and 100.
  • Collision handling: The time-of-impact implementation splits each time step into pre-impact motion using old velocity and post-impact motion using new velocity.

G ADDITIONAL TIPS ON GRADIENT BEHAVIORS

Gradient-based physical optimization can fail when objectives are flat, locally minimized, discontinuous, or numerically unstable. These behaviors arise from collision outcomes, poor initialization, falling robots, and singular physical models.

  • Zero gradients on flat billiards objectives prevent gradient descent from making progress without proper initialization.Most initial angles produce a flat objective because the white ball does not collide and trigger the chain reaction; a local minimum also occurs near (−5, 0.03).
  • Complex billiards collisions create many local minima, making the objective difficult to optimize even beyond flat regions.
  • Falling robots in mass_spring and rigid_body simulations create non-trivial local minima where controller changes produce no further progress.
  • As r →0 in electric and mass_spring models, gradients of inverse-distance terms become inaccurate because numerical precision problems worsen relative to the primal computation.Safeguarding r is critically important for gradient stability.
  • Discontinuous collision responses can result when slight rotation changes which rigid-body corner hits the ground first.Different normal and friction impulses then produce a discontinuity in the final position, whose y coordinate is the loss.
Loading 1910.00935v3…