Подводные камни Learned Motion Matching: опыт реализации / Хабр

Автор: Георгий Маркелов, разработчик Softellion

В прошлой статье я сделал обзор на существующие на данный момент методы реализации анимации персонажей с применением машинного обучения. Однако существует ряд нераскрытых приемов, без которых финальный результат будет неудовлетворительным. В этой статье я поделюсь опытом реализации Learned Motion Matching.

Предисловие

Все формулы справедливы для левосторонней (left‑handed) системы координат:

Кватернионы хранятся в формате (скаляр на последнем месте)

Quaternion unroll

Поскольку один и тот же угол может быть представлен двумя кватернионами и , может получиться так, что появляются сильные разрывы в данных (-179 и + 181 образуют одинаковый поворот) и одна и та же поза (или похожие) кодируется совершенно по-разному. При обучении нейросети это совершенно теряет всякий смысл, и при попытке интерполяции таких противоположных значений также не выйдет ничего хорошего. Поэтому необходимо сделать разворот кватерниона:

def unroll(x):
   y = x.clone()
   for i in range(1, len(x)):
       d0 = torch.sum( y[i] * y[i-1], dim=-1)
       d1 = torch.sum(-y[i] * y[i-1], dim=-1)
       y[i][d0 < d1] = -y[i][d0 < d1]
      
   return y

Суть в том, что, если скалярное произведение с кватернионом с предыдущего кадра отрицательное, то его нужно развернуть. В таком случае на протяжении всей последовательности будет соблюдаться условие $q_i⋅q_{i−1} ≥ 0$ .

Наглядный пример значений компонентов кватерниона до и после разворота:

После разворота кватернионов следующей важной задачей становится создание новой искусственной кости, описывающей траекторию движения — root transformation.

Root transformation

Эта кость может быть как создана вручную, например в Blender, так и сгенерирована процедурно, но присутствовать должна обязательно — в противном случае весь алгоритм не работает. При ее создании важно помнить, что она должна как можно более точно передавать характер движения персонажа. После добавления делается корневой для всего скелета. Позиция находится так: проецируем верхнюю кость позвоночника на землю и добавляем сглаживающий фильтр Савицкого-Голея:

from scipy import ndimage, signal

g_pos, g_rot = quat_utils.fk(y_pos, y_rot, parents) # forward kinematics

sim_bone_pos = torch.from_numpy(
    signal.savgol_filter(g_pos[:, spine_node_id] * torch.tensor((1.0, 1.0, 0.0)), 
                         31, 3, axis=0, mode='interp')
)

На датасете со спокойной ходьбой данный фильтр может давать разницу в ~0.6 метра между исходной и сглаженной позицией, что довольно внушительно.

Затем берем направление «вперед» кости бедра, проецируем на землю и также сглаживаем фильтром. После чего конвертируем во вращение вокруг вертикальной оси — так вычисляется вращение кости. Также применяем фильтр Савицкого-Голея:

# direction taken from projected hip forward direction
# 0.0, 0.0, 1.0 is local up axis
sim_bone_dir = (torch.tensor((1.0, 1.0, 0.0)) *
                    quat_utils.mul_vec(g_rot[:, hips_node_id:hips_node_id+1], 
                                       torch.tensor((0.0, 0.0, 1.0)).unsqueeze(0).unsqueeze(0))
                    )
sim_bone_dir = quat_utils.normalize(sim_bone_dir)
sim_bone_dir = torch.from_numpy(signal.savgol_filter(sim_bone_dir, 61, 3, axis=0, mode='interp'))
sim_bone_dir = quat_utils.normalize(sim_bone_dir)

sim_bone_rot = quat_utils.normalize(
        quat_utils.between(torch.tensor([1.0, 0.0, 0.0]).unsqueeze(0).unsqueeze(0), sim_bone_dir)
    ).squeeze(1)

Обязательно не забываем нормализовывать кватернионы после каждой операции.

После чего назначаем новую кость корневой для всего скелета и вычисляем позиции всех остальных костей с учетом новой иерархии, не забывая также обновить и саму иерархию:

# Transform first joint to be local to sim and append sim as root bone
y_pos[:, 0] = quat_utils.inv_mul_vec(sim_bone_rot, y_pos[:, 0] - sim_bone_pos)
y_rot[:, 0] = quat_utils.inv_mul(sim_bone_rot, y_rot[:, 0])

y_pos = torch.concat([sim_bone_pos.unsqueeze(1), y_pos], dim=1)
y_rot = torch.concat([sim_bone_rot.unsqueeze(1), y_rot], dim=1)

parents = torch.concat([torch.tensor((-1,)), parents + 1])

После добавления новой кости стоит задача вычисления скоростей всех костей скелета.

Скорость костей

Сильные выбросы значений могут приводить в дальнейшем к дрожанию анимации, поэтому, хотя скорость и можно вычислить по стандартной формуле, для исключения сильных всплесков лучше использовать интерполяцию методом центральной разности (central difference) $v_i=(x_{i+1}−x_{i−1}) / 2h$

vel[1:-1] = (pos[2:] - pos[1:-1] +
            pos[1:-1] - pos[:-2]) * 0.5 * fps
vel[0] = vel[1] - (vel[3] - vel[2])
vel[-1] = vel[-2] + (vel[-2] - vel[-3])

Метод использует симметричное окно вокруг одной точки. Это дает три эффекта:

шумы и скачки в одном кадре компенсируются соседними кадрами;
оценка получается второго порядка точности , а не первого, как в forward/backward difference;
скорость получается более гладкой и менее чувствительной к единичным выбросам;

Однако полностью выбросы конечно же не исключаются — снижается чувствительность к шуму.

Дополнительно отдельно вычисляются значения в первой и последней точке методом линейной экстраполяции, потому что, очевидно, там нет соседних точек с обеих сторон.

Разобравшись с вычислением скорости, приходит черед представления вращения костей и определение контакта ног с поверхностью.

2-column rotation matrix

Как было сказано ранее (см. quaternion unroll), предсказывание кватернионов нейросетью не является оптимальным решением, поэтому вращения костей в качестве входных и выходных данных для нейросети представляются в виде относительных векторов «вперед» и «вверх».

def to_xform_xz(x):
   qx, qy, qz, qw = x[..., 0:1], x[..., 1:2], x[..., 2:3], x[..., 3:4]

   x2, y2, z2 = qx + qx, qy + qy, qz + qz

   # Precompute products for the rotation matrix
   xx, yy, wx = qx * x2, qy * y2, qw * x2
   xy, yz, wy = qx * y2, qy * z2, qw * y2
   xz, zz, wz = qx * z2, qz * z2, qw * z2

   return torch.concatenate([
       torch.concatenate([1.0 - (yy + zz), xz + wy], dim=-1)[..., torch.newaxis, :],
       torch.concatenate([xy + wz, yz - wx], dim=-1)[..., torch.newaxis, :],
       torch.concatenate([xz - wy, 1.0 - (xx + yy)], dim=-1)[..., torch.newaxis, :],
   ], dim=-2)

Контакт ног с поверхностью

Можно вычислить через скорость с заданием порогового значения — если скорость ниже, то считаем, что нога стоит на месте. L2-норма 3D скорости сводит значение к скаляру. Добавляем медианный фильтр для устранения слишком коротких «окон» (ложных срабатываний):

contact_velocity_threshold = 0.15 * 100  # to cm
# L2 norm of 3D velocity collapses all motion into single scalar speed
contact_velocity = torch.norm(g_vel[:, [left_toe_id, right_toe_id]], dim=-1)

extra = contact_velocity < contact_velocity_threshold

for ci in range(extra.shape[1]):
   extra[:, ci] = torch.from_numpy(ndimage.median_filter(extra[:, ci], size=6, mode='nearest'))

Пример работы фильтра:

Наконец, последней задачей аугментации данных перед тренировкой нейросетей выступает формирование траектории движения.

Траектория движения

Траектория вычисляется в точках 20, 40, 60 кадров в будущем, поэтому необходимо учесть окончание ролика и корректно его обработать:

def clamp_frame(frame, offset):
   for i in range(len(range_starts)):
       if range_starts[i] <= frame < range_stops[i]:
           return min(max(frame + offset, range_starts[i].item()), range_stops[i].item() - 1)


for i in range(start, stop):
   future_frame_20 = clamp_frame(i, 20)
   future_frame_40 = clamp_frame(i, 40)
   future_frame_60 = clamp_frame(i, 60)

Так, по приближению к окончанию ролика, мы плавно ограничиваем будущий кадр последним.

Позиции и вращения траектории можно было бы просто получить как разность, однако, если дополнительно умножить на обратный кватернион Root transformation, то мы развернем персонажа так, что он на каждом кадре будет смотреть всегда в одном направлении. Таким образом достигается инвариантность по отношению к направлению движения.

traj_pos_20 = quat_utils.inv_mul_vec(y_rot[i, 0], y_pos[future_frame_20, 0] - y_pos[i, 0])
traj_pos_40 = quat_utils.inv_mul_vec(y_rot[i, 0], y_pos[future_frame_40, 0] - y_pos[i, 0])
traj_pos_60 = quat_utils.inv_mul_vec(y_rot[i, 0], y_pos[future_frame_60, 0] - y_pos[i, 0])

# Rotate the local forward vector (1.0, 0.0, 0.0) because X is forward to get world-space forward
traj_dir_20 = quat_utils.inv_mul_vec(y_rot[i, 0], quat_utils.mul_vec(y_rot[future_frame_20, 0], torch.tensor((1.0, 0.0, 0.0))))
traj_dir_40 = quat_utils.inv_mul_vec(y_rot[i, 0], quat_utils.mul_vec(y_rot[future_frame_40, 0], torch.tensor((1.0, 0.0, 0.0))))
traj_dir_60 = quat_utils.inv_mul_vec(y_rot[i, 0], quat_utils.mul_vec(y_rot[future_frame_60, 0], torch.tensor((1.0, 0.0, 0.0))))

То же самое касается и вычисления позиций и вращений костей, входящих в набор для описания параметров позы (features):

left_foot_pos = quat_utils.inv_mul_vec(rot[:, 0], g_pos[:, left_foot_id] - g_pos[:, 0])
right_foot_pos = quat_utils.inv_mul_vec(rot[:, 0], g_pos[:, right_foot_id] - g_pos[:, 0])
left_foot_vel = quat_utils.inv_mul_vec(rot[:, 0], g_vel[:, left_foot_id])
right_foot_vel = quat_utils.inv_mul_vec(rot[:, 0], g_vel[:, right_foot_id])
hips_vel = quat_utils.inv_mul_vec(rot[:, 0], g_vel[:, hips_node_id])

На это аугментация данных заканчивается.

Projector nearest neighbor search

Напомню, что нейросеть Projector эмулирует поиск методом ближайшего соседа для . Задача сводится к нахождению евклидового расстояния и является довольно ресурсоемкой, а также занимает большую часть времени всего алгоритма.

$d(p,q) = \sqrt{(p_1 - q_1)^2 + (p_2 - q_2)^2 + \ldots + (p_n - q_n)^2} = \sqrt{\sum_{k=1}^n(p_k - q_k)^2 }$

Для ускорения поиска можно использовать структуру sklearn.neighbors.BallTree:

from sklearn.neighbors import BallTree


tree = BallTree(x)
nearest = tree.query(x_hat, k=1, return_distance=False)[:, 0]

Можно использовать и другие решения, например, faiss, он даст еще больший прирост производительности, но и его применение не такое простое, а также нет официальной поддержки Windows.

Сохранение и интеграция обученных моделей

После обучения моделей LMM возникает задача их интеграции в конечную среду. Важно выбрать формат, который будет одновременно удобен для инференса и совместим с движком.

Здесь есть 2 варианта решения:

Сохранить модели в формате ONNX
Записывать веса каждого слоя и параметры нормализации (СКО и среднее значение) непосредственно в бинарный формат

ONNX удобен тем, что, если конечная среда предоставляет инструменты инференса моделей в таком формате, то ничего больше не требуется делать. Во втором же случае нужно будет дополнительно реализовать нормализацию данных, матричное умножение и функцию активации (ReLU в случае LMM) и сделать это эффективно.

Единственная проблема, которая возникла при работе с ONNX - совместимость поддержки конечной средой ONNX opset version и opset_version, который использовался при экспорте модели в этот формат.

В качестве рабочей среды для интеграции LMM было решено использовать Unreal Engine 5, который имеет встроенные инструменты инференса ONNX моделей - UNNERuntime. Для работы с ним нужно только загрузить модель и описать входные/выходные данные.

Сам экспорт в ONNX выполняется тривиально, нужно только не забыть переключить модель в режим вычисления и отключить градиенты:

compressor.eval()

with torch.no_grad():
   input = (torch.cat([
           y_pos[:, 1:].reshape([num_frames, -1]),
           y_txz[:, 1:].reshape([num_frames, -1]),
           y_vel[:, 1:].reshape([num_frames, -1]),
           y_ang[:, 1:].reshape([num_frames, -1]),
           q_pos[:, 1:].reshape([num_frames, -1]),
           q_txz[:, 1:].reshape([num_frames, -1]),
           q_vel[:, 1:].reshape([num_frames, -1]),
           q_ang[:, 1:].reshape([num_frames, -1]),
           y_root_vel.reshape([num_frames, -1]),
           y_root_ang.reshape([num_frames, -1]),
           y_extra.reshape([num_frames, -1]),
       ], dim=-1) - compressor_mean_in) / compressor_std_in

   z = compressor(input.to(device))

torch.save(z.detach().cpu(), root_dir / 'z.pt')

torch.onnx.export(
   compressor,
   (compressor_mean_in,),
   f"{root_dir}/compressor.onnx",
   input_names=["y_q"],
   output_names=["z"],
   dynamo=True
)

Аналогично и для других моделей за тем исключением, что не нужно сохранять никакие промежуточные значения.

После того как модель сохранена, её можно загрузить в UE5 и выполнить инференс.

Ниже приведен пример инференса модели Decompressor:

bool NNModels::RunDecompressorModel(TArray<float>& features, TArray<float>& latent) {
   if (DecompressorModel.IsValid() == false || InputBindings.Num() == 0) {
      return false;
   }

   TArray<float> input;
   input.SetNumZeroed(features.Num() + latent.Num());

   for (int i = 0; i < features.Num(); i++) {
      input[i] = features[i];
   }

   for (int i = 0; i < latent.Num(); i++) {
      input[features.Num() + i] = latent[i];
   }

   InputBindings[0].Data = input.GetData();
   InputBindings[0].SizeInBytes = input.Num() * sizeof(float);

   OutputBindings[0].Data = DecompressorOutputData.GetData();
   OutputBindings[0].SizeInBytes = DecompressorOutputData.Num() * sizeof(float);

   return DecompressorModel->RunSync(InputBindings, OutputBindings) == UE::NNE::EResultStatus::Ok;
}

Заключение

Применив перечисленные выше приемы, можно получить стабильно работающую систему Learned Motion Matching. Однако не стоит считать, что они единственно верные: так, можно использовать другие методы интерполяции, формирования root transformation, представления данных и т.д.