ICLR2022：清华、腾讯AI Lab共同提出等变图力学网络，实现多刚体物理体系摹拟

近年来，AI for Science 操纵人工智能步骤与物理、化学、生物等自然科学进行交叉融合，在一些重要的科学问题上（如蛋白质结构展望）取得了瞩目的进展。鉴于物理学科的基础性与重要性，AI+Physics 无疑是 AI for Science 不可缺失的一环。为此，清华 AIR、计算机系与腾讯 AI Lab 合作共同发表论文《Equivariant Graph Mechanics Networks with Constraints》, 提出等变图力学网络，实现了理论力学中的一类重要任务—多刚体体系摹拟。论文已被 ICLR2022 接收。这一项目也收到了来自腾讯 AI Lab 犀牛鸟专项研究计划的资助。 

论文地址：https://openreview.net/forum?id=SHbhHHfePhP

开源代码：https://github.com/hanjq17/GMN

一、引言多体问题（N-body problem）是物理学的基础问题，广泛存在于诸多任务中，例如电子绕原子核运动、份子动力学摹拟、机械体系（如机械臂）的动力控制、宇宙天体的轨迹展望等，都可以表示成多体问题。多体问题通常不存在解析解，其求解需要冗长的数值计算。近年来，机器学习研究人员操纵图神经网络（Graph Neural Network, GNN）从数据中挖掘多体体系的演化规律，为多体问题的求解提供了一个全新的思路。该类步骤把物体建模成节点，把物体之间的交互建模成连边，训练好的 GNN 直接展望体系的未来情态。

图 1 从多体问题到多刚体体系摹拟尽管如此，现有步骤仍然没有很好实现对多刚体体系（Multi-rigid system）的摹拟。多体问题一般假设每一个物体是独立粒子，而在实际问题中，物体之间很可能存在若干自在，形成多刚体体系。实际上，多体问题是多刚体体系的一种特殊形式。如图 1（b）所示，粒子和粒子之间通过连杆连接（连杆两端的粒子只能相对转动），不同连杆可以通过搭钮连接（搭钮两端的连杆只能相对转动）。多刚体体系是理论力学模型中的最常见形式，在机械体系动力学控制中尤为普遍。为了更好地刻画多刚体体系的动力学规律，本文提出一种新颖的图神经网络——图力学网络（Graph Mechanics Network, GMN）。GMN 具备几个优点:1.GMN 满足刚体的若干自在，输出不会改变输入的若干属性（如棍子的长度、搭钮的连接等）。2.GMN 是欧氏等变（Equivariant）的，即对输入做任何平移、旋转、翻转，输出都相应地改变。这操纵了物理规律关于空间是均匀对称的这一特点，作为模型的归纳偏置，很好地提高了模型的泛化性和解释性。3. 理论上，GMN 具有良好的表达能力，能很好地逼近任何等变的函数。我们构造了一个由一定数量的球、棍子和搭钮组成的虚拟物理体系，GMN 比其他步骤更准确地展望这些体系演变，并满足上述所说的性质。此外，在实在的使用场景包括份子动力学摹拟和人体骨架轨迹展望等，GMN 也表现出最好的展望效果。

二、步骤简介

图 2 刚体体系的两种表示步骤：左子图操纵笛卡尔坐标描述体系情态，不同坐标存在若干自在；右子图操纵广义坐标描述体系情态，不同坐标完全解耦我们的任务是给定多刚体体系的初始情态，操纵提出的图力学网络展望其未来情态。在理论力学中，为了满足刚体体系的若干自在，需要引入广义坐标来表示体系的情态。如图 2 所示，我们以棍子为例，在左子图中，棍子两端球 1、球 2 的笛卡尔坐标分别是。因为存在棍子长度自在，所以两个球的运动过程中需要满足两者相对距离不变；换句话说，操纵笛卡尔坐标系表示体系的情态，需要引入额外的自在条件，不便于采用神经网络来端到端地表征体系演化规律。为了克服这个困难，我们引入广义坐标系。如图 2 右子图所示，可以通过棍子的中心 0 点的笛卡尔坐标，以及棍子的相对转角，来表示棍子的情态。同时，可以独立变化，始终不改变棍子的长度，特别利于使用神经网络进行建模。具体而言，为了刻画多刚体体系的变化过程，我们的步骤主要包含 4 部分：1. 在笛卡尔坐标空间计算每一个球受到其他刚体的作用力；2. 操纵逆向动力学把每一个刚体的受力转化成广义坐标的加速度变化；3. 对广义坐标进行革新；4. 操纵前向运动学，将革新后的广义坐标映射到笛卡尔坐标空间。值得指出的是，受力的计算、逆向动力学、前向运动学都是通过等变神经网络来实现，极好地满足了物理学的对称性。如果我们把整个体系的初始情态旋转 90 度，我们的模型对未来情态的展望也会精准地旋转 90 度。通过引入这样的物理先验，我们能提高模型的泛化性和可解释性。下面简要介绍每一部分过程。1. 相互作用力的计算。在本文中，我们认为不同刚体中不同的球会发生相互作用，即所有的球组成全连接图。对于每一个球i，它的受力的计算为：，即，每两个球的受力与其相对距离有关。2. 逆向动力学。基于每一个刚体上所有受力及其分布，我们计算广义坐标的加速度为：，是相对位置，是相对速度。直观上，这个公式是基于刚体上每一个球的受力及其在刚体局部坐标下（也就是广义坐标系）的情态推算出广义坐标系的加速度。3. 广义坐标革新。有了加速度，我们就可以革新广义坐标的速度和位置：。4. 前向运动学。得到了革新后的广义位置和速度，我们就可以通过前向运动学对每一个刚体内每一个球的笛卡尔坐标进行革新：。需要指出的是，FK 可以是手工设置（Hard version），保证笛卡尔坐标革新始终满足自在；FK 也可以是自动学习的（Soft version），从数据中学习自在，但是不能精确保证自在条件。为了简便，以上过程没有考虑每一个节点的隐层变量，和广义坐标角度。读者可以通过查阅原论文得到更多这方面的细节。三、理论分析

上文提到的函数均为欧氏等变函数。不失一般性，我们记正交等变函数。我们定义，

其中，为参数是的任意多层感知机（MLP）。很容易可以证明上述定义的等变性，即对于任意正交矩阵。实际上，我们这里的定义是 EGNN[1]的更一般形式，从向量形式扩展为矩阵形式。我们还可以证明，这样定义的正交等变函数具有很好的表达能力。【定理 1】如果，行满秩，即，那么对于任意连续的正交等变函数，一定存在某个使得充分逼近，即,任意小。三、实验结果1. 为了验证 GMN 的有效性，我们构造了一个仿真数据集 Constrained N-body。该数据集的每一个数据（也就是每一个体系）是由一定比例的孤立球、棍子、搭钮组成。下图，GMN 操纵连续的初始情态展望未来多帧情态。通过对比 GMN 展望效果（蓝色）与实在情态（红色），可以看到，GMN 能非常准确地拟合实在的动态过程。

图 3 多刚体体系仿真结果（蓝色为展望结果，红色为实在情态）以下表格提供了在不同刚体组合的体系中，GMN 与当前步骤的展望精度对比，可以看出 GMN 输出更小的展望误差，显著优于其他步骤。 

2. 除了仿真数据，我们还在两个实在数据集：人体运动捕捉 CMU Motion Capture 和份子动力学轨迹 MD17，验证 GMN 在实际使用中的潜力。2.1 CMU Motion Capture 是一个记录人类动作捕捉轨迹的数据集。在这个数据中，人体骨架可以看成是一种刚性自在。在 CMU Motion Capture 上，GMN 的展望轨迹几乎与实在轨迹重合。

图 4 人体骨架运动捕捉（蓝色为展望结果，红色为实在情态）以下是不同步骤的展望误差：

2.2 MD17 是份子动力学摹拟的经典数据集，其包含 8 种份子的动力学运动轨迹，以下展示了 8 个份子的的份子结构：

在这个数据集上的任务是，根据当前情态份子构象展望间隔 t 时间后份子的情态。下表是 GMN 和各个步骤的对比结果。

四、结论本文介绍了图力学网络（GMN）的构造与理论分析。借助广义坐标，GMN 能有效刻画若干自在；借助等变神经网络，GMN 能满足物理对称性。在多刚体仿真体系 Constrained N-body、人体骨架展望 CMU Motion Capture、份子动力学摹拟 MD-17 等任务上都验证了 GMN 的有效性。未来，GMN 可以用于更复杂、更大体系的物理体系动力学摹拟。[1] Satorras V G, Hoogeboom E, Welling M. E (n) equivariant graph neural networks[C]//International Conference on Machine Learning. PMLR, 2021: 9323-9332.关于 AIR清华大学智能产业研究院（Institute for AI Industry Research, Tsinghua University，英文简称 AIR）是面向第四次工业革命的国际化、智能化、产业化的使用研究机构。AIR 的使命是操纵人工智能技术赋能产业升级、推动社会进步。通过大学与企业创新双引擎，突破人工智能核心技术，培养智能产业领军人才，推动智能产业跨越式发展。AIR 于 2020 年由多媒体及人工智能领域的世界级科学家、企业家张亚勤院士创建。智慧交通（AI＋Transportation）、智慧物联（AI＋IoT）、智慧医疗（AI＋Healthcare）是清华大学智能产业研究院的三个重点研发方向。扫一扫，关注我们：

腾讯 AI Lab腾讯的企业级 AI 实验室，于 2016 年 4 月在深圳成立，目前有 100 多位顶尖研究科学家及 300 多位使用工程师。借助腾讯丰富使用场景、大数据、计算力及一流人才方面的长期积累，AI Lab 立足未来，开放合作，致力于不断提升 AI 的认知、决策与创造力，向 “Make AI Everywhere” 的愿景迈步。腾讯 AI Lab 强调研究与使用并重发展。基础研究关注机器学习、计算机视觉、语音技术及自然语言处理等四大方向；技术使用聚焦在游戏、数字人、内容、社交四大领域，并初探 AI 在工业、农业、医疗、医药、生命科学等领域的研究使用。详情可访问：ai.tencent.com