Синтез обучения с подкреплением и классического планирования: как выиграть соревнование CVPR Habitat Challenge 2023

Всем привет! Меня зовут Алексей Староверов, работаю научным сотрудником в AIRI и в составе нашей команды (вместе с Кириллом Муравьевым, Татьяной Земсковой, Дмитрием Юдиным и Александром Пановым) мы выиграли соревнование Habitat Challenge, которое проводилось в рамках крупнейшей конференции по компьютерному зрению CVPR 2023. Мы смогли эффективнее других команд научить робота навигироваться до целевых объектов в новых помещениях с использованием только RGB-D камеры, датчика GPS и компаса. Сейчас это является очень важной задачей при создании роботов-помощников, выполняющих задачи по инструкциям на естественном языке. В этой заметке я расскажу, как это у нас получилось.

Мотивация и постановка задачи

В последнее время у научного сообщества появился большой интерес к задачам воплощенного искусственного интеллекта (ВИИ / Embodied AI). Один из вариантов такого воплощенного интеллектуального агента — домашний ассистент, который не только может взаимодействовать с вами через голосовые команды, но и способен перемещаться по дому, взаимодействуя с окружающими объектами, выполняя бытовые задачи. Представьте себе сочетание таких устройств, как робот пылесос iRobot, станция с голосовым помощником и манипулятор (как Canadarm2 на МКС). 

Для реализации этой большой идеи, ее можно поделить на три группы задач: визуальная навигация (visual navigation), воплощенное зрение и язык (embodied vision-and-language) и манипуляция объектами (manipulation). 

Задача визуальной навигации до целевого объекта (Visual ObjectGoal Navigation) заключается в следующем: роботу нужно найти объект определенного типа (например, холодильник), доехать до него и остановиться рядом с ним. Именно такая задача была поставлена на соревновании Habitat ObjectNav Challenge 2023. Никакой карты и других подсказок роботу не давалось, ему нужно было выполнить задачу, полагаясь лишь на картинки с бортовой RGB-D камеры и данные о своем местоположении.

В задаче визуальной навигации до целевого объекта мы выделили две подзадачи: исследование пространства в поисках целевого объекта (Exploration) и доезд до цели (Goal Reaching). Обе этих подзадачи могут решаться как классическим подходом, с одновременным построением карты и локализацией (Simultaneous Localization and Mapping, SLAM), так и методами обучения с подкреплением (Reinforcement Learning, RL). 

В случае подхода, основанного на SLAM, робот в процессе движения строит карту помещений и наносит на эту карту целевой объект, как только обнаружил его. После того, как целевой объект нанесен на карту, робот строит до него маршрут по полученной карте и движется по этому маршруту. В случае подходов, основанных на RL, информация об окружающей среде вместо карты запоминается в скрытые слои рекуррентной нейронной сети (RNN). Нейросеть принимает на вход данные с робота и выдает действие, оптимальное в данный момент для решения задачи — проехать вперед, повернуть налево или направо, остановиться по достижении целевого объекта.

И SLAM, и RL имеют свои достоинства и недостатки. Например, RL-подходы позволяют быстро и эффективно объехать помещение за счет целенаправленного обучения на миллионах примеров. Однако рекуррентные нейросети склонны «забывать» информацию с предыдущих шагов, из-за чего робот может начать ездить по помещению кругами. В SLAM-подходах полная информация о всех помещениях хранится на карте, что позволяет объехать все помещения полностью без повторений и финишировать четко возле целевого объекта. Однако такие подходы чувствительны к качеству карты (которое зависит от наличия шума в данных с робота). Также, из-за ограниченного поля зрения камеры, на карту могут быть нанесены не все объекты, и робот может «споткнуться» о ненанесенное препятствие и застрять. Чтобы извлечь максимум из обоих подходов, соединив их сильные стороны, мы разработали комбинированный метод SkillFusion.

В SkillFusion помимо «сырой» информации с сенсоров используется информация о семантике сцены — распределении по категориям объектов пикселей на RGB изображении, полученном с камеры робота. Семантическая сегментация сцены позволяет получить точную информацию о местоположении объектов в ней, а также об их категориях для планирования перемещений робота. В соревновании Habitat не предоставляется информация о семантике сцены, поэтому возникает необходимость в точном распознавании объектов по RGB изображениям с помощью нейронных сетей для семантической сегментации. 

Про данные и симулятор

Соревнование проводилось на симуляторе Habitat, на множестве сцен из датасета Habitat Matterport 3D Dataset. Большая часть сцен представляла собой жилые помещения. Задача состояла в визуальной навигации до целевого объекта одного из шести типов: стул, диван, кровать, комнатное растение, телевизор и туалет. Если задуматься, то это почти все домашние объекты, с которыми каждый человек регулярно взаимодействует (не хватает кухни, конечно). 

Для выполнения этой навигационной задачи виртуальный робот (агент) получал на каждом шаге изображения и глубины с RGB-D камеры с ограничением глубины от 0.5 до 5 метров, а также точные данные о местоположении и ориентации относительно стартовой точки (GPS+Compass). Агент мог на каждом шаге выполнить действие одного из шести типов: проехать 25 см вперед (move_forward), повернуться на 30 градусов влево (turn_left), повернуться на 30 градусов вправо (turn_right), наклонить камеру на 30 градусов вниз (look_down), поднять камеру на 30 градусов вверх (look_up), остановиться и завершить эпизод (stop). Эпизод считался успешно завершенным, если агент выдал действие stop в радиусе 1 метра от целевого объекта, и целевой объект находился в прямой видимости. На выполнение давалось 500 секунд.

Чтобы сделать задачу поинтереснее, организаторы оснастили агента камерой с горизонтальным углом обзора в скромные 42 градуса и подвесили эту камеру на высоте 1.31 м. Чтобы охватить такой камерой все окружающее пространство, а также увидеть низкие объекты, агенту приходилось много вращаться на месте, наклонять и поднимать камеру. Также в действия был добавлен шум — например, действие move_forward двигало агента не ровно на 25 см вперед, а с небольшим случайным отклонением. Дополнительно задачу усложняли дефекты симуляционных сцен, такие как «дырки» в полу и стенах, а также зеркала и стеклянные стены в некоторых сценах.

Всего в симуляторе Habitat присутствуют 1000 сцен, из которых в 216 представлена разметка для инстанс сегментации. Наличие размеченных сцен позволяет собирать большие наборы данных для обучения сегментационных моделей нейросетей. В симуляторе присутствуют два типа разметки: около 1600 изначальных категорий объектов и соответствующие им 40 категорий объектов из набора данных Matterport 3D.

Перед нами стояла задача собрать данные в симуляторе для дообучения нейросетевой модели сегментации. Для этого мы попробовали различные стратегии выбора кадров, а также различные наборы категорий. Стратегия сбора датасета в случайных точках сцен показала лучшие результаты по сравнению с эвристическими стратегиями выбора точек обзора категорий объектов, представляющих навигационные цели. При этом оказалось, что выбор перечня категорий для разметки данных также оказывает влияние на качество распознавания семантики сцены. В финальном наборе данных для обучения модели сегментации мы использовали метод, комбинирующий псевдолейблинг для 150 категорий набора данных ADE20K и ground truth разметку. Для категорий объектов, общих для ADE20K и Matterport3D, использовались ground truth маски семантической сегментации, а для остальных категорий использовались маски, предсказанные с помощью эффективного трансформерного метода семантической сегментации — SegFormer. Общее количество изображений в наборе данных составило 511 тысяч изображений, из которых 27 тысяч случайных изображений использовалось в качестве валидационной выборки, а все остальные в качестве — обучающей.

Такой метод сбора данных для дообучения нейронной сети для сегментации позволил задействовать предобученные веса, т.к. многие SOTA методы для семантической сегментации валидируются на наборе данных ADE20K. Ещё одним достоинством категорий набора данных ADE20K является их разнообразие и более высокий уровень взаимоисключаемости по сравнению с категориями Matterport3D. В категориях Matterport3D различные объекты объединены в категории “objects”, “misc” и “appliance”, что приводит к сильному дисбалансу классов. 

Особенности SLAM

Наличие точных данных о позиции робота и глубине изображений позволяло строить идеально точную карту методами SLAM. Для этого мы сначала строили локальную карту в виде трехмерного облака точек методом обратной проекции с использованием карты глубин. Далее мы проецировали локальную трехмерную карту на плоскость и встраивали в глобальную карту препятствий. Препятствиями считались все объекты, расположенные на высоте от 25 до 150 см, разрешение карты составляло 7 см.

Для поиска целевого объекта использовался алгоритм PONI. PONI выбирает промежуточную цель на построенной карте с помощью потенциальной функции, предсказанной нейросетью. Потенциальная функция точки состоит из двух компонент — оценка площади, которую можно откартировать с данной точки, и оценка вероятности наличия цели рядом с данной точкой. Для предсказания второй компоненты на карту наносятся объекты 16 типов — помимо целевых объектов, также стол, плита, раковина и т.д., и полученная 16-слойная карта подается на вход нейросети. Нейросеть, обученная на сотнях тысяч примеров подобных карт, оценивает вероятность наличия цели по взаимному расположению объектов на построенном участке карты. Промежуточная цель для агента выбирается как точка с наибольшим значением потенциальной функции. Для движения к промежуточной цели использовался алгоритм Fast Marching Method (FMM).

После того, как на карту был нанесен один или несколько целевых объектов, агент строил по карте путь до ближайшего из них с помощью метода поиска пути А* и двигался к цели по построенному пути. Так как истинные метки объектов агенту не давались, то объекты на входных изображениях размечались с помощью нейросетевой семантической сегментации, которая имеет значительный процент ошибки. Из-за ошибки нейросети, агент мог нанести на карту ложный целевой объект (например, в случае, когда нейросеть приняла край дивана за стул, или картину на стене за телевизор). По мере приближения к такой ложной цели, она становилась лучше видна камерой агента, и нейросеть переставала помечать ее как цель, однако она уже была нанесена на карту, и агент доезжал до нее и завершал эпизод с нулевым результатом.

Чтобы решить проблему ложных целей, мы ввели фильтрацию семантической карты с помощью эрозии, накопления и затухания. Так, сначала все нанесенные на данном шаге на карту целевые объекты «истончались» с помощью метода `cv2.erode()`, а затем оставшиеся объекты расширялись обратно методом `cv2.dilate()`. Это позволяло отсечь мелкие выбросы семантической сегментации. Далее, значение семантической карты в каждой клетке оставшегося после эрозии целевого объекта увеличивалось на 1. Значения в клетках, не являющихся клетками целевого объекта, но попавших в поле зрения камеры, умножались на коэффициент α < 1. Клетка считалась содержащей целевой объект, если значение в ней превышало 1.

Особенности RL

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

Особенностью DDPPO над PPO является эффективная работа с распараллеленными средами. Для максимизации этого эффекта и обучения самого сложно RL навыка Exploration, мы взяли симулятор BPS вместо Habitat. Его авторы реализовали рендеринг многих запросов агента в одном процессе симулятора за один шаг, вместо создания копии всего симулятора под каждый процесс агента. Это позволило увеличить скорость сбора траекторий для обучения на два порядка (13300 vs 140 FPS на одной видеокарте). Рисунок ниже показывает как происходит взаимодействие симулятора со стратегией агента. Симуляция сенсоров агента делится на два этапа, определение состояния агента и рендеринга изображений камер из этих состояний. Первый этап считается на CPU, второй на GPU с использованием библиотек Nvidia Vulkan. Сцены для этого симулятора мы брали также из набора HM3D, что бы для агента не было отличий входных данных от симулятора Habitat. Использовать симулятор Habitat приходилось для обучения RL навыка GoalReacher, так как ему на этапе обучения нужен сенсор сегментации, который на данный момент не реализован в симуляторе BPS.

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

Объединение многих навыков агента в единую архитектуру обучения позволило получить более репрезентативные LSTM-энкодеры, так как каждый навык требовал от агента разного поведения, вся необходимая информация для которых как раз оставалась в LSTM блоках. Так, навык Exploration должен иметь хорошее понимание о пространственной структуре сцены, для Flee наиболее важно понимать направление движения относительно прошлых состояний, а PointNav должен больше сосредотачивается на нахождении кратчайшего пути к точке. Функция вознаграждения каждого навыка мотивирует эти виды поведения. 

В качестве энкодеров изображений с камеры, мы взяли модель CLIP, так как по количеству обучаемых параметров эти блоки самые большие, а предобученный CLIP может извлекать достаточно хорошие векторные представления и не будет переобучаться под особенности симулированных сцен. Чтобы показать преимущество данного выбора, мы сравнили обучение навыка GoalReacher c замороженным CLIP против обучаемых ResNet блоков. С CLIP агент учится гораздо быстрее, оставляя преимущество в генерализации на все виды сцен. Чтобы показать, что агент использует именно репрезентации CLIP, а не полагается на GPS или Depth сенсоры (которые обрабатывались обучаемыми не через CLIP блоками), мы обучили агента только на CLIP эмбеддингах от RGB сенсора, и агент успешно справился с этим.  

Как и для чего комбинировать RL и SLAM

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

Из-за высоты и малого угла обзора камеры не все низко расположенные препятствия наносились на карту — агент мог не увидеть условную табуретку, уткнуться в нее и застрять. Чтобы это учесть, мы также наносили на карту препятствие перед агентом в случае его застревания (т.е. когда позиция не менялась после выполнения действия move_forward). Однако в случае протяженных невидимых препятствий (например, стеклянная стена) или когда препятствия окружали агента с нескольких сторон, выход из препятствий таким способом занимал очень долгое время, снижая эффективность навигации. Поэтому в случае застревания агента мы наклоняли камеру и переключались со SLAM на RL на 10 шагов для выхода из застревания.

Также, из-за ограничения карты глубин в 5 м, агент со SLAM-подходом часто проезжал мимо целевого объекта, т.к. видел его издалека и не нанес на карту. Поэтому для предварительного нацеливания и подведения к целевому объекту мы использовали RL. Как только агент подъезжал к целевому объекту достаточно близко и наносил его на карту, мы переключались обратно на SLAM для точного достижения цели. Такой комбинированный подход позволил нам максимально эффективно достигать целевого объекта в условиях ограниченной глубины и неточной семантической сегментации.

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

def agent_step(observation):
    goal_semantic_mask = get_semantic(observation['rgb'])

    update_map(observation,goal_semantic_mask)
    update_RL_Exploration(observation)
    update_RL_GoalReacher(observation)

    classic_choose_subgoal_on_a_map(observation)
    classic_create_path_to_map_subgoal(observation)

    found_goal = sum(goal_semantic_mask)>0
    goal_is_mapped = sum(map.semantic_map)>0

    RL_explore_score = get_RL_explore_score(RL_explore.critic)*(not found_goal)
    Classic_explore_score = get_Classic_explore_score(map,path)*(not goal_is_mapped)
    RL_GoalReacher_score = get_RL_GR_score(RL_GR.critic)*found_goal*(not goal_is_mapped)
    Classic_GoalReacher_score = get_Classic_GR_score(map,path)*goal_is_mapped

    Available_skills = [RL_exploration, Classic_exploration, RL_GoalReacher, Classic_GoalReacher]
    Skill_values = [RL_explore_score, Classic_explore_score, RL_GoalReacher_score, Classic_GoalReacher_score]

    Skill_to_use = Skill_values.index(max(Skill_values))
    action = Available_skills[Skill_to_use]()

    return action

Про семантическую сегментацию

В качестве нейросетевого подхода для семантической сегментации мы выбрали OneFormer — один из SOTA методов для семантической сегментации. Особенностью данного метода является использование единой трансформерной модели для семантической, инстанс и паноптической сегментации.

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

Серьёзным недостатком модели для применения в задаче визуальной навигации является её размер (219М параметров) и скорость работы (~220 мсек на инференс на одном изображении на GPU RTX 2080 Ti). Поэтому перед окончательным выбором данного метода для семантической сегментации в нашем подходе мы провели сравнение OneFormer с более легковесной моделью SegFormer B5. 

Мы сравнили модели OneFormer и SegFormer B5, обученные на одинаковом наборе данных, состоящим из 160 тысяч изображений из симулятора Habitat, размеченных на 40 категорий Matterport3D. При этом т.к. для решении задачи визуальной навигации нам необходима только семантическая сегментации, то обучение OneFormer производилось в режиме с фиксированной задачей семантической сегментации. По результатам сравнения (см. Таблицу) OneFormer показал значительно более высокое качество сегментации, поэтому мы остановились на данном методе. 

Сравнение качества сегментации методов OneFormer и SegFormer B5.

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

Влияние качества семантической сегментации на метрики качества навигации

Как видно из данной таблицы, выбор метода семантической сегментации вносит существенный вклад в качество визуальной навигации до цели.  

Ещё одной отличительной особенностью OneFormer по сравнению с SegFormer является использование парадигмы классификации масок для семантической сегментации. Вместо классификации каждого пикселя выход OneFormer состоит из предсказания бинарных масок и их распределения по классам, которые затем объединяются в многоканальную маску для получения семантической сегментации. Такой подход увеличивает качество формируемых масок, однако критическим становится вопрос их правильной классификации. При создании обучающей и тестовой выборки желательно, чтобы сцены, в которых собраны данные не пересекались. Это позволяет оценить качество сегментации на разных объектах одинаковых категорий, при этом снизив вероятность переобучения на валидационной выборке. 

Подробнее об особенностях архитектуры SkillFusion можно узнать из видео ниже:

Особенности соревнования

В соревновании тестирование решений делилось на три фазы: Minival, Test-Standard и Test-Challenge. Фаза Minival состояла из 30 эпизодов и предназначалась для быстрой предварительной проверки решения. Фазы Test-Standard и Test-Challenge состояли из 1000 эпизодов. На фазе Test-Standard можно было отправить до 100 попыток, но одновременно в системе могло тестироваться только три попытки. На фазу Test-Challenge можно было отправить только 5 попыток, при этом две попытки одновременно тестироваться не могли. Результаты оценивались по фазе Test-Challenge в закрытом режиме. Фаза Test-Standard была открыта, но результат можно было скрыть из таблички по желанию участников. В итоге все участники, кроме нас и еще одной команды, скрыли свои результаты.

Качество решений измерялось по трем метрикам: Success, SPL и SoftSPL. Метрика Success представляла собой процент успешно завершенных эпизодов. Метрика SPL (Success weighted by Path Length) представляла собой среднее отношение длины кратчайшего пути до цели к длине пройденного агентом пути, домноженное на успешность завершения. Метрика SoftSPL считалась так же, как SPL, но вместо бинарной успешности использовалось значение прогресса в направлении цели.

Для участия в соревновании нужно было отправить свое решение в виде Docker-контейнера, в котором прописан запуск кода агента. Размеры образа для отправки со всеми установленными библиотеками и весами нейросетей не должны были превышать 40 Гб. Отправка происходила через систему EvalAI.

Основная сложность отправки решений состояла в том, что тестирование на фазах Test-Standard и Test-Challenge занимало несколько дней, а очередь на тестирование могла длиться больше недели. Таким образом, за неделю до дедлайна соревнования шансов отправить решение на тестирование, получить результат, исправить и отправить заново уже не было — мы успевали отправить только одну попытку. Даже очередь на Minival, предназначенный для быстрой предварительной проверки решения, занимала сутки или больше.

Также общее ограничение по времени тестирования на все 1000 эпизодов составляло 48 часов, что намного меньше, чем 500 000 секунд — сумма ограничений по времени на каждый эпизод, из‑за чего сначала наши решения не прошли по времени, и мы думали ставить «отсечку» по времени на каждый эпизод. За день до дедлайна нам по запросу расширили лимит времени с 48 до 96 часов, и мы смогли отправить решение на Test‑Standard и Test‑Challenge без отсечки. После без малого двух недель ожидания конца тестирования, мы увидели свои результаты.

Результаты

Результаты на фазе Test-Challenge соревнования Habitat Navigation Challenge 2023 получились следующие: Success 0.53, SPL 0.27, SoftSPL 0.33. С такими результатами мы и заняли первое место. Результаты на фазе Test-Standard получились похожие: Success 0.53, SPL 0.28, SoftSPL 0.34. На втором месте в таблице по фазе Test-Standard оказалась команда ICanFly с чуть меньшим значением SPL и значительно меньшим значением Success.

Для отбора лучших решений и их отправки в систему соревнования мы отобрали выборку из датасета HM3D, содержащую 108 эпизодов на 36 сценах. На этой выборке мы прогоняли локально разные варианты решений и оценивали качество по метрикам Success, SPL и SoftSPL. Также мы оценивали долю эпизодов, завершенных у ложной цели (т.е. когда агент выдал действие stop не возле целевого объекта). На этой выборке мы сравнили итоговую версию комбинированного подхода с чистым SLAM-подходом и чистым RL-подходом.

Результаты подходов SLAM и RL по отдельности в сравнении с комбинированным подходом на локальной выборке из 108 эпизодов

Комбинированный подход (SkillFusion) значительно опередил SLAM и RL-подходы по всем метрикам. Особенно большой прирост наблюдался в метрике SPL, за счет выхода из застревания с помощью RL и эффективного исследования пространства с помощью потенциальных функций PONI. Доля завершений  у ложной цели также значительно уменьшилась за счет фильтрации семантической карты. 

Код нашего подхода мы сделали общедоступным по ссылке.

Для нас участие в соревновании Habitat Challenge стало хорошим опытом, как победители соревнования мы были приглашены на воркшоп Embodied AI на конференции CVPR’2023. И в ближайшем будущем мы хотим внедрить получившиеся модели и решения в общую архитектуру системы управления воплощенным агентом, выполняющим сложные языковые инструкции в виртуальным и реальных средах.

Код модели и процесса обучения будет опубликован в открытом репозитории.