DiffTaichi: Differentiable Programming for Physical Simulation

TL;DR

DiffTaichi presents a differentiable programming language built on Taichi to enable high-performance, end-to-end gradient computation for physical simulators. It combines megakernel fusion, imperative parallel programming, and flexible indexing with a two-scale AD approach: within-kernel differentiations via source-code transforms and end-to-end gradients using a lightweight tape. The framework demonstrates substantial productivity and speed gains across ten simulators, including elastic continua, incompressible fluids, and rigid bodies, often outperforming hand-tuned CUDA and mainstream autodiff frameworks. A key contribution is the robust handling of gradient computation in irregular, collision-based simulations through TOI-based techniques and complex-kernel overrides. Overall, DiffTaichi lowers the barrier to building differentiable physics engines suitable for neural-control loops and optimization tasks in robotics and related domains.

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.

Figures (16)

• Figure 1: Left: Our language allows us to seamlessly integrate a neural network (NN) controller and a physical simulation module, and update the weights of the controller or the initial state parameterization (blue). Our simulations typically have $512\sim 2048$ time steps, and each time step has up to one thousand parallel operations. Right: 10 differentiable simulators built with DiffTaichi.
• Figure 2: Left: The DiffTaichi system. We reuse some infrastructure (white boxes) from Taichi, while the blue boxes are our extensions for differentiable programming. Right: The tape records kernel launches and replays the gradient kernels in reverse order during backpropagation.
• Figure 3: Simple IR preprocessing before running the AD source code transform (left to right). Demonstrated in . The actual Taichi IR is often more complex. Containing loops are ignored.
• Figure 4: How gradients can go wrong with naive time integrators. For clarity we use a large $\Delta t$ here. Left: Since collision detection only happens at multiples of $\Delta t$ ($2\Delta t$ in this case), lowering the initial position of the ball (light blue) leads to a 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.9\Delta t$ (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 $\Delta t$ is small), the naive time integrator gives an incorrect gradient of $1$, but adding TOI yields the correct gradient. Please see our https://youtu.be/Z1xvAZve9aE for a better demonstration. [Reproduce: python3 rigid_body_toi.py]Right: When zooming in, the loss of the naive integrator is decreasing, and the saw-tooth pattern explains the positive gradients. [Reproduce: python3 rigid_body_toi.py zoom]
• Figure 5: Adding TOI greatly improves gradient and optimization quality. Each experiment is repeated five times. [Reproduce: python3 [mass_spring/rigid_body.py] [1/2] plot && python3 plot_losses.py]