H-JEPA учит модели мира строить планы на разных уровнях абстракции

Зачем моделям мира JEPA нужна иерархия

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

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

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

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

H-JEPA решает эту проблему, обучаясь строить представления на разных уровнях. Для планирования на большом горизонте можно использовать информацию, отличную от той, что нужна для управления ближайшими действиями.

Как H-JEPA обучается и планирует

H-JEPA обучают на последовательностях наблюдений и связывающих их действиях. Нижний уровень кодирует визуальные наблюдения и предсказывает следующее скрытое состояние на основе предыдущих состояний и действий. Уровни выше получают представления нижнего уровня и предсказывают переходы через более длинные интервалы.

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

Исследователи обучали иерархию целиком. Каждый уровень сравнивает предсказанное представление с представлением будущего, полученным из наблюдений. Механизм регуляризации не даёт модели выдавать одно и то же представление для любых наблюдений. Сигналы обучения с верхних уровней также проходят через кодировщики нижних.

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

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

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

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

В лабиринте верхний уровень может сосредоточиться на движении к месту назначения, а нижние — на коротких переходах и физических движениях.

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

Проверка H-JEPA

Исследователи оценили H-JEPA в четырёх симулированных средах для навигации и манипуляции: FourRoom Distractors, Visual AntMaze, Push-T и OGBench Cube. Модель сравнили с LeWM — одноуровневой JEPA — и HWM, другой моделью мира, которая строит иерархические планы в общем скрытом пространстве.

В FourRoom, AntMaze и Cube добавление уровней — вплоть до трёхуровневой иерархии — обычно повышало долю успешных попыток и снижало вычислительные затраты на планирование. Отдельные представления давали наибольшее преимущество в AntMaze: трёхуровневая H-JEPA справилась почти вдвое успешнее HWM.

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

Затем команда испытала H-JEPA на DROID — наборе видео реальных манипуляций роботов с разными предметами, освещением и фоном. При обучении только на предсказаниях модель, как правило, сохраняла неподвижные детали сцены, но теряла информацию о движущемся роботе. Исследователи добавили цель «обратной динамики»: предсказывать действие, связывающее два наблюдаемых состояния. Это помогло сохранить сведения, важные для действий. В офлайн-планировании H-JEPA точнее базовых моделей и при этом требует от планировщика меньше вычислений.

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

Исследователи предлагают связать представления верхних уровней с языком. Это могло бы позволить агентам получать цели в виде инструкций, а не целевых изображений. Пока H-JEPA показывает, что для ближайших движений и отдалённых задач можно обучать разные представления.