Добрый день, уважаемыее!
Я хочу поделиться с вами очень интересным проектом, над которым работал в последнее время.
В первой статье я не буду сильно углубляться в технические подробности, а вместо этого постараюсь провести вас по пути, который я прошел при реализации своего пайплайна для обучения нейросеток, сражающихся друг с другом на арене. Весь код доступен на моем GitHub и готов к использованию, поэтому вы сразу сможете обучить чемпиона и поучаствовать в сражении!
Готовы? Тогда - вперед!
Вступление: зачем и почему
Меня всегда завораживали нейросетевые агенты, действующие в некоторой среде для достижения своих целей. Завораживали они прежде всего своей внутренней сложностью, степенями свободы своего искусственного разума.
Судите сами - внутреннее состояние большинства игровых объектов в любой игре, в том числе и в файтингах - горстка переменных. Какие бы текстуры на них ни натягивали, в какой бы нарратив ни вписывали, внутри - это все те же грустные 5-10-20 переменных, определяющие очень простое поведение. Файтинги мне вообще с детства казались странными - можно миллион раз повторить комбинацию нажатий кнопок и боец миллион раз выполнит одно и то же действие. У меня никогда не было ощущения, что персонаж действительно находится там, на арене. Он слеп, он глух, он - механическая рука из автомата-хваталки, повторяющая движения джойстика. Он - анимация, проигрываемая при нажатии комбинации клавиш.
Когда я впервые увидел нейросетевых агентов, управляющих своим телом в физической симуляции, я сразу прочувствовал эту разницу. Здесь нет предзаписанных анимаций, но есть проприоцепция агента (его восприятие своего тела). Здесь нет комбинаций клавиш, но есть его мотивация в виде функции наград, которую он стремится оптимизировать. Наконец, даже если два агента не отличаются 3д-моделькой и текстурой, поведение каждого из них определяется несколькими сотнями тысяч весовых коэффициентов. Каждый агент уникален, предсказать, как он себя поведет можно только приблизительно, как будто это - такая удивительная волшебная зверюшка, которую можно призвать в наш мир силами машинного обучения.
С этого началась моя работа над проектом под кодовым названием "Мужики из MuJoCo" (The MuJoCo Men). С одним таким товарищем, бодро семенящем в светлое будущее, вы уже познакомились в заголовке статьи - это достаточно известная модель Humanoid в физическом симуляторе MuJoCo, который некоторое время назад был выкуплен DeepMind'ом и переведен в open source.
По ссылке вы можете ознакомиться с документацией на модель и среду, в которой она существует. Как вы можете увидеть, у нее есть набор действий (actions), представляющих собой 17 float-переменных, управляющих моментами на моторах в его конечностях. Каждый тик дискретного времени симулятор просчитывает физику - гравитацию, столкновения объектов - поэтому минимальный работающий код не нуждается ни в каких нейросетях, можно подать в качестве управляющего сигнала вектор из 17 случайных float'ов и смотреть, как Однако если мы в самом деле хотим зажечь эту искорку искусственной жизни, нам придется обратиться к обучению с подкреплением - подразделу машинного обучения, в котором нейросетевой агент формирует этот самый управляющий сигнал (вектор из 17 флоатов в нашем случае), стремясь максимизировать получаемую от среды награду (reward). Сам MuJoCo - это просто физический симулятор, никаких сигналов награды он не предоставляет, поэтому для реализации пайплайна обучения нужно обернуть код, пинающий расчет физики, в код, который бы получал от MuJoCo состояние системы (например, координаты агента и его составных частей), рассчитывал на основании этого некоторый сигнал, соответствующий нашим представлением о правильном выполнении задач, и отдавал его пайплайну обучения вместе с наблюдениями (observations). В упомянутой среде от OpenAI в качестве награды используется значение, расчитанное на основании скорости центра масс агента (чтобы побудить его бежать вперед), активации его моторов (чтобы побудить его экономить силы и не размахивать руками-ногами почём зря), его вертикальности (чтоб не падал) и т.п. Наблюдения формируются из собранных в один вектор состояний, полученных от симулятора (агенту не обязательно видеть все состояние физической системы целиком!)Про стандартную модель Humanoid
его корёжит Humanoid совершает случайные движения всеми своими сочленениями (joint'ами).
Немного покопавшись в интернете, я обнаружил, что один из наиболее эффективных алгоритмов для обучения этого хуманоида бегу - алгоритм Soft Actor-Critic (SAC) все от того же OpenAI. В своей реализации я ориентировался на реализацию SAC от CleanRL - очень простой, понятный код без излишеств. Но почти непригодный для реального применения из-за того, что все считается последовательно, в одном потоке.
Тогда же я познакомился с невероятно могучим инструментом под названием JAX.
JAX, великий и ужасный
Если очень коротко - это система JIT-компиляции матричных преобразований, позволяющая писать код на питоне (с некоторыми ограничениями), который при запуске будет автоматически скомпиллирован в высокоэффективный код под GPU, TPU либо CPU. Представьте себе работу с матрицами в numpy, но с возможностью автоматически вычислять градиенты функций, одной строчкой добавлять работу с батчами, использовать современные архитектуры нейронных сетей и оптимизаторов - и все это сходу работающее на любых высокопроизводительных системах, от современных видеокарт до TPU.
Чтобы лучше вникнуть в особенности работы SAC, а главное - набить руку в использовании нового для меня JAX, я решил реализовать пайплайн для обучения c нуля, с дальнейшими планами на реализацию, собственно, арены, на которой эти агенты смогут сражаться. Кроме того, меня очень интересовала возможность реализации непрямого управления, которая бы позволила, наконец, отойти от этих фиксированных комбинаций клавиш и перейти к динамическому заданию мотивации агента в духе "держись ближе к центру арены", "бейся агрессивнее", "будь осторожнее".
Не все из этого удалось реализовать на нужном мне уровне, пришлось скорректировать планы, но об этом я расскажу в заключении. А сейчас - перейдем к практике!
Устанавливаем и запускаем пайплайн
Так как пайплайн я писал с нуля, у него довольно мало зависимостей. Главное, что вам нужно будет установить - JAX с поддержкой CUDA и MuJoCo с его JAX'овой версией (да-да, его портировали на тот же джакс, благодаря чему теперь можно на видеокарте считать тысячи инстанций в параллель!)
Полный набор зависимостей моего пайплайна можно найти в файле requirements.txt.
requirements.txt
jax[cuda12]>=0.4.30 Сам JAX
matplotlib>=3.5.3 Библиотека для рисования графиков
mujoco>=3.1.6 Симулятор MuJoCo с биндингами к питону
mujoco-mjx>=3.1.6 Его JAX-версия
numpy>=1.24.4 Старый добрый Numpy для работы с матрицами на CPU
optax>=0.2.3 Оптимизаторы от JAX
flax>=0.8.5 Библиотека для сборки нейросетей на JAX
wandb>=0.18.5 Логирование результатов экспериментов в Weights & Biases
tensorboard==2.11.2 Альтернативой логированию в Wandb может быть Tensorboard. Можно поставить что-то одно, пайплайн позволяет выбрать, куда рапортовать.
Предлагаю вам обратиться к гайду по установке JAX - если у вас машина на linux, то, скорее всего, все сведется к выполнению строки pip install -U "jax[cuda12]"
Под Windows поддержки GPU, к сожалению, нет - вы сможете запустить инфиренс, посмотреть, как дерутся агенты в реалтайме, но для обучения придется плясать с бубном и пытаться запустить под WSL - я сам этого не пробовал, поэтому не могу сказать, насколько это тяжело.
Второй критически важный элемент - сам симулятор MuJoCo с биндингами к питону и его JAX-версия - MuJoCo MJX.
Скорее всего, достаточно будет сказать pip install mujoco mujoco-mjx
Если эти два компонента поставились без проблем и увидели вашу CUDA, то дальше все будет элементарно. Если нет... То вашей целью будет поднять их так, чтоб они начали выполнять код на вашей видеокарте, большинство проблем разобрано в гайдах этих двух библиотек.
Третий компонент, опциональный, но очень полезный - трекинг экспериментов Weights & Biases
Вы можете стартовать пайплайн без логирования в wandb, но отслеживать прогресс без него очень неудобно. Для подключения wandb нужно бесплатно зарегистрироваться на их сайте, получить ваш API-ключ, после чего (не забыв предварительно поставить библиотеку через pip install wandb) сказать в консоли wandb login и предоставить ключ API. В качестве альтернативы можно поставить Tensorboard, но мне, честно говоря, он нравится намного меньше. В настройках пайплайна (о них ниже) можно будет выбрать, куда логировать данные - только в консоль, в wandb, в Tensorboard.
Теперь осталось выполнить
git clone https://github.com/r-aristov/arena.git
cd arena
python sac_my_flax.py
Если все сделано правильно, скрипт начнет медленно пробуждаться, JIT-компилируя все необходимое (довольно медленный процесс). Все настройки - внутри стартового скрипта, поэтому пока он стартует, я бегло опишу, что же собой представляет пайплайн.
Арена
Первое, от чего пришлось отказаться - это, собственно, Мужики Из MuJoCo. Я довольно быстро обучил хуманоида бегать, и даже отвечать на мой управляющий сигнал.
Но как оказалось, в стандартной среде отключен рассчет коллизий для всего, кроме самых "стоп" гуманоида. В случае с бегом этого вполне достаточно. Но для битвы нужно двое, причем со включенными коллизиями. Даже один гуманоид при включенных коллизиях настолько медленно считался на моей старенькой GTX 1080, что пришлось корректировать планы. В итоге я остановился на более простой модельке - MuJoCo Ant. Разумеется, среду я полностью реализовал свою, позаимстовав только физическую модель этого четырехногого вошика.
Таким образом, среда представляет собой арену размерами 2.5x2.5x1.5, на которой сражаются два четырехногих агента, пытаясь скинуть друг-друга в пропасть.
Полностью включенные коллизии делают даже этих вошиков тяжелыми для рассчета, поэтому я остановился на варианте "включены коллизии между каждым из агентов и ареной, а также между торсами агентов" - их ноги проходят друг через друга, но все же этого достаточно, чтобы они могли активно друг-друга пинать.
Физическая модель среды в XML-формате, требуемом MuJoCo лежит в папке models/arena.xml, там же лежит вариант с полными коллизиями - models/arena_all_collisions.xml. Это то, из чего физически состоят агенты - набор примитивов, соединенных управляемыми сочленениями. У MuJoCo есть огромный документ про этот их язык описания. Не могу сказать, что мне доставило большое удовольствие копаться в нем, но на пару с ИИ мы как-то разобрались и собрали арену с двумя четырехногими вшами и ограниченными коллизиями. Вам вряд ли придется что-то менять в этих файлах, но если вдруг захотите посмотреть, как ведут себя агенты с потяжелевшими задницами, можно добавить density="25.0" (или что-то побольше) в параметры торсов агентов на строке 23 для первого агента и на строке 77 для второго.
Но если вы хотите поэкспериментировать с мотивацией вашего чемпиона (а я очень хотел бы увидеть ваших чемпионов на арене!) - вам нужно отправиться в глубины файла arena.py, который содержит код среды, а также код для теста и визуализации обученных агентов.
Там, на строке 121 вы найдете функцию compute_reward, которая получает на вход полный прошлый и текущий стейты среды, сигнал "был ли контакт между агентами" в прошлом и текущем тике, а также - насколько приросло максимальное удаление агента0 и агента1 от центра арены.
Ее код - результат моих проб и ошибок, порождающий сигнал награды, с которым получился вполне сносный боец. Его ревард состоит из награды за сближение (результат расчета movement_projection(0,1) - проекций векторов смещения агента в направлении противника), награды за пинок (dd_reward(0,1), то самое приращение расстояния от центра, при условии, что был зафиксирован контакт), а также награды/штрафа за падение с арены.
Оказалось очень трудно найти баланс между безбашенным берсерском (который отлично, агрессивно пинает врага, но улетает вместе с ним) и ссыкуном очень осторожным агентом, который не рискует лишний раз приближаться к краю, чтобы не упасть.
Текущий чемпион более-менее сбалансирован, но все же слегка берсерк, довольно часто улетает с арены вместе с противником. Может быть, у вас получится лучше?
Функцию validation_reward трогать не нужно - она выражает самую простую и понятную награду - агент получает -1 каждый тик, что он провел ниже плоскости пола арены, если его противник при этом не упал. И -1, если все наоборот. Если же оба ниже пола, никакой награды никто не получает. Эта функция нужна чтобы отслеживать объективный прогресс - именно по ней я отбирал чемпионов, именно ей будут оцениваться ваши агенты, если вы решите поучаствовать в сражении. Но если ее выбрать в качестве основной функции наград, агенты учатся намного хуже.
Чтобы посмотреть на битву обученных агентов, просто запустите arena.py - при старте будет загружен агент agent.flax из корня проекта как агент0 (оранжевый) и один из моих старых чемпионов в качестве его противника (синий). Ну или можете передать в качестве параметра пути к одному или обоим загружаемым агентам: python arena.py --agent0="path/to/your/agent0" --agent1 "path/to/your/agent1"
Стартовый скрипт и запуск обучения
Вот мы и подобрались к самому главному - стартовому скрипту sac_my_flax.py, содержащему все параметры текущей сессии обучения.
Прежде всего обратите внимане на dict с конфигурацией config на 38й строке.
config = dict(
seed = 42,
worker_device=worker_device,
trainer_device=trainer_device,
ref_agents_count=len(ref_agents_paths),
trainer_batch_size=4096*2,
worker_batch_size=1024*4,
buffer_size=4*4.096e6,
random_steps_count=2000,
last_agents_count=16,
self_play_lag = 2501,
initial_alpha = 1.0,
autotune_alpha = True,
tau = 0.005,
gamma = 0.995,
q_lr = 0.001,
p_lr = 0.001,
total_steps=2_000_000,
warmup_steps=40_000,
report_to_tensorboard=False,
report_to_wandb=True
)Пойдем по порядку:
seed - инициализирующее значение для всех генераторов случайных чисел в системе, для воспроизводимости экспериментов обычно фиксируется.
worker_device, trainer_device - устройства, на которых будет запущено обучение, определится автоматически, если JAX встал нормально, это будет cuda:0. Если у вас целых две видеокарты, то, во-первых, поздравляю, а во-вторых - пайплайн автоматически разбросает свои компоненты по разным GPU.
ref_agents_count - тоже определяется автоматически. Чтобы ваше обучение быстрее принесло плоды, в обучающие битвы будут добавлены несколько моих референсных агентов, хранящихся в legacy-agents/, это ощутимо улучшает результат. Если хотите учить с нуля, замените выше строку 36 ref_agents_paths = validation_agent_paths*4 на ref_agents_paths = []
А вот дальше - всё то, что можно крутить в поисках лучшей конфигурации.
Размер батча тренера trainer_batch_size и воркера worker_batch_size - я не стал в этой статье вдаваться в подробности архитектуры, но в ней два основных компонента - тред worker, который проводит симуляции и собирает данные в буфер и тред trainer, который берет данные из буфера и учит агента и его q-функцию. У меня уменьшение размера worker_batch_size ощутимо било по эффективности агента. И оно же больше всего влияло на производительность worker'а - симуляция физики - штука затратная.
buffer_size - размер буфера с наблюдениями, действиями агентов и наградами (replay buffer), располагается в RAM. Подбирал под свой размер RAM, 32GB, так что если не влезает - делайте меньше.
random_steps_count - Количество шагов случайным агентом перед началом обучения. Нужно для того, чтобы обучение не начиналось с одинаковых действий необученного агента. При использовании референсных агентов, наверное, не настолько актуально.
last_agents_count - как и положено в историях о становлении воина, самый сложный бой - с самим собой. Наш агент сражается с last_agents_count своих копий, взятых с задержкой self_play_lag итераций каждая. С этими значениями экспериментировал мало, можете смело пробовать менять.
Далее идет блок метапараметров тренировки, чтобы лучше понимать, за что они отвечают, рекомендую ознакомиться с описанием алгоритма SAC у OpenAI, посмотреть по коду или подождать следующих статей.
Если кратко: initial_alpha - коэффициент штрафа за низкую энтропию выдаваемых агентом действий, побуждает больше исследовать пространство состояний. 1.0 - это очень много, но далее идет опция autotune_alpha, которая включает оптимизацию этого значения. Я учил своих агентов в автоматическом режиме, с 1.0 в качестве начального значения. Можно экспериментировать, особенно, когда в обучении участвуют уже обученные референсные агенты и обучающемуся не грозит заскучать.
tau - с такой скоростью в целевые q-функции дистиллируется информация из текущих q-функций. Экспериментировал, но мало. В основном учил с этим значением.
gamma - старый-добрый discount factor из обучения с подкреплением - насколько далеко во времени тянется влияние прошлых наград. При gamma=0.0 агент учитывает только награды, которые получил непосредственно в текущий момент.
q_lr, p_lr - начальные значения для learning rate при оптимизации q-функции и самого агента (policy). На текущий момент я использую cosine learning rate scheduler, а именно - нарастание LR до начального значения в течение первых warmup_steps шагов обучения, а дальше - медленное снижение к нуля по косинусу.
total_steps - сколько всего итераций должен выполнить trainer. Сейчас стоит 2 млн, что на моей новой RTX 4070 занимает около 30 часов. На старой 1080 занимало примерно 5 суток. RL - штука небыстрая. Особенно с физикой. warmup_steps - столько шагов в начале обучения ваши q_lr и p_lr будут нарастать с 0.0 до заданных вами значений.
report_to_tensorboard, report_to_wandb - куда рапортовать графики ревардов, лоссов, валидаций. Можно выбрать что-то одно, я предпочитаю wandb. Можно даже отключить и то и то, тогда лог будет выводиться только в консоль, но этот вариант не рекомендую, очень трудно будет понять динамику обучения.
После запуска скрипта вы увидите примерно следующее:
Если включен вывод в wandb, то скрипт инициализирует новую сессию в вашем воркспейсе под названием sac-my-flax-arena, после чего выполнит 2000 шагов случайным агентом, а потом начнется, собственно, обучение. Основной график, за которым надо следить - validator-last-reward, показывающий ту самую валидационную награду, которую агент получает в сражении против набора референсных "старейшин".
График validator-best-reward - почти то же самое, только отображает не последнюю, а лучшую валидационную награду. Когда агент бьет свой рекорд, пайплайн сохраняет его самого (agent.flax), его q-функции (q1.flax, q2.flax) в папку checkpoints/имя_агента/номер_итерации. Там же, в папке last лежит последний сохраненный агент, q-функции и состояние оптимизаторов (сохраняется каждые 250 итераций).
Заключение
Статья получилась довольно объемная, но не охватывающая и десятой доли всего, что я хотел бы рассказать на тему этого проекта. Если дочитали досюда - поздравляю, вы справились! Обучайте вашего бойца и присылайте мне - можете на Хабре, можете - в комментариях на моем небольшом телеграм-канале, я там отвечаю всем. Если наберется хотя бы пара участников - сделаем соревнования, отрендерим красивое видео битвы.
Дальше проект можно развивать в огромном количестве направлений - вариантов экспериментов с ревард-функциями и вовсе бесконечно много. Я лично хотел бы все-таки подступиться к задаче непрямого управления - сделать у сетки дополнительный вход, на который идут сигналы, задающие "мотивацию" агента в терминах его ревард-функции. Я уже начал экспериментировать, но задачка трудная, особенно с учетом того, что один эксперимент занимает 30 часов.
Интересно было бы попробовать выделить какой-то необходимый минимум в виде предобученных фрагментов сетки (может быть, автоэнкодер?), чтобы снизить время на эксперимент и демократизировать доступ к этому замечательному развлечению - пока это все-таки забава для программистов. Если найдете более эффективные архитектуры или набор мета-параметров, снижающих время на обучение или сильно улучшающих результат - обязательно пишите!
Что я обязательно попробую - так это обучить агентов со включенными коллизиями. Это замедлит обучение раза в три, но возможностей у них будет намного больше (не только головами бодаться).
Также было бы неплохо, конечно, все же вернуть Мужиков из MuJoCo и посмотреть, как сражаются гуманоиды, но я уже представляю, сколько мне придется ждать, пока они обучатся, так что это пока в отдалённой перспективе.
Ну и напоследок - опрос!
Автор: Ariman
