ICLR2022:清华、腾讯AI Lab共同提出等变图力学网

时间:2022-10-07 12:37:10 作者:雷火电竞登录 来源:雷火电竞直播平台

  清华 AIR、计算机系与腾讯 AI Lab 合作,共同提出等变图力学网络,实现了理论力学中的一类重要任务—多刚体系统模拟。

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

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

  尽管如此,现有方法仍然没有很好实现对多刚体系统(Multi-rigid system)的模拟。多体问题一般假设每个物体是独立粒子,而在实际问题中,物体之间很可能存在几何约束,形成多刚体系统。实际上,多体问题是多刚体系统的一种特殊形式。如图 1(b)所示,粒子和粒子之间通过连杆连接(连杆两端的粒子只能相对转动),不同连杆可以通过铰链连接(铰链两端的连杆只能相对转动)。多刚体系统是理论力学模型中的最常见形式,在机械系统动力学控制中尤为普遍。

  为了更好地刻画多刚体系统的动力学规律,本文提出一种新颖的图神经网络——图力学网络(Graph Mechanics Network, GMN)。GMN 具备几个优点:

  1.GMN 满足刚体的几何约束,输出不会改变输入的几何属性(如棍子的长度、铰链的连接等)。

  2.GMN 是欧氏等变(Equivariant)的,即对输入做任何平移、旋转、翻转,输出都相应地改变。这利用了物理规律关于空间是均匀对称的这一特点,作为模型的归纳偏置,很好地提高了模型的泛化性和解释性。

  我们构造了一个由一定数量的球、棍子和铰链组成的虚拟物理系统,GMN 比其他方法更准确地预测这些系统演变,并满足上述所说的性质。此外,在真实的应用场景包括分子动力学模拟和人体骨架轨迹预测等,GMN 也表现出最好的预测效果。

  图 2 刚体系统的两种表示方法:左子图利用笛卡尔坐标描述系统状态,不同坐标存在几何约束;右子图利用广义坐标描述系统状态,不同坐标完全解耦

  我们的任务是给定多刚体系统的初始状态,利用提出的图力学网络预测其未来状态。在理论力学中,为了满足刚体系统的几何约束,需要引入广义坐标来表示系统的状态。如图 2 所示,我们以棍子为例,在左子图中,棍子两端球 1、球 2 的笛卡尔坐标分别是

  。因为存在棍子长度约束,所以两个球的运动过程中需要满足两者相对距离不变;换句话说,利用笛卡尔坐标系表示系统的状态,需要引入额外的约束条件,不便于采用神经网络来端到端地表征系统演化规律。为了克服这个困难,我们引入广义坐标系。如图 2 右子图所示,可以通过棍子的中心 0 点的笛卡尔坐标

  值得指出的是,受力的计算、逆向动力学、前向运动学都是通过等变神经网络来实现,极好地满足了物理学的对称性。如果我们把整个系统的初始状态旋转 90 度,我们的模型对未来状态的预测也会精准地旋转 90 度。通过引入这样的物理先验,我们能提高模型的泛化性和可解释性。

  1. 相互作用力的计算。在本文中,我们认为不同刚体中不同的球会发生相互作用,即所有的球组成全连接图。对于每个球i,它的受力的计算为:

  是相对速度。直观上,这个公式是基于刚体上每个球的受力及其在刚体局部坐标下(也就是广义坐标系)的状态推算出广义坐标系的加速度。

  4. 前向运动学。得到了更新后的广义位置和速度,我们就可以通过前向运动学对每个刚体内每个球的笛卡尔坐标进行更新:

  。需要指出的是,FK 可以是手工设置(Hard version),保证笛卡尔坐标更新始终满足约束;FK 也可以是自动学习的(Soft version),从数据中学习约束,但是不能精确保证约束条件。

  的任意多层感知机(MLP)。很容易可以证明上述定义的等变性,即对于任意正交矩阵

  。实际上,我们这里的定义是 EGNN[1]的更一般形式,从向量形式扩展为矩阵形式。

  1. 为了验证 GMN 的有效性,我们构造了一个仿真数据集 Constrained N-body。该数据集的每个数据(也就是每个系统)是由一定比例的孤立球、棍子、铰链组成。下图,GMN 利用连续的初始状态预测未来多帧状态。通过对比 GMN 预测效果(蓝色)与真实状态(红色),可以看到,GMN 能非常准确地拟合真实的动态过程。

  以下表格提供了在不同刚体组合的系统中,GMN 与当前方法的预测精度对比,可以看出 GMN 输出更小的预测误差,显著优于其他方法。

  2. 除了仿真数据,我们还在两个真实数据集:人体运动捕捉 CMU Motion Capture 和分子动力学轨迹 MD17,验证 GMN 在实际应用中的潜力。

  2.1 CMU Motion Capture 是一个记录人类动作捕捉轨迹的数据集。在这个数据中,人体骨架可以看成是一种刚性约束。在 CMU Motion Capture 上,GMN 的预测轨迹几乎与线 人体骨架运动捕捉(蓝色为预测结果,红色为真实状态)

  2.2 MD17 是分子动力学模拟的经典数据集,其包含 8 种分子的动力学运动轨迹,以下展示了 8 个分子的的分子结构:

  在这个数据集上的任务是,根据当前状态分子构象预测间隔 t 时间后分子的状态。下表是 GMN 和各个方法的对比结果。

  本文介绍了图力学网络(GMN)的构造与理论分析。借助广义坐标,GMN 能有效刻画几何约束;借助等变神经网络,GMN 能满足物理对称性。在多刚体仿真系统 Constrained N-body、人体骨架预测 CMU Motion Capture、分子动力学模拟 MD-17 等任务上都验证了 GMN 的有效性。未来,GMN 可以用于更复杂、更大体系的物理系统动力学模拟。

  清华大学智能产业研究院(Institute for AI Industry Research, Tsinghua University,英文简称 AIR)是面向第四次工业革命的国际化、智能化、产业化的应用研究机构。AIR 的使命是利用人工智能技术赋能产业升级、推动社会进步。通过大学与企业创新双引擎,突破人工智能核心技术,培养智能产业领军人才,推动智能产业跨越式发展。AIR 于 2020 年由多媒体及人工智能领域的世界级科学家、企业家张亚勤院士创建。智慧交通(AI+Transportation)、智慧物联(AI+IoT)、智慧医疗(AI+Healthcare)是清华大学智能产业研究院的三个重点研发方向。

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

  原标题:《ICLR2022:清华、腾讯AI Lab共同提出等变图力学网络,实现多刚体物理系统模拟》