バネのダイナミクス予測 #バネ

こちらの文献を参考に、下図の状況のようなバネにつながれた物体のダイナミクスのNNによる学習・予測を試してみます。

具体的には時刻$t_n$までの各物体の位置$\mathbf{r(t_n)}$、速度$\mathbf{\dot{\mathbf{r}}(t_n)}$、加速度$\mathbf{\ddot{\mathbf{r}}(t_n)}$を入力として、一期先時刻$t_{n+1}$での加速度変化($\mathbf{\ddot{\mathbf{r}}(t_{n+1})}$-$\mathbf{\ddot{\mathbf{r}}(t_n)}$)を予測します。
※文献の詳細は読めていないので、思想的なところだけ参考に試していきます。

大まかな流れ

運動方程式から数値計算を行いていくつかの初期条件で教師データを生成
1で作成した教師データを用いてNNモデルを学習
適当な状態からその先の状態をNNモデルから逐次的に予測し、それらしい運動になるか確認する。

1. 教師データの生成

簡易のため、適当な初期状態を生成後、運動方程式からEuler法により逐次計算を行います。

位置$\mathbf{r}_{i, j}$の物体の運動方程式

上下左右4方向からの力の和で、$m$は質量、$\mathbf{r}^{eq}_{i,j,i\prime,j\prime}$は$(i,j)$の物体から$(i\prime, j\prime)$の物体を向いた平衡位置のベクトル、$k$はばね定数です。端っこの連結がない部分は除外します。

上記式から、以下Euler法で一期先の状態を計算します

物体定義と初期化

@dataclass
class Obj:
    idx: int
    pos: np.array
    vel: np.array
    accel: np.array
    force: np.array
    connected_objs: List = None
    mass: float = 1.0
    energy: float = 0.0

def init_objs_2d():
    objs = {}
    # 物体作成
    for i in range(4):
        for j in range(4):
            objs[(i,j)] = Obj(
                idx=i*4+j,
                pos=np.array([i, j])+0.05*np.random.randn(2),
                vel=0.01*np.random.randn(2),
                accel=np.zeros(2),
                force=np.zeros(2),
            )

    # バネ連結情報付与
    for i in range(4):
        for j in range(4):
            connected_objs = []
            if i-1 >= 0:
                connected_objs.append(objs[(i-1,j)])
            if i+1 <= 3:
                connected_objs.append(objs[(i+1,j)])
            if j-1 >= 0:
                connected_objs.append(objs[(i,j-1)])
            if j+1 <= 3:
                connected_objs.append(objs[(i,j+1)])

            objs[(i,j)].connected_objs = connected_objs
    return objs

各種状態更新関数

def force(pos1: np.array, pos2: np.array):
    """pos2の物体からpos1の物体に働くフックの法則による力"""
    dist = np.sqrt(((pos1-pos2)**2).sum())
    r_normed = (pos2-pos1)/dist
    equilibrium_vec = EQUILIBRIUM_LENGTH*r_normed
    f = -K*(equilibrium_vec-(pos2-pos1))
    return f

def _update_force_on_obj(obj: Obj):
    """物体に働く力の更新"""
    forces = [force(obj.pos, obj2.pos) for obj2 in obj.connected_objs]
    obj.force = sum(forces)

def _update_accel(obj: Obj):
    """物体の加速度更新"""
    obj.accel = obj.force / obj.mass

def _update_velocity(obj: Obj):
    """物体の速度更新"""
    obj.vel = obj.vel + DT*obj.accel

def _update_pos(obj: Obj):
    """物体の位置更新"""
    obj.pos = obj.pos + DT*obj.vel

def update_states(objs: List[Obj]):
    """全物体の状態更新"""
    [_update_force_on_obj(obj) for obj in objs.values()]
    [_update_accel(obj) for obj in objs.values()]
    [_update_velocity(obj) for obj in objs.values()]
    [_update_pos(obj) for obj in objs.values()]
    [_calc_energy(obj) for obj in objs.values()]

計算結果の可視化

赤い矢印が物体に働いているフックの法則による力のベクトルで、緑の矢印が物体の速度ベクトルです。

全力学的エネルギーの時間変化も確認してみます。

保存されていないみたいですが、気にせず進みます。

2. ニューラルネットワークモデルの学習

1で生成したデータから、各物体の時刻$t_n$までの各物体の状態を入力、各物体の一期先$t_{n+1}$の加速度変化を出力としたグラフニューラルネットワークを学習します。

モデル構造

モデル全体像イメージは冒頭で述べた文献記載の下図モデルを大まかに参考にし、GNがグラフニューラルネットワーク(GNN)モジュール、エンコーダー部分の詳細は分からなかったのでエンコーダーは除外、デコーダーは単純に全結合層にしました。
グラフにおいて、物体がノード、物体をつなげてるバネがエッジで、ノード(物体)に特徴量(位置、速度等)が紐づいているイメージです。

ref

GNNモジュール

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GNNModule(torch.nn.Module):
    def __init__(self, n_features, out_dim, intermediate_dim=24):
        super(GNNModule, self).__init__()
        self.conv1 = GCNConv(n_features, intermediate_dim)
        self.conv2 = GCNConv(intermediate_dim, out_dim)

    def forward(self, x, edge_index):
        x2 = self.conv1(x, edge_index)
        x3 = F.relu(x2)
        x4 = F.dropout(x3, training=self.training)
        x5 = self.conv2(x4, edge_index)
        return x5

(バッチサイズ, ノード数(物体数), n_features)のfeatureと、(2, エッジ数(バネの数))のedge_indexを入力すると、(バッチサイズ、out_dim)のテンソルが出力されます。

全体ネットワーク

class SpringDynamicsNet(torch.nn.Module):
    def __init__(self, n_features, out_dim, n_gnn_module=4, intermediate_dim=24):
        super(SpringDynamicsNet, self).__init__()
        self._intermediate_gnn_modules = nn.Sequential(
            *[GNNModule(n_features, n_features) for _ in range(n_gnn_module)]
        )
        self._fc = nn.Linear(n_features, out_dim)

    def forward(self, x, edge_index):
        for sping_module in self._intermediate_gnn_modules:
            residual = x
            gnn_out = sping_module(x, edge_index)
            x = residual + gnn_out
        return self._fc(x)

入出力データ

入力の状態は時刻$t_n$における各物体の位置・加速度、$t_n$から$N$期前までの各物体の速度とし、今回は$N$は3で行い、出力は１期先の加速度変化としました。
具体的に書くと、物体１つに対し、入力は
$[x(t), y(y), v_x(t), v_y(t), v_x(t-1), v_y(t-1), v_x(t-2), v_y(t-2), v_x(t-3), v_y(t-3), a_x(t), a_y(t)]$
出力は
$[a_x(t+1)-a_x(t), a_y(t+1)-a_y(t)]$
といった感じです。
これが物体数(16)あるので、入力が(バッチサイズ、16, 12)、出力が(バッチサイズ、16, 2)のテンソルになります。
モデル入力には上記に加え、下記の物体連結情報のエッジデータも渡します。

エッジデータ作成

連結している物体のインデックスペアリストを作ります。

edge_index = {}
for obj in objs.values():
    for obj2 in obj.connected_objs:
        edge_index[(obj.idx, obj2.idx)] = 1

edge_index = list(edge_index.keys())
edge_index = torch.tensor(edge_index).T

学習

spring_dynamics_net = SpringDynamicsNet(
    n_features=FEATURE_DIM,
    out_dim=2,
    n_gnn_module=4,
    intermediate_dim=24
)

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(spring_dynamics_net.parameters(), lr=0.01, weight_decay=5e-4)

BATCH_SIZE = 32
N_EPOCHS = 20

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
spring_dynamics_net.to(device)
spring_dynamics_net.train()

loss_histroy = []

for epoch in pb(range(N_EPOCHS)):
    epoch_loss = 0
    for i in range(feats.shape[0]//BATCH_SIZE):
        feat_batch = feats[i*BATCH_SIZE:(i+1)*BATCH_SIZE, :, :]
        out_batch = outs[i*BATCH_SIZE:(i+1)*BATCH_SIZE, :, :]
        optimizer.zero_grad()
        preds = spring_dynamics_net(feat_batch, edge_index) 
        loss = criterion(out_batch, preds)
        loss.backward(retain_graph=True)
        optimizer.step()
        epoch_loss += loss.item()*feat_batch.shape[0]
    loss_histroy.append(epoch_loss)

損失時系列

学習できてるようです。

バネのダイナミクス予測