Transformer Transformer: ИИ сам спроектировал робота и снизил ошибку слежения на 73%

Команда представила модель Transformer Transformer, диффузионный трансформер, который решает задачу «совместного проектирования» робота (co-design): по демонстрации целевого движения конечного эффектора и функции вознаграждения модель генерирует полную конструкцию робота, каждое звено, шарнир, мотор и параметры инерции, а затем сама же управляет получившимся роботом, чтобы проверить дизайн. Одна и та же сеть выступает генератором, критиком и кросс-платформенным контроллером.

В основе, RoboTokens, единая токенизация, которая описывает и конструкцию робота (неизменную во времени), и его динамику (состояния и действия во времени) одной последовательностью токенов. Модель обучили на 11 роботах из библиотеки MuJoCo Menagerie, от 0,65-килограммовой кисти-манипулятора до 67,5-килограммового четвероногого робота, от 6 до 35 подвижных шарниров; каждый робот превратился в последовательность из 28, 101 токена.

Ключевой прием назвали Dynamics Self-Guidance («самонаведение по динамике»): модель предсказывает не только конструкцию, но и её динамику, поэтому предсказанную траекторию можно подставить в пользовательскую функцию вознаграждения и получить дифференцируемую оценку. Градиент этой оценки встраивается обратно в процесс диффузии и на каждом шаге подталкивает генерацию к более высокому вознаграждению, по аналогии с classifier-guided diffusion. Так модель зеро-шотно оптимизирует конструкцию под функции вознаграждения, которые не видела при обучении: например, добавление штрафа за крутящий момент делает робота компактнее и легче при сохранении точности движения.

Метод проверили в двух сценариях. Первый, на подмножестве данных UMI по двурукой мойке посуды, где 26 траекторий отложили для проверки: модель генерировала конструкции, оптимизированные сразу под всё распределение движений, комбинируя дискретный выбор (жёсткий/скользящий/гибкий позвоночник робота) и непрерывные параметры (размеры рук, точки крепления). Против базовых методов, случайной генерации и эволюционного оптимизатора CMA-ES, модель выигрывала и по итоговому вознаграждению, и по скорости работы.

Второй сценарий, проверка на реальном железе. Модель оптимизировала конструкцию двурукого робота ALOHA2 под задачу динамического расправления ткани (тест выбрали намеренно жёстким: и кинематика, рука должна дотягиваться над тканью на высокой скорости, и непросчитываемая динамика вроде аэродинамического сопротивления). Изменённую конструкцию изготовили физически; по сравнению с оригинальным ALOHA она снизила ошибку слежения за траекторией на 73% и максимальную скорость шарниров на 30%.

Авторы сами перечисляют ограничения. RoboTokens пока описывают только примитивную геометрию, а не произвольные меши, деформируемые объекты или контекст сцены. Кросс-платформенный контроллер коррелирует с эталонным RL-контроллером с коэффициентом Пирсона r = 0,53 на сгенерированных конструкциях, результат обнадёживающий, но с заметным запасом для улучшения; на отдельных нетипичных конструкциях контроллер плохо обобщается. В отличие от больших языковых моделей, у которых качество растёт с увеличением времени на «размышление», выигрыш от увеличения вычислений на этапе вывода у Transformer Transformer выходит на плато примерно через минуту. Наконец, обучающие данные для RL-контроллеров дороги в получении, около 16 часов на одном GPU A100 на одну конструкцию, и именно это авторы называют главным препятствием для масштабирования метода на новые классы роботов.

Ключевые факты

Диффузионная модель Transformer Transformer генерирует по демонстрации движения полную конструкцию робота (звенья, шарниры, моторы, инерция) и сама же управляет ей для проверки, одна сеть в трёх ролях: генератор, критик, контроллер
На физическом двуруком роботе ALOHA2 сгенерированная под задачу расправления ткани конструкция снизила ошибку слежения за траекторией на 73% и максимальную скорость шарниров на 30% против оригинала
Метод Dynamics Self-Guidance превращает предсказанную моделью динамику в дифференцируемую оценку вознаграждения и направляет диффузию к лучшим конструкциям, зеро-шотно, под функции вознаграждения, не встречавшиеся при обучении
Модель обучена на 11 роботах из MuJoCo Menagerie (0,65, 67,5 кг, 6, 35 шарниров) и превзошла базовые методы, случайную генерацию и эволюционный оптимизатор CMA-ES, по качеству и по скорости работы
Авторы честно называют ограничения: контроллер коррелирует с эталонным RL-контроллером лишь с r = 0,53, выигрыш от вычислений на этапе вывода выходит на плато примерно за минуту, а данные для обучения RL-контроллеров стоят около 16 часов A100 на конструкцию

Почему это важно

Обычно в робототехнике конструкцию робота проектируют вручную, а под готовое «железо» отдельно обучают систему управления, и удачная политика управления может упираться в неудачную конструкцию. Transformer Transformer сводит проектирование конструкции, оценку и управление в одну модель, которая обучается совместно на динамике робота. Это не просто демонстрация в симуляции: результат проверили на физическом роботе и получили измеримое улучшение (73% меньше ошибка слежения, 30% меньше максимальная скорость шарниров), редкий для академических работ по co-design случай, когда сгенерированную конструкцию действительно изготовили и протестировали на реальном железе.

Кому это важно

Работа адресована исследователям робототехники и co-design (совместного проектирования конструкции и управления), инженерам, которые проектируют роботов-манипуляторов под конкретную задачу (например, промышленную или складскую автоматизацию), и разработчикам платформ вроде ALOHA, для которых предложенный подход, способ быстро подобрать вариант конструкции под новую задачу без долгой ручной инженерии.

Как это применить

Схема использования, которую описывают авторы: пользователь предоставляет демонстрацию целевого движения конечного эффектора (например, запись движения человека через систему UMI) и функцию вознаграждения; модель генерирует конструкцию робота под эту задачу и сама же управляет ей, чтобы проверить дизайн. Метод протестирован на 11 типах роботов из открытой библиотеки MuJoCo Menagerie и на открытом наборе данных UMI по двурукой мойке посуды. В тексте проекта нет сведений о публичном релизе кода, лицензии или доступности модели, есть только страница проекта и 17-минутный видеоразбор.

Можно ли доверять

Результаты подкреплены не только симуляцией: сгенерированную конструкцию для ALOHA2 изготовили физически и сравнили с оригиналом на реальном роботе, с конкретными числами (73% и 30%). Метод также сравнили с двумя базовыми подходами, случайной генерацией и эволюционным оптимизатором CMA-ES, на 11 роботах разного масштаба и на отдельной задаче мойки посуды с отложенными на проверку траекториями. При этом сами авторы указывают на умеренную (не высокую) корреляцию r = 0,53 между их контроллером и эталонным RL-контроллером на сгенерированных конструкциях, то есть обобщение контроллера на новые, нестандартные конструкции пока работает не идеально.

Риски и подводные камни

Авторы прямо перечисляют ограничения. Токенизация RoboTokens пока охватывает только примитивную геометрию, без произвольных мешей, деформируемых объектов (тканей, мягких материалов как объекта манипуляции) и контекста окружающей сцены. Кросс-платформенный контроллер обобщается на новые конструкции не идеально (корреляция с эталоном r = 0,53), и на отдельных нетипичных конструкциях заметны сбои. Прирост качества от увеличения вычислений на этапе вывода выходит на плато примерно через минуту, в отличие от языковых моделей, где длительные рассуждения продолжают давать выигрыш. Наконец, получение обучающих данных для RL-контроллеров остаётся дорогим (около 16 часов на GPU A100 на одну конструкцию), и именно это авторы называют главным барьером для расширения метода на новые классы роботов.

«Как должен измениться этот робот, чтобы увеличить вознаграждение?»

— так авторы описывают то, что модель фактически «спрашивает у себя» в процессе Dynamics Self-Guidance