Retrieval-Augmented Decision Transformer: External Memory for In-Context RL

really cool idea по поводу ин-контекст рл, которую захотелось заимплементить (в том числе оттого, как понятно авторы описали свою идею)

в основном работает интуиция, что чем сложнее среда, тем бОльший контекст нужен модели для перформанса на near-optimal/optimal level (особенно когда речь идет про ПОМДП). и вот тут, как и в нлп, возникает загвоздка - улетаем в требования по памяти. и здесь довольно много есть путей, как такое решать - как костыльные, так и более-менее красивые 👀

авторы же попробовали прикрутить индекс датабазу эмбеддингов и получилось очень даже прикольно - но есть нюансы

- будем делать кнн по внешней базе эмбеддингов и затем подавать выданные через кросс аттеншн, окей: откуда взять репрезентативные эмбеддинги? можно претрейнить Decision Transformer, но так же использовать и ЛМки, что очень красиво
- делается это при помощи рандомной матрицы (которая уже использовалась здесь и здесь), чтобы подогнать размерности к хидденам ЛМки
- а эмбеддинги языковых моделей остаются релевантными и для рл данных из-за своей плотной полисемантичности (и потому что первый автор уже разбирался в этом вопросе)
- при этом при контексте модели в С токенов в индекс базе для ретривала сохраняются эмбеддинги подтраекторий длиной 2С, чтобы агент мог так же получать интуицию о последствиях действий в будущем
- ретривятся эмбеддинги по схожести, но так же потом отбираются еще дополнительно по понятию utility, которая может быть связана либо же с повышенным ретерном траектории, либо по проверке того, одинаковая ли таска с траекторией-запросом (query)

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

есть только несколько вопросов - почему на помдп (например в прокгене или мейзраннере) в некоторых местах дпт перформит лучше ад, в то время как дпт заявляется как MDP-like algorithm, может ли это говорить о чем-то другом кроме как того что не совсем правильно завели какой-то из алгосов + на даркруме (мдп) ад лучше по резам чем дпт (здесь предполагаю, что оставили результаты с 2 эпизодами в качестве контекста, в то время как дпт может по теории перформить хоть с single-episode context, а ад из-за способности к помдп нет)

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

👀LINK

подслушано здесь