Эпизод 1. Трансформер идёт в Atari Frostbite

Какие информационные процессы протекают в нашем мозгу? Какие цели он преследует? Какой моделью данных и какой алгеброй можно передать его работу? Почему мы подвержены разным когнитивным искажениям? Зачем нам сон и можно ли без него? Как использовать наш мозг на полную катушку?

Меня зовут Михаил Кочеров и я хочу понять, как работает наш мозг. А что значит понять? Вот знание о морфологии мозга или о его развитии от момента зачатия до момента смерти — это понято или нет? Каждому своё. Для меня понять — значит, уметь воспроизвести пусть и в упрощённом виде. Постичь ключевые механизмы и знать, как из них собрать работающую систему, воспроизводящую интеллектуальные свойства мозга в целостном виде.

Люди, птицы, рыбы и т.д. — это автономные интеллектуальные субъекты с мозгом. Познают мир в меру своих возможностей, адаптируются к нему, вырабатывают оптимальное поведение на основе проб, ошибок и редких, но столь желанных наград. Это не что иное, как обучение с подкреплением, или Reinforcement Learning (RL). Отсюда идея: разрабатываем RL-агента для нетривиальной проблемы и по ходу сталкиваемся с теми же трудностями и задачами, с которыми столкнулась живая природа, развивая мозг. А следовательно, решая эти задачи, познаём и сам мозг, его устройство и механизмы.

Такой нетривиальной RL-проблемой видится разработка агента, который научается играть в компьютерные игры. Агент должен сам разобраться с тем, как воспринимать внешнюю среду (например, видеопоток), постичь внутреннюю логику игр, выработать политику поведения для каждой и поддерживать память, чтобы не забыть прошлый опыт. Для этих экспериментов шикарно подходит Atari Learning Environment (ALE) [8][9] + Gymnasium [10].

В ALE много разных игр: от простых аркадных (типа Breakout) до стратежно многоходовых (типа Montezuma’s Revenge). Я выбрал Atari Frostbite. Эскимос должен прыгать по льдинам, уворачиваться от врагов, собирать ледяной домик (иглу) и потом в нём прятаться, чтобы пережить заморозки. Игра очень динамичная: враги меняют своё поведение и скорость, форма и движение льдин мутируют, день сменяется ночью и наоборот, и т.д. Награда распределена неравномерно: есть небольшие очки за конкретное действие (съел рыбу, собрал кирпичик иглу), но гораздо больше обламывается агенту, когда он проходит уровень. Налицо нестационарная игра с разреженной наградой и иерархическим целеполаганием — всё как в жизни!

Итак, план такой:

Важная оговорка

Вопрос. А почему бы не взять какую-нибудь LLM типа GPT-6 Astra или Claude Sonnet для получения такого играющего агента?

Ответ. Разные цели. Если надо решить задачу здесь и сейчас для бизнеса, то так и надо сделать — взять готовый инструмент и прикрутить его. Только тогда не получится проникнуть в суть происходящих процессов. Мы просто сыграем в интеграцию а-ля CRUD: возьми данные из игры, передай в модель (чёрный ящик), верни ответ обратно и не расплескай по пути.

Хотя сейчас даже и CRUD не нужен — Claude Code или Openclaw сами во всём разберутся, и останется только объявить себя владельцем результата. В этом же путешествии хочется увидеть и познать скрытые механизмы мозга, а раз так, то придётся идти более тернистым путём и обдирать себе бока =)

Почему 136k+ очков?

Frostbite — вытоптанная площадка в RL теме и включена в известные бенчмарки:

• Atari-57 – “насколько ты силён, если учиться тебе не жалко?” В пуле 57 игр и щедрый бюджет: стандартный протокол даёт 200 млн кадров на игру (~920 часов игры), а чтобы обойти человека на всех играх, уходят на миллиарды кадров. Мерят потолок мастерства: до чего агент дотянется, если дать ресурсов от души.

• Atari-100k – “насколько ты быстро соображаешь, если данных почти нет?” В пуле 26 игр и всего 100 000 шагов или два часа реальной игры, один вечер за приставкой. Мерят сэмпл-эффективность: сколько мастерства выжимается из крох опыта.

Свежие цифры такие:

_{* шагов = кадров / 4 (frameskip=4)}

У меня нет своего датацентра с сотней другой жирных GPU, поэтому бюджет в миллиарды кадров для меня закрыт. А вот про 200 млн ещё как-то можно думать.

Короче, побить SOTA модель MEME – заманчивая абмиция!

Ещё важная оговорка

Вопрос. А почему тогда не взять за базу ту же MEME и не развить её?

Ответ. Это тоже срезание углов, только более тонкое, чем если поручить работу современным LLM. Важно самому пройти путь от начала и до конца, иначе проблемы будут замаскированы уже готовыми ответами. Надо столкнуться с проблемами, прочувствовать их, попытаться самому решить – только так идёт познание. Как говорил Х.Д. Борхес: “Нет такого интеллектуального упражнения, которое в итоге не принесло бы пользы” (“Пьер Менар, автор «Дон Кихота»”).

При этом отказываться от заимствования идей или запчастей тоже глупо – дажи при оголтелом заимствовании в конце концов всё равно выйдет как-то своя история =)

Азбука PPO

Будем использовать Proximal Policy Optimization (PPO) [1] — популярный RL-алгоритм. Большой плюс в том, что он хорошо масштабируется, и цикл обучения получается быстрее, чем, например, у DQN [2].

PPO и DQN как сборка пазла

DQN — бережливый: каждую деталь (шаг в игре) кладёт в коробку и возвращается перебирать её снова и снова, поэтому выжимает из деталей больше. PPO — расточительный: схватил горсть деталей, разложил, а старые выбросил, чтобы взять свежие. Когда что важно:

• Дорогой шаг (реальный робот, медленный симулятор, платный API) → лучше выбрать DQN, SAC (off-policy).

• Дешёвый и параллелизуемый шаг (Atari, быстрые симуляторы) → выигрывает PPO (on-policy) с массовой параллельностью. Нам подходит.

Для работы PPO нам нужна сама модель (агент) и батарея игровых окружений (инстансов игры). Чем больше батарея, тем быстрее получится собрать траектории. Например, 32 одновременных окружения — хорошее стартовое число. Ключевые методы в интерфейсе окружения очень простые:

• reset() — сбросить игру (типа нажали кнопку “Reset” на игровой приставке). Делается в самом начале, или когда случился game over, чтобы начать заново.

• step(action) — выполнить очередной шаг. 99% времени будет вызываться этот метод.

Каждый из этих методов возвращает observation — то, что игрок должен увидеть на экране. Для ALE — это RGB-картинка размером 210×160. Т.к. у нас N окружений, то и на вход надо подать массив из N действий, а взамен мы получим N картинок.

Вместе с observation метод step ещё возвращает награду (reward) и флаг done — закончилась ли игра или нет. Для Frostbite reward будет ненулевым, когда выполняется какое-то полезное промежуточное действие (например, агент добыл кирпичик для иглу), или когда агент прошёл уровень и началась его зачистка (таймер времени убавляется, иглу разбирается, очки начисляются). В остальное время reward равен нулю, т.е. агент что-то делает, но ничего не получает.

На стадии раскатки мы берём и прогоняем нашу модель шаг за шагом через все эти окружения одновременно. На вход модели подаём наблюдение (что сейчас на экране). Для этого у модели должен быть какой-то входной интерфейс по приёму картинок — Encoder. На выходе модель должна выдать action (какую кнопку надо нажимать, например, UP или FIRE) и value (сколько в данной игре получится заработать очков). Для этого в модели должны быть выходные интерфейсы Action Head и Value Head. По сути мы делаем так, чтобы агент, каким бы глупым или умным он ни был, играл в игру и при этом на каждом шаге возвращал оценку очков, которые он может заработать. В итоге накапливаются кортежи: (наблюдение, действие, награда, value, done). Набор кортежей называется траекторией. Если игра в каком-либо окружении закончилась, мы эту игру сбрасываем и всё заново — получается непрерывный майнинг траекторий.

Вместе с траекториями на раскатке ещё рассчитываются advantages — то, насколько исполненное действие оказалось лучше ожиданий агента. Например, агент думал, что, исполнив действие A, получит 10 очков награды, а получил 50. Значит, advantages будет это учитывать в виде положительного коэффициента. В обратную сторону тоже — если исполненное действие принесло меньше, чем планировалось, то коэффициент будет отрицательным.

Раскатка закончилась, переходим к оптимизации. Здесь берём и за несколько эпох пробегаемся по накопленным траекториям и advantages. Выходной интерфейс агента Action Head выдаёт вектор. Каждый элемент вектора соответствует какой-то кнопке (FIRE, UP, DOWN и т.д.). Поэтому размерность вектора равна количеству кнопок, которое игра поддерживает, например 18 или 6. Значение элемента вектора — это ненормализованная вероятность (логит). Поэтому цель оптимизации так крутить эти логиты, чтобы наиболее вероятным становилось то действие, которое соответствует более высокому advantage. Или наоборот, если advantage отрицательный (т.е. плохое было действие), то надо крутить в обратную сторону. Эти advantages — как маяк: показывают, куда надо крутить политику, чтобы агент действовал лучше и лучше. Но сильно резко крутить логиты тоже нельзя, т.к. в RL резкие изменения политики могут её необратимо поломать — менять политику нужно мелкими шагами. Поэтому в PPO встроен “предохранитель”, или clipping [1]. Мы насильно обрезаем направление изменения, например коэффициентом 0.2.

Чтобы правильно считать advantages, нужно правильно считать value. На оптимизации мы одновременно с Action Head оптимизируем и Value Head. Для этого берутся value из кортежа траектории и сопоставляются с тем, сколько реально агент заработал по траектории. Это расхождение между “агент думал, что заработает” и “агент по факту заработал” используется для уточнения Value Head.

Интерфейсы Encoder, Action Head и Value Head — суть части одной модели (в которой ещё много скрытых частей), поэтому они крутятся все сразу и вместе. Т.е., подкручивая логиты и предсказание value, мы сразу оптимизируем скопом и всю модель: заставляем её зрение (Encoder) выделять значимые элементы на сцене, Value Head — учиться правильно предсказывать награду, Action Head — какие действия в той или иной ситуации нужно предпринимать.

Резюме

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

Лейтмотив модели агента (мозг)

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

Каждое наблюдение — это картинка. Сырая картинка ALE — RGB 210×160 (uint8). Энкодер преобразует эти картинки в вектора более низкой размерности — латентные вектора, или латенты. Например, кадр 210×160 (uint8) конвертируется в латент размерности 384 (float). Сжатие примерно в 66 раз (3x210x160x1 / 384×4).

Что находится в латенте наблюдения?

Работает алгебра высоких размерностей. Возьмём даже размерность 384. Можно сгенерировать два таких случайных вектора — и о, чудо! — они будут практически ортогональны друг другу. То есть для размерности 384 можно нагенерить миллионы, если не миллиарды почти ортогональных векторов [11]. Каждый вектор — это независимое направление или ось, которое может описывать ту или иную фичу. Например, положение игрока, сколько кирпичей иглу собрано, счёт и т.д. Можно взять несколько десятков таких фичей и сложить в латент, и всё там прекрасно будет храниться и оттуда восстанавливаться. Конечно, тысячу векторов в один вектор уже запихнуть не получится, т.к. пусть слабая, но всё-таки корреляция осей кумулятивно внесёт много шума; но несколько десятков — легко. Поэтому вектор размерности 384 может легко описывать то, что происходит на сцене.

Хотя латенты и упорядочены физически в памяти, для трансформера они всё равно выглядят как просто мешок. Всё из-за того, что трансформер оперирует контентом — ему всё равно, по какому адресу латент расположен. Чтобы трансформер мог корректно различать, где какой латент во времени, их нужно пропатчить — к латентам добавляется вектор с временным кодированием (temporal encoding) [6]. Теперь трансформер готов потреблять данные. Он берёт набор латентов с временным кодированием и прогоняет через свои слои. На каждом слое отрабатывают последовательно MultiHeadAttention и MLP [6]. MHA работает сразу со всеми латентами, вычисляет отношения между ними — латенты как бы разговаривают друг с другом (cross-talk) и обмениваются информацией [6]. За счёт этого трансформер научается понимать физику игры, с какими скоростями кто движется, и т.д. В MLP-слое накапливается экспертиза — что делать в конкретной ситуации [7]. Проходя между слоями, понимание сцены мутирует от более простых вещей к более абстрактным.

На выходе трансформера получаем также набор латентов, но обогащённых опытом. В них уже сидит вся информация, необходимая для принятия решения. Рабочая схема — это просто потребовать, чтобы эта информация осела в последнем латенте, который соответствует последнему наблюдению. Value Head и Action Head берут этот латент и извлекают из него информацию:

• Value Head извлекает оценку награды — сколько можно получить до конца игры;

• Action Head рассчитывает вектор логитов, из которого потом сэмплится наиболее вероятное действие и отдаётся в окружение.

Почему ставка на трансформер

1. заточенность на работу с последовательностями. Здесь у нас как раз последовательность наблюдений (точнее, их латентов);

2. отсутствие состояния (stateless). Скорость тренировки в разы, если не на порядки, выше, чем если бы использовался какой-нибудь рекуррентный мотив;

3. отдельные слоты для экспертизы (MLP) и для вычисления отношений (MHA). По идее, ёмкость для впитывания опыта огромна;

4. на задачах обработки языка показывают себя феноменально, что указывает на их “умение” извлекать полезные абстракции и делать глубокие обобщения. Прямо как наш мозг.

Хочется всего и сразу: Transformer-XL

Вдохновляемся статьёй “Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agents” [5]. Авторы взяли архитектуру Transformer-XL и прикрутили её через PPO для научения агента играть в игры, где нужна длинная и точная память. Например, в одной из игр агент должен пробираться по бесконечному лабиринту, который генерируется на лету. Сама архитектура Transformer-XL решала задачу, как дёшево расширить контекстное окно трансформера без квадратичной сложности [4].

Особенно приятно, что в библиотеке CleanRL нашлась референсная реализация этой концепции [3] — не надо изобретать велосипед и можно быстро стартануть. Получается картина ниже.

В ванильном мотиве на каждый шаг нужно сделать одно и то же: взять цепочку наблюдений, прогнать через энкодер, потом отдать трансформеру. При обработке нового наблюдения надо обновить буфер цепочки по FIFO-принципу и повторить. Налицо повторная обработка одних и тех же входных данных. Если длина входной цепочки 10 наблюдений, то на каждом шаге мы будем обрабатывать 90% старой информации. Неэкономно и медленно.

В концепции Transformer-XL латенты кешируются, а не вычисляются на каждом шагу. Это очень здорово и позволяет замахнуться на длину входной цепочки в сотню латентов и больше. Правда, в реальном прогоне окно пришлось урезать до 16 наблюдений — сотня обрабатывается очень долго. Каждые 4 кадра эмулятора сливаются в один для скорости, поэтому 16 наблюдений соответствуют 2.1 секундам игрового времени. Не густо, но Frostbite — игра полностью наблюдаемая. Здесь нет, как в Starcraft, невидимых областей, нет тумана войны, все враги — всё на экране. Поэтому даже пары секунд истории должно быть достаточно, чтобы осознать происходящее и правильно среагировать.

Сборка готова: https://github.com/kocherovms/neurolab/blob/main/17_rl/17d_ppo_trxl_atari_mp_01.ipynb. Отрезвление быстрое. Оно как бы работает, но быстро обнаруживается потолок. Агент научается прыгать по льдинам и разбирать их, чтобы продлить свою жизнь. Но в собранный иглу не заходит. Заработок агента – 526 очков (лучший запуск 17d_ppo_trxl_atari_mp_01/3) — это очень мало.

Скорее всего проблема с обучаемостью кроется в закешированных латентах. И тут две проблемы:

1. чтобы их закешировать и переиспользовать, их надо заморозить, т.е. отцепить от графа вычисления градиентов. В итоге они работают только на read-only, а на оптимизации ведущую роль играет только последний латент — градиенты как бы текут только через него. Если бы латенты были разморожены, то они бы все участвовали в расчёте градиентов, что способствовало бы более быстрому познанию агентом внутренней логики игры, а так агент становится очень близоруким;

2. десинхронизация замороженных латентов с весами модели. На шаге оптимизации после первой же эпохи веса модели обновляются, и закешированные латенты уже не совместимы с ними. В той или иной степени. После 2-й или 3-й эпохи расхождение становится ещё больше. В итоге эти закешированные латенты превращаются в шум и никак не помогают.

В исходной работе по Transformer-XL авторы использовали комбинацию: замороженные латенты + достаточно длинная живая цепочка, что нивелировало вышеописанные проблемы. Здесь же живая цепочка выродилась до одного элемента, и проблемы показали себя в полный рост.

А почему тогда работало у ребят из “Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agents”? [5] Скорее всего дело в стационарности. Динамика и объём происходящего в Memory Gym не меняется так быстро, как в Frostbite. Размерность наблюдения (картинка) тоже несравненно меньше. Поэтому вероятнее всего, что проблемы кеширования и десинхронизации проявляются, но в меньшей степени.

Каждому своё

Решения, работающие в одних условиях, не всегда будут работать в других. Когда система небольшая (в модели всего несколько миллионов параметров), то даже небольшое внешнее возмущение заставляет её штормить. Короче, всё надо перепроверять.

Возвращение к истокам

Одного живого латента недостаточно, значит логичная идея — расширить это число до, скажем, 4. Таким образом станем чуть ближе к корректной реализации, которая была в оригинальной работе по Transformer-XL. Живых латентов держать много дорого: требование к памяти (VRAM) и времени растёт нелинейно, а тренировка нескольких десятков агентов в параллель — это много часов даже на нескольких GPU. Получаются какие-то вилы: добавить 4 живых латента к уже имеющимся замороженным – это сделать ещё медленней, а не добавлять нельзя, так как страдает обучение. Но раз так, то тогда проще вообще избавиться от замороженных латентов и посмотреть, что получится на исключительно живых. Бонусом мы получаем более чистую архитектуру без дополнительных сложностей.

Сказано — сделано: https://github.com/kocherovms/neurolab/blob/main/17_rl/17e_ppo_tr_atari_mp_03.ipynb. Делаем пару десятков прогонов. Ну уже получше, один из агентов нахоботил аж 779 очков (лучший запуск — 17e_ppo_tr_atari_mp_03/8). Это недалеко от 980 очков для Agent57 на 1 млрд. бюджета – есть первое положительное подкрепление! Хотя видно, что агенты выходят на плато и дальше особо ничему не учатся.

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

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

Больше очков != понял игру

Локальный минимум — не баг агента, а сигнал о том, что награда за «зайти в иглу» слишком редка, чтобы её нашёл случайный поиск.

Чтобы пройти первый уровень, агент каким-то образом должен перестать разбирать иглу. Трёх жизней достаточно, чтобы, даже просто тупо прыгая по льдинам вверх-вниз, достроить иглу — иглу переживает потерю жизни. А потом он должен каким-то образом вылезти на берег и потопать в сторону иглу, и зайти в дверь (надо нажать на кнопку UP, будучи в створе двери). Какова вероятность, что агент это сделает случайным образом? Очень близко к 0. Просто нет никаких нативных подкреплений, зачем ему так делать.

Выводы

Формула [трансформер + PPO] выглядит рабочей для Frostbite. Заработано 779 очков на 30 млн шагов, тогда как Atari57 на 1 млрд. зарабтал 980.

Агент успешно научается базовым механикам игры, но не способен постичь иерархическое целеполагание. Возможно, что это можно побороть грубой силой — просто выполнить пару сотен, а может и тысяч запусков. Глядишь, какой-нибудь счастливчик и появится, который случайным образом соберёт иглу, зайдёт в него и закрепит это знание. В конце концов Atari57 так и делает и у него получается.

Но лучше поискать более управляемое решение.

Будем искать. В следующем эпизоде.

Github: https://github.com/kocherovms/neurolab

Технические заметки: https://kocherovms.github.io/

Литература

1. Schulman J., Wolski F., Dhariwal P., Radford A., Klimov O. Proximal Policy Optimization Algorithms. — arXiv:1707.06347, 2017. arxiv.org/abs/1707.06347

2. Mnih V. et al. Human-level control through deep reinforcement learning (DQN). — Nature 518:529–533, 2015. nature.com/articles/nature14236

3. Huang S., Dossa R. F. J., Ye C., Braga J., Chakraborty D., Mehta K., Araújo J. G. M. CleanRL: High-quality Single-file Implementations of Deep Reinforcement Learning Algorithms. — JMLR 23(274):1–18, 2022. jmlr.org/papers/v23/21-1342.html

4. Dai Z., Yang Z., Yang Y., Carbonell J., Le Q., Salakhutdinov R. Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. — ACL 2019. arxiv.org/abs/1901.02860

5. Pleines M., Pallasch M., Zimmer F., Preuss M. Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agents. — JMLR 26, 2025. arxiv.org/abs/2309.17207

6. Vaswani A. et al. Attention Is All You Need. — NeurIPS 2017. arxiv.org/abs/1706.03762

7. Geva M., Schuster R., Berant J., Levy O. Transformer Feed-Forward Layers Are Key-Value Memories. — EMNLP 2021. arxiv.org/abs/2012.14913

8. Bellemare M. G., Naddaf Y., Veness J., Bowling M. The Arcade Learning Environment: An Evaluation Platform for General Agents. — JAIR 47:253–279, 2013. jair.org/index.php/jair/article/view/10819

9. Machado M. C., Bellemare M. G., Talvitie E., Veness J., Hausknecht M., Bowling M. Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents. — JAIR 61:523–562, 2018. arxiv.org/abs/1709.06009

10. Towers M. et al. Gymnasium: A Standard Interface for Reinforcement Learning Environments. — arXiv:2407.17032, 2024. arxiv.org/abs/2407.17032

11. Kanerva P. Hyperdimensional Computing: An Introduction to Computing in Distributed Representation with High-Dimensional Random Vectors. — Cognitive Computation 1(2):139–159, 2009. link.springer.com/article/10.1007/s12559-009-9009-8