Generalization to New Sequential Decision Making Tasks with In-Context Learning [2023]

Продолжаем разбираться в In-Context лоре. Играться с In-Context трансформерами (ICT) можно разными способами. Мы уже разбирали два из них:

1) Cценарий Gato, когда на вход подаётся траектория в рамках одного эпизода в среде, и модель учат предсказывать действие эксперта в данном контексте.
2) Cценарий Algorithm Distillation / AdA, в котором алгоритм получает на вход не одну траекторию, а всю историю взаимодействия алгоритма со средой, и учится либо предсказывать действие алгоритма, либо на награду с помощью RL в квадрате.

Авторы данной работы хотят получить ICT, быстро адаптирующийся к новой задаче, но облегчают проблему, допуская, что у нас есть несколько экспертных траекторий в данной среде. Это близко к in-context решению задач, которое мы проворачиваем с LLM, когда добавляем в контекст несколько примеров правильных пар вход-выход, можно даже сказать, что близко к сценарию ARC.

При этом, авторы говорят о сильной out-of-distribution генерализации, т.е. тестированию алгоритма на средах, совсем непохожих на обучающие. В общем, цели я одобряю. В чём заключается подход авторов?

На вход трансформеру подаётся текущая траектория агента в среде - состояния, действия, награды. Помимо этого, туда подаются несколько экспертных траекторий из этой же среды. В отличие от Algorithm Distillation / AdA, это не история обучения алгоритма на данной задаче. Это просто засэмплированные из датасета траектории уже обученного PPO, применённого на этой задаче.

В рамках сред существуют "уровни" - конкретные расположения всех объектов. Авторы всегда помещают как минимум одну экспертную траекторию из того же самого уровня, на котором идёт тестирование, в контекст. Учитывая, что среды сами по себе детерминированные (за исключением случайного залипания действия), это довольно сильная подсказка.

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

Другой бейзлайн - это чистый Gato, но они его называют Behavioural Cloning. К сожалению, среди бейзлайна нет самого эксперта (PPO), не считая его калеку-представителя в виде Hashmap.

В результате, ICT с экспертными траекториями в контексте побеждает и Gato, и Hashmap, при этом сам Gato больше, чем на половине графиков, проигрывает Hashmap. Получается, что ICT получилось обобщить за пределы задач, на которых его обучали, но все они были настолько простыми, что особо порадоваться я не могу. При этом, нельзя сказать, что авторы придумали что-то сильно новое и чего-то качественно добились. Их подход вряд ли позволяет получить обобщающийся ICT с большей вероятностью, чем Algorithm Distillation / AdA.

@knowledge_accumulator