UniMate: как управлять скелетами, не привлекая внимания санитаров

Привет, Хабр! Чтобы гуманоидный робот решал реальные задачи, нужно научить его воспринимать окружающую среду, принимать решения на основе контекста и в соответствии с ними уже двигаться в пространстве. Последнее до сих пор остается узким местом в Physical AI, и один из вариантов его преодоления — использование универсальных моделей, генерирующих движения для роботов.

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

Разберемся, какая начинка у UniMate, и что в нем может быть интересного для Phisical AI.

Сложность при работе с кинематическими моделями

Практически вся современная компьютерная графика строится на 3D-объектах, которые в последнее время становятся все более востребованными из-за необходимости обучать роботов разного типа действиям в реальном мире.

ИИ-модели неплохо выполняют риггинг по изображению объекта. Риггинг — это создание кинематической модели объекта, которая знает, как отдельные части объекта могут двигаться относительно друг друга, то есть, в нашем понимании, подвижного скелета.

Но даже если у нас уже есть готовый скелет, это еще не значит, что мы научились его анимировать так. Чтобы скелет двигался логично, естественно и при этом еще и выполнял поставленную человеком задачу, по-прежнему требуется либо ручная работа специалиста по анимации, либо специализированные ИИ-модели, заточенные под определенные типы объектов, либо вообще нужно использовать актера и технологию захвата движения. 

Модель, умеющая генерировать движения человека, не сможет анимировать собаку, дракона или какую-нибудь кракозябру с фантастическим строением. У каждого объекта своя топология — уникальное строение скелета и связей в нем. Поэтому для разных категорий персонажей часто нужны отдельные шаблоны, специализированные модели или дополнительное обучение под новую структуру.

Вот бы было здорово, если бы ИИ было можно скормить готовый скелет и сказать, какое движение выполнить?

От текста к движению

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

Допустим, у нас есть 3D-модель — человек, робот, животное или механизм. Модель уже подготовлена к анимации: у нее есть скелет, определены суставы и связи между ними. Дальше пользователь просто говорит, какое движение или последовательность движений должен выполнить персонаж: «Робот идет вперед», «Собака поворачивает голову и машет лапой» или «Дракон нападает и смотрит в сторону». А UniMate генерирует видео с указанной последовательностью движений для любого персонажа.

Здесь важно не перепутать функциональность. UniMate не создаст 3D-персонажа из обычного рисунка и не построит скелет для 3D-модели. На вход она получает готовую 3D-модель с заданной системой сочленений, а нейросеть уже решает, как этот скелет должен двигаться, и генерирует анимацию по текстовым командам.

Фишкой UniMate является универсальная модель, работающая со скелетам разной топологии. Она обучена на огромном количестве разных объектов и их скелетах, и поэтому может сгенерировать движения нового для себя персонажа.

Как это работает

В основе UniMate лежит архитектура нейросети Topology-Aware Diffusion Transformer (TADiT). Эта модель должна одновременно принимать решения на основе текстовой команды и понимать, как устроен персонаж, которого нужно анимировать.

Для генерации движений авторы используют подход flow matching, в котором нейросеть учится преобразовывать случайный шум в осмысленные данные, в данном случае — в последовательность движений. Но ключевая особенность UniMate не столько в самом способе генерации, сколько в том, как модель учитывает строение скелета. Для этого разработчики предусмотрели три механизма:

Внимание с учетом связей между суставами

Для того чтобы определять, какие части входных данных важны и как они соотносятся относительно друг друга, авторы применяют механизм внимания self-attention. Он помогает нейросети оценить изображения, но сам по себе этот метод не умеет работать со строением и кинематикой скелета. UniMate дополняет его информацией о структуре скелета, сочленениях, их возможному взаимному расположению и движению. Благодаря этому модель учитывает как локальные связи между соседними суставами, так и общую координацию движения всего объекта.

Позиционное кодирование, не привязанное к конкретному скелету

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

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

Глобальное описание скелета

Кроме отдельных связей, модели нужно понимать общую структуру персонажа. Для этого UniMate собирает информацию о скелете в единое представление — своего рода структурный профиль, который учитывается на всех этапах генерации. Технически это реализовано с помощью механизмов attention-pooling, объединяющего признаки, и AdaLN-Zero — адаптации работы слоев трансформера под входные данные. В итоге модель учитывает устройство скелета в целом.

Без хороших данных универсальную модель не обучить

Нейросети нужны данные, на которых можно научиться работать с разными скелетами и движениями. И здесь у авторов UniMate возникла проблема: собрать готовый качественный обучающий датасет оказалось не так просто. Последовательности трехмерных моделей во времени (4D-данные) есть в открытом доступе, но их мало, они все похожи, а некоторые вообще с ошибками. В одних скелетах некорректно задан корневой сустав (broken root), в других встречаются лишние или ошибочные сочленения (phantom joints). Различаются системы координат, и названия суставов, и способы организации скелетов. 

Получается, что похожие данные могут быть устроены совершенно по-разному. Если напрямую собрать из них датасет, то нейросеть не сможет его переварить и модель будет выдавать черт знает что.

Чтобы решить эту проблему, авторы проекта собрали уникальный датасет из 13 006 анимационных последовательностей UniML3D. В него вошли люди, разные двуногие персонажи, животные, птицы, морские обитатели, насекомые, змеи и даже механизмы.

Все данные, которые вошли в датасет UniML3D, прошли:

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

— Аннотирование. Анимации получают текстовые описания, необходимые для обучения генерации движений по запросу. Для этого используются большие языковые модели (LLM), которые помогают создавать подписи и приводить названия суставов к единому стандарту. Это особенно важно, когда одинаковые по смыслу части скелета обозначены в исходных файлах по-разному.

— Канонизацию. Скелеты приводят к согласованному представлению: нормализуют их структуру и формат описания чтобы модель могла сравнивать и обрабатывать скелеты разного устройства в рамках единого подхода.

Авторы опубликовали датасет на Hugging Face, чтобы другие исследователи тоже могли использовать его в своих работах.

Как работает Unimate: кейсы от авторов проекта

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

В первом примере авторы показывают, как UniMate справляется с разными топологиями скелетов и разными командами для них. Авторы тестируют анимацию совершенно разных персонажей — роботизированную руку, робота Валли, роботизированную собаку с манипулятором, мультяшного робота, гуманоидного робота и других. Структура скелета у каждого сильно различается: количество суставов, их расположение и связи между ними. Для каждого персонажа — своя задача. Здесь проявляется ключевая идея UniMate: для анимации каждого используется универсальная модель без дополнительного дообучения под конкретного персонала. 

В следующем примере авторы задают разные задания для одного типа скелета, а UniMate генерирует разные движения в зависимости от промпта. Здесь мы видим, насколько гибко модель интерпретирует команды и умеет управлять разными частями тела, показывая естественное движение объекта.

Дальше становится интереснее, авторы показывают, как модель работает с одним и тем же промптом для разных скелетов: текстовая команда одна, а персонажи отличаются строением — у одного две ноги, у другого еще есть две руки, у третьего совсем другая система сочленений. Простая команда — «Объект движется вперед». UniMate должна не просто повторить одно и то же движение, а выполнить команду с учетом возможностей каждого персонажа и его кинематической модели.

Еще авторы для разных персонажей задают промпт «Объект атакует вперед». Модель должна интерпретировать, как каждый из объектов может это сделать. Это показывает адаптацию моделью движения к устройству конкретного объекта.

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

Бонусы при работе UniMate

Кроме основных функций анимации заданных команд, авторы UniMate получили еще несколько интересных возможностей модели.

Редактирование движения. Требуется дополнить движение, которое персонаж уже выполняет. Можно описать нужную правку текстом и UniMate генерирует измененный фрагмент с учетом двух обоих промптов. В качестве примера авторы приводят робота, который сначала просто поворачивает голову налево, а потом к этому добавляется движение вперед.

Создание переходов между движениями. Если задать UniMate несколько разных действий персонажа, модель умеет сама генерировать переходы между ними: крайние позы задаются пользователем, а модель заполняет промежуток плавным и правдоподобным движением.

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

Заключение и полезные ссылки

Для Physical AI подобные сценарии интересны тем, что роботам при выполнении команд также необходимо корректировать движения, восстанавливать недостающие участки траекторий и продолжать действие в изменившихся условиях. UniMate демонстрирует общий принцип решения подобных задач: генеративную модель можно использовать для гибкой работы с движением, а не только для воспроизведения заранее заданных последовательностей.

Если хотите разобраться в нем глубже, загляните на официальный сайт проекта, изучите датасет UniML3D на Hugging Face и исходный код на GitHub.