Structured State Space Models for In-Context Reinforcement Learning

Трансформеры в ин-контекст лернинге в рл (и нлп) опробовали, почему бы не опробовать и ссм?

Отлично, только хорошо бы не просто их запустить из-под коробки, а как-то внести в структуру модели понимание того, какое действие хорошо в новой таске, а какое плохо

А как этого достичь, да так, чтобы не сохранить скорость stat space models? Возможно, понадобится как-то видоизменить параллельный скан, которым так славится ссм

Так и есть - авторы добавляют в ассоциативный оператор их ссм, которая называется S5, флаг о том, что эпизод закончился (а это классика поскольку в рл средах этот флаг поступает на вход при каждом шаге) -> сохраняется 'ссм-ность', и можно проверять, могут ли осуществлять ин-контекст рл не только трансформеры и рнн

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

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

И это работает! Более того, это повышает вариативность тренировочного датасета (потому что мы более не привязаны к размерностям состояний и действий) + это позволяет адаптироваться к ООД задачам

супергуд! разве что пока возможно такое только на джаксе сделать

👀LINK