可微编程-自上而下的产品形态 6 DiffTachi：针对物理模拟的可微编程

论文地址： DiffTachi：针对物理模拟的可微编程

说明文档： DiffTachi

范例: DiffTachi

名词解析 ：

AD： 自动微分

Tape: tensorflow 提供 tf.GradientTape api来实现自动求导功能。只要在 tf.GradientTape() 上下文中执行的操作，都会被记录与“tape”中，然后 tensorflow 使用反向自动微分来计算相关操作的梯度。

注意：当前框架没有使用Tensorflow，但是对tape的使用和定义都是等价的。

特性 ：

与生成不可变输出缓冲区的TensorFlow等工具不同，太极中采用的命令式编程范式允许程序员自由修改全局张量。为了使自动微分在这种情况下得到很好的定义，我们对可微程序的太极程序作如下假设：

全局数据访问规则：

如果一个全局张量元素被多次写入，那么从第二次写入开始，写入必须以原子加法（“累加”的形式出现，使用ti.atomic_add或简单地+=）。

在完成全局张量元素的累积之前，不会对其进行读取访问。

内核简单性规则：内核体由多个简单嵌套的for循环组成。一、 例如，每个for循环可以只包含一个（嵌套的）for循环（没有其他语句），或者包含一组没有循环的语句。

例子：

1,介绍

可微物理仿真器是机器学习系统的有效组成部分。例如，de Avila Belbute Peres等人。（2018a）和Hu等人。（2019b）已经表明，带可微模拟器的控制器优化收敛速度比无模型快一到四个数量级强化学习算法。可微物理模拟器在这些应用程序的内环中的存在使得它们的性能至关重要。不幸的是，使用现有的工具很难实现高性能的模拟器。

提出了一种新的可微程序设计语言DiffTaichi，用于CPU和GPU上的高性能物理仿真。它基于太极编程语言（Hu等人，2019a）。DiffTaichi自动识别系统是为适应语言的主要特点而设计的物理模拟所需，但在现有的可微编程工具中经常缺少，具体如下：

Megakernels我们的语言使用一种“megakernel”方法，允许程序员自然地将多个计算阶段融合到一个内核中，然后使用源代码转换和即时编译对其进行区分。与线性代数算子的比较TensorFlow（Abadi et al.，2016）和PyTorch（Paszke et al.，2017）DiffTaichi核具有更高的算术强度，因此对于物理模拟任务更有效。

命令式并行编程与现代深度学习中流行的函数数组编程语言（Bergstra等人，2010；Abadi等人，2016；Li等人，2018b）相比，大多数传统的物理模拟程序是用命令式语言编写的，如Fortran和C++。DiffTaichi同样采用命令式方法。该语言提供并行循环和控制流（如“if”语句），它们是物理模拟中广泛使用的结构：

它们简化了常见的任务，如处理碰撞、计算边界条件和构建迭代求解器。使用命令式样式可以更容易地将现有的物理模拟代码移植到DiffTaichi。

灵活引用现有的并行可微编程设计系统在相同形状的数组上提供了元素操作，例如c[i，j]=a[i，j]+b[i，j]。然而，许多物理模拟操作，如数值模板和粒子网格的相互作用并不是按元素进行的。常用的模拟模式如y[p[i]*2，j]=x[q[i+j]]只能表示在这些现有的系统中，存在着非直观的分散/聚集操作，这些操作不仅效率低下，而且很难开发和维护。另一方面，在DiffTaichi中，程序员通过任意索引直接操作数组元素，从而允许全局数组和使这些常见的模拟模式自然地表现出来。显式索引语法还使编译器易于执行访问优化（Hu等人，2019a）。这三个要求促使我们设计了一个定制的双尺度自动微分系统，这使得DiffTaichi特别适合于开发复杂的、高性能的、可区分的物理模拟器，可能有神经网络控制器（图1，左）。使用我们的语言，我们能够快速实现并自动区分10个物理模拟器1，覆盖刚体、可变形物体和流体（图1，右）。在附录A中对difftaichi和其他可微编程工具进行了全面的比较。

2背景：太极编程语言

DiffTaichi基于太极编程语言（Hu等人，2019a）。Taichi是一个嵌入在C++ 14中的命令式编程语言。它在现代硬件上提供高性能和高生产率。太极拳区别于其他命令式程序的关键设计明文如C++和CUDA是从数据结构中解耦的计算。这使得程序员可以轻松地在不同的数据布局之间切换，并使用索引（即x[i，j，k]）访问数据结构，就好像它们是普通的密集数组一样，而不管底层布局如何。

然后，太极编译器获取数据结构和算法信息来应用性能优化。太极拳提供“并行for”循环作为一级构造。这些设计使得太极特别适合于编写高性能的物理模拟器。更多细节，读者可参考Hu等人。（2019a年）。DiffTaichi语言前端嵌入到Python中，并编译Python AST transformer区分太极代码到太极中间表示（IR）。与Python不同，DiffTaichi语言是编译的、静态类型的、并行的和可微的。我们扩展了Taichi编译器以进一步编译并自动将生成的Taichi IR区分为正向和反向可执行文件。

我们使用一个质量弹簧模拟器来演示该语言，它有三个弹簧和三个质量点，如右图所示。在本节中，我们将介绍前向模拟器使用太极的扩散式太极前端，这是一种更容易使用的太极包装C++ 14前端。

首先分配全局变量我们分配一组全局张量来存储模拟状态。这些张量包括float32类型的标量损失，2D张量x，v，大小步长×núu弹簧和32x2型浮子，以及用于弹簧特性的1D尺寸núu弹簧阵列：弹簧ú锚úa（int32）、弹簧锚固（int32）、弹簧长度（浮动32）。定义核质量弹簧系统用胡克定律建模

其中k是弹簧刚度，F是弹簧力，xa和xb是两个质点的位置，l0是剩余长度。以下内核在所有弹簧上循环，并将力分散到质量点：

对于每个粒子i，我们使用带阻尼的半隐式Euler时间积分：

当

分别是粒子I在时间步t处的速度、位置和质量。a为阻尼系数，核如下：

使用这些组件组装正向模拟器，我们定义正向时间积分：

3,太极运动模拟器的三种自动判识方法

DiffTaichi的自动微分（AD）系统的主要目标是自动生成梯度模拟器，对传统的前向模拟器进行最小的代码修改。

设计决策源代码转换（SCT）（Griewank&Walther，2008）和跟踪（Wengert，1964）是设计广告系统时的常见选择。在我们的设置中，使用SCT来区分一个具有数千个时间步的整个模拟器，会导致高性能但较差的结果

灵活性强，编译时间长。另一方面，由于在反向传播过程中没有保留“巨核”结构，单纯地采用跟踪提供了灵活性，但性能较差。为了获得性能和灵活性，我们开发了一个双尺度自动区分系统（图2）：我们使用SCT在内核内进行区分，并使用一个轻量级Tape，它只存储用于端到端模拟区分的函数指针和参数。全局张量是梯度计算的自然检查点。

与生成不可变输出缓冲区的函数式编程语言不同，命令式编程允许程序员自由修改全局张量。为了使自动微分在该设置下得到很好的定义，我们对im操作核做了以下假设：

全局数据访问规则：

1） 如果一个全局张量元素被多次写入，那么从第二次写入开始，写入必须以原子加法（“累加”）的形式出现。

2） 在全局张量元素的累积完成之前，不会对其进行读取访问。

在前向模拟器中，程序员可能会做出细微的更改以满足规则。例如，在质量弹簧模拟示例中，我们记录了x和v的整个历史，而不是只保留最新的值。由此引起的内存消耗问题可以通过检查点来缓解，如后面附录D所述。有了这些假设，内核将不会覆盖彼此的输出，AD的目标很明确：给定一个原始内核f作为输入

并且输出为(或累加为)

,生成的梯度（伴随）核f*应作为输入

和

并累积梯度贡献为

，其中每一个

都是

的adjoint.

伴随张量的存储控制用户可以使用太极数据结构描述语言（Hu等人，2019a）指定伴随张量的存储，就好像它们是原始张量一样。我们还提供ti.root.lazy_grad（）来自动将伴随张量放置在它们的布局之后初选。

3.1本地AD：使用源代码转换区分太极kernel

一个典型的太极核心由多个水平的for循环和一个身体模块组成。为了使以后的AD更容易，我们引入两个基本的代码转换来简化循环体，如下所述。

图3：运行AD源代码转换（从左到右）之前的简单红外预处理。在C++中演示。真正的太极IR往往更复杂。将忽略包含循环。物理模拟分支中的展平分支是常见的，例如，在实现绑定条件和碰撞时。为了简化反模式AD pass，我们首先用三元运算符select（cond，value_if_true，value_if_false）替换“if”语句，它的梯度是清晰的定义（图3，中间）。这是程序矢量化中常见的转换（例如Karrenberg&Hack（2011）；Phar&Mark（2012））。

消除可变局部变量去掉分支后，我们最终得到直线回路体。为了进一步简化IR并使过程真正的单一赋值，我们应用一系列局部变量存储转发变换，直到可变局部变量可以完全消除（图3，右）。在这两个自定义的IR简化转换之后，DiffTaichi只需要区分不带可变变量的直线代码，这是通过使用标准源代码转换的反向模式AD实现（Griewank&Walther，2008）。有关此转换的详细信息，请参见附录B。

物理模拟中的大多数环都是并行环，在模拟过程中我们保留了并行环的结构。对于未显式标记为并行的循环，我们在AD转换期间反转循环顺序。我们不支持携带变异局部变量的循环，因为这样会需要复杂而昂贵的运行时堆栈来维护局部变量的历史记录。相反，用户被指示使用满足全局数据访问规则的全局变量。

并行性和线程安全性对于正向模拟，我们继承了太极拳的“并行For”结构，将每个循环迭代映射到CPU/GPU线程上。程序员使用原子操作来保证线程安全。我们的系统可以自动区分这些原子操作。梯度后向kernel的贡献通过原子加和累积到伴随张量。

3.2全局AD：使用轻型Tape进行端到端反向传播

我们构造了一个内核执行的Tape（图2，右），这样梯度内核可以按相反的顺序重放。Tape非常轻：由于中间结果存储在全局张量中，因此在正向模拟过程中，Tape只记录内核名称和（标量）输入参数，而不像其他可微函数数组系统，所有中间缓冲区都必须由Tape记录。无论何时启动DiffTaichi内核，我们都会将内核函数指针和参数附加到Tape。在计算梯度时，我们会遍历反转的Tape，然后使用记录的参数调用渐变内核。注意，DiffTaichi AD是根据输入全局张量而不是输入参数来计算梯度。

现在，我们重新讨论了质量弹簧的学习/优化问题，并使其具有可微性。假设目标是优化弹簧的剩余长度，以便在模拟结束时三个弹簧形成的三角形区域变为0.2。我们首先定义损失函数：

程序员使用ti.Tape来存储正向内核启动。它自动地反向重放这些核的梯度，以便反向传播。最初，弹簧的长度为[0.1,0.1,0.14]，优化后，其余长度为[0.600,0.600,0.529]。这意味着弹簧将根据胡克定律展开三角形并形成一个更大的三角形：[复制：mass_spring_simple.py]

复杂的内核有时用户可能想要覆盖编译器提供的渐变。例如，当使用迭代求解器对三维奇异值分解进行微分时，最好使用手动设计的SVD导数子例程以获得更好的稳定性。我们再提供两个decorator ti.complex_kernel和ti.complex_kernel_grad来覆盖默认的自动区分，如附录C所述。除了自定义渐变外，还可以使用复杂内核来实现检查点，如附录D所述。

4评估

我们在10个不同的物理模拟器上评估扩散太极拳，涵盖大规模连续体和小型刚体模拟。所有结果都可以用提供的脚本进行复制。动态/优化过程在补充视频中可视化。在本节中，我们将重点讨论三个模拟器。更多关于模拟器的细节见附录E。

4.1弹性物体的可微连续介质力学

首先，我们建立了一个适用于软机器人应用的可微连续模拟。物理系统受动量和质量守恒的控制，即。

0 我们遵循ChainQueen的实现（Hu等人，2019b）并使用移动最小二乘材料点法（Hu等人，2018）来模拟系统。我们能够很容易地将最初的CUDA模拟器翻译成DiffTaichi语法。使用这个模拟器和开环控制器，

我们可以很容易地训练一个软机器人向前移动（图1，diffmpm）。性能和生产率与（Hu等人，2019b）中的手动渐变实现相比，在DiffTaichi中获取渐变是毫不费力的。因此，DiffTaichi实现是代码行数缩短了4.2倍，运行速度几乎与TensorFlow相同；与TensorFlow相比，DiffTaichi代码分别缩短了1.7倍和188倍（表1）。由于tf.gather nd/scatternd和数组转置和广播的大量使用，Tensorflow实现是冗长的。表1:NVIDIA GTX 1080 Ti GPU上的diffmpm性能比较。我们使用6.4K粒子在2D中进行基准测试。对于代码行，我们只包括基本实现，不包括样板代码。[复制：python3 diffmpm_benchmark.py]

4.2可微不可压缩流体模拟器[烟雾]

我们用半拉格朗日平流实现了烟雾模拟器（图1，烟雾）（Stam，1999）以及内隐压力投影，以Autograd为例（Maclaurin等人，2015）。通过对初始速度场进行梯度下降优化，我们能够找到一个速度场，该速度场将流体的模式改变为目标图像（附录中的图7a）。我们比较一下表现我们的系统对抗表2中的PyTorch、Autograd和JAX。注意，作为Autograd库中的一个例子，这个基于网格的模拟器被有意地简化，以适应传统的基于数组的程序。例如，使用周期性边界条件，以便Autograd可以表示它使用numpy.roll，没有任何分支。尽管如此，太极拳的性能仍然高于这些基于阵列的系统。整个程序在GPU上运行DiffTaichi需要10秒，在JIT上运行需要2秒。JAX-JIT编译需要2分钟。

4.3可微刚体模拟器[刚体]

我们建立了一个基于脉冲的（Catto，2009）可微刚体模拟器（图1，刚体），用于优化机器人控制器。该模拟器支持刚体碰撞和摩擦、弹簧力、关节和驱动。模拟是端到端可微的，除了不连续性。有趣的是，尽管前向模拟器工作得很好，但由于刚体碰撞，用DiffTaichi对其进行简单的区分会导致完全错误的梯度。我们在下面讨论这个问题的原因和解决方法。改进碰撞梯度考虑图4（左）中的刚性球示例，其中刚性球与无摩擦地面碰撞。忽略重力，由于动能守恒，即使在弹性碰撞后，球仍保持恒定速度。在正向模拟中，使用一个小的∏t通常会得到一个合理的结果，就像在许多物理模拟器中所做的那样。降低初始球的高度将增加最终球的高度，因为在球落地前和落地后的距离较小（见图4中的损失曲线，右中）。然而，使用一个简单的时间积分器，无论∏t有多小，最终高度w.r.t.初始高度的评估梯度将是1而不是-1。这种违反直觉的行为是因为时间离散化本身并没有被编译器区分。图4说明这一效果更为详细。

我们提出了一个增加连续碰撞分辨率的简单解决方案（例如，参见Redon等人。（2002年），考虑精确的撞击时间（TOI）到前进程序（图4，左中）。虽然它几乎没有改善正向模拟（图4，右中），但梯度将是有效纠正（图4，右）。连续碰撞检测的细节见附录F。在实际的模拟器中，我们发现TOI技术可以显著改善控制器优化任务中的梯度质量（图5）。是否有TOI几乎不影响正向模拟：在补充视频中，我们显示在有TOI的模拟器中优化的机器人控制器在没有TOI的模拟器中实际工作良好。结果是，区分物理模拟器并不总是产生被模拟物理系统的有用梯度，即使模拟器很好地进行了正向模拟。在附录G中，我们讨论了我们遇到的一些额外的梯度问题。

5相关工作

可微编程 深度学习的兴起推动了深度NNs可微编程库的发展，最显著的是自动微分框架，如Theano（Bergstra et al.，2010）、TensorFlow（Abadi et al.，2016）和PyTorch（Paszke et al。，2017年）。然而，物理模拟由于其固有的计算不规则性，需要复杂的、可定制的操作。使用上述框架，程序员必须将基本操作粒度化为所需的复杂操作。这样做往往导致不满意的表现。早期关于自动微分的工作集中在转换现有的标量代码以获得导数（例如，Utke等人。（2008年），Hascoet&Pascual（2013年），Pearlmutter&Siskind（2008年）。现代编程语言最近出现了一种趋势，即通过注释来支持可微函数转换（例如Innes等人。（2019），Wei等人。（2019年）。这些框架能够区分通用编程语言，但它们提供的并行性有限。可微数组编程语言，如Halide（Ragan Kelley et al.，2013；Li et al.，2018b）、Autograd（Maclaurin et al.，2015）、JAX（Bradbury et al.，2018）和Enoki（Jakob，2019）对数组而不是标量进行操作，以利用并行性。DiffTaichi没有在im 可变的数组上操作，而是使用具有灵活索引的命令式风格，以使移植现有物理模拟算法更容易。

可微物理仿真器 为机器人学和机器学习构建可微仿真器最近越来越流行。无可微规划，Battaglia等人（2016），Chang等人。（2016）和Mrowca等人。（2018）使用NNs近似物理过程，并使用NN梯度作为近似模拟梯度。Degrave等人。（2016）和de Avila Belbute Peres等人。（2018b）分别用该方法和PyTorch方法建立了不同的刚体模拟器。Schenck&Fox（2018）使用自定义CUDA内核区分基于位置的流体。波波维奇等人。c（2000）使用可微刚体模拟器操纵基于物理的动画。ChainQueen可微弹性物体模拟器（Hu等人，2019b）用手写的CUDA核实现了连续介质力学的前向和梯度版本，其性能比纯张量流高出两个数量级实施。Liang等人。（2019）建立了一个用于材料估计和运动控制的可微布料模拟器。深度学习社区还经常包含可区分的渲染操作（OpenDR（Loper&Black，2014）、N3MR（Kato et al.，2018）、redner（Li et al.，2018a），Mitsuba 2（Nimier David等人，2019年）学习3D场景。

6结论

我们提出了DiffTaichi，一种新的可微程序设计语言，专门用于构建高性能可微物理模拟器。出于支持的需要超大内核、命令式编程和灵活的索引，我们开发了一个定制的双尺度自动区分系统。我们使用difftachi构建了10个模拟器，并将它们集成到深层神经网络中，证明了difftachi在现有系统上的性能和生产率。我们希望我们的编程语言能大大降低针对机器学习和机器人社区中的可微物理模拟未来研究的障碍。