Давненько хотел написать бота, который играет в какую‑то сложную игру. Негамакс с альфа‑бета отсечением, конечно, работает неплохо, но на достаточно сложных играх тонет. В связи с этим, хотелось бы как‑то это ограничение обойти, что и привело меня к обучению с подкреплением. Стандартные Q‑Table, опять же, неприменимы для чего‑то сложнее крестиков‑ноликов, так что переходим к нейронным сетям. А точнее, к одному конкретному алгоритму: AlphaZero.

Немного о сути алгоритма:

В основе AlphaZero лежит алгоритм Monte Carlo Tree Search (далее MCTS). Если коротко, данный алгоритм можно описать как вероятностный минимакс. В отличие от обычного минимакса, который исследует всё дерево вширь до заданной глубины, MCTS исследует чётко ограниченное количество узлов, которые считает перспективными.

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

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

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

Подводные камни:

Естественно, всё не так просто на практике. Основная проблема, с которой мы здесь сталкиваемся — это переобучение, которое проявляется самым возмутительным образом. Если итоговая работа MCTS слишком сильно совпадает с предсказаниями нейронной сети, то мы попадаем в ловушку: MCTS не генерирует новых данных → сети не на чем учиться → качество игры стагнирует. И, к сожалению, это может случиться на любом этапе обучения, а не только в самом конце, когда сеть уже достигла высокого уровня игры.

В связи с этим становится позарез необходимо обеспечить разнообразные данные для нейронной сети, что достигается следующими образами:

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

   Естественно, у данного подхода есть проблема: ходы, которые по мнению MCTS не лучшие, зачастую реально не лучшие, и учить на них сеть как‑то сомнительно. Поэтому, чтобы качество игр оставалась высоким, стандартной практикой является выбор ходов с заданной вероятностью до какого‑то хода партии, а потом уже фиксировано выбирать только лучших ход. Это обеспечивает разнообразие ходов в начале и ювелирную точность в конце.

2. Сделать MCTS менее уверенным. Если сеть предсказала, что ход плохой, то MCTS туда никогда не зайдёт. Что полезно, когда ход реально плохой, но если он на самом деле хороший, то MCTS об этом никогда не узнает, сколько поисков ему не выдавай.

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

   Проблема в том, что добавление шума после каждого узла сильно режет эффективность самого MCTS, поскольку поиски тратятся на то, что нейронная сеть уже отсеяла, поэтому шум применяется не к каждому узлу, а только к корню — актуальной игровой позиции.

Данные костыли вероятностные, поэтому, чтобы они работали, необходимы большие числа. Соответственно перед тем, как обучать сеть, необходимо сгенерировать огромное количество симуляций, что приводит нас к следующему вопросу:

Чему учить сеть?

Оригинальный алгоритм разрабатывался Гуглом для го, в которой на тот момент компьютеры всё ещё не могли победить человека, а после успеха был обобщён на сёги и шахматы. В чём проблема? В том, что для генерации и обучения данных использовался кластер из 1000+ TPU. Ой. Как оказалось, алгоритм немного кусается по затратам, так что придётся взять за основу игру, которая несколько проще. И вот тут я узнал о существовании игры, для которой как раз по каким‑то причинам и существует огромное количество любительских реализаций AlphaZero — Connect 4 (или по‑русски 4-в‑ряд).

Данная игра очень похожа на крестики‑нолики, за исключением того, что для победы надо поставить не 3 фишки в ряд, а 4, игра ведётся на поле 6×7 и ставить фишку можно только в самую низкую незанятую клетку столбца. Всё бы хорошо, но игра была полностью решена в 1987 году и выложенный в открытом доступе солвер Паскаля Понса идеально её решает (да ещё и показывает, как быстро каждый ход приведёт к поражению) за доли секунды. Упс. Слишком просто, чтобы тратить время? Как бы да, но как бы нет.

1. То, что игра такая простая гарантирует, что моих скромных ресурсов в виде GTX 1650 на неё хватит.

2. Солвер Понса идеально работает за счёт книги опенингов. Без неё он на просчёт первого хода тратит огромное количество времени. Так что, технически, идеальная сеть может работать быстрее, чем алгоритм без книги, к тому же занимая меньше места, чем сама книга.

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

   Кроме того, солвер Понса даёт роскошь, которую нет в шахматах и других нерешённых играх: объективная оценка силы модели. Нам не нужно гадать, улучшается модель или нет, просто проанализировать, как она играет через мгновенный солвер, и результат налицо.

Таким образом, моя цель здесь НЕ просто изобрести велосипед, к тому же кривой и медленно работающий (увы, обогнать идеальный солвер вероятностным подходом не получится никак). Моя цель: взять этот велосипед и заставить его ездить максимально быстро на GTX 1650. Ведь, как оказалось в дальнейшем, генерировать огромное количество игр задача не самая тривиальная.

Почему я ушёл с Python.

Итак, для начала я взял следующую реализацию алгоритма (а точнее его более старой версии, разработанной конкретно для го и отличающийся тем, что в ней модель для генераций игр обновляется только после того, как она победила предыдущую. В конечной итерации от этого отказались, и модель обновляется сразу) на Python из книги Deep Reinforcement Learning HandsOn: https://github.com/PacktPublishing/Deep‑Reinforcement‑Learning‑Hands‑On‑Third‑Edition/tree/main/Chapter20.

Сделав парочку исправлений в файле с самой сетью, её можно даже запустить. После этого мы сталкиваемся с очевидной проблемой: одна единственная игра генерируется несколько секунд, при переходе на GPU скорость у меня даже упала в два раза, а качество игры при оригинальных параметрах быстро выходит на плато.

При попытке начать генерировать больше игр за итерацию, время на генерацию становится недопустимо долгим. В итоге мы не можем ни тренировать нормально, ни даже использовать GPU. Отсюда возникает простой вывод: алгоритм нужно ускорить, и самое очевидное решение получить бешеную скорость — переписать на плюсы.

Действительно, основное время занимает не обработка нейронной сетью игровой позиции, а построение дерева, что на Python происходит непростительно медленно. Сказано, сделано, и после простого переноса кода один в один (насколько это возможно), без использования дополнительных оптимизаций, скорость генерации внезапно взлетела, а использование GPU наконец‑то дало прирост в производительности вместо того, чтобы её ухудшать.

Чтобы не быть голословным, вот конкретные цифры на Python cpu (я в тот момент экспериментировал с большим количеством поисков (1280), потому что моя модель никак не хотела учиться играть, так что время на генерацию несколько медленное):

steps  17, leaves 3398, steps/s  1.66, leaves/s 331.46, best_idx 0, replay 18
steps   6, leaves 1324, steps/s  1.69, leaves/s 373.82, best_idx 0, replay 25
steps   7, leaves 1702, steps/s  1.70, leaves/s 413.27, best_idx 0, replay 33
steps  36, leaves 3708, steps/s  1.71, leaves/s 176.49, best_idx 0, replay 70
steps   8, leaves 1661, steps/s  1.63, leaves/s 338.48, best_idx 0, replay 79
steps   6, leaves 1160, steps/s  1.59, leaves/s 307.38, best_idx 0, replay 86

А вот на с++ cpu:

steps  41, leaves 2753, steps/s  9.42, leaves/s 632.29, best_idx 0, replay 42
steps  14, leaves 1580, steps/s  7.39, leaves/s 833.77, best_idx 0, replay 57
steps  14, leaves 1450, steps/s  7.92, leaves/s 820.60, best_idx 0, replay 72
steps  21, leaves 1868, steps/s  8.14, leaves/s 723.75, best_idx 0, replay 94
steps   7, leaves  648, steps/s  7.74, leaves/s 716.81, best_idx 0, replay 102
steps   6, leaves  753, steps/s  6.85, leaves/s 859.59, best_idx 0, replay 109

И на c++ gpu:

steps   8, leaves  656, steps/s 10.05, leaves/s 824.12, best_idx 0, replay 9
steps   6, leaves  476, steps/s 12.02, leaves/s 953.91, best_idx 0, replay 16
steps  18, leaves 2184, steps/s 11.31, leaves/s 1372.72, best_idx 0, replay 35
steps  19, leaves 1877, steps/s 11.76, leaves/s 1161.51, best_idx 0, replay 55
steps  15, leaves 1706, steps/s 11.94, leaves/s 1358.28, best_idx 0, replay 71
steps  25, leaves 2140, steps/s 11.84, leaves/s 1013.26, best_idx 0, replay 97

Нехилая такая разница, раз в 5 точно. И это только начало!

Борьба за параллельную генерацию игр.

Если поглядеть на оригинальный код, то можно заметить очевидную вещь. При генерации нескольких игр за раз они генерируются ПОСЛЕДОВАТЕЛЬНО, хотя друг от друга не зависят вообще никак.

for _ in range(PLAY_EPISODES):
  _, steps = model.play_game(mcts_store, replay_buffer, 
                             best_net.target_model, best_net.target_model, 
                             steps_before_tau_0=STEPS_BEFORE_TAU_0, 
                             mcts_searches=MCTS_SEARCHES, 
                             mcts_batch_size=MCTS_BATCH_SIZE, device=device)
  game_steps += steps

Решение простое: нужно подключить многопоточность. Для python многопоточность для ускорения CPU‑bound вычислений работает плохо из‑за GIL, что объясняет, почему параллельной генерации не было в оригинальном коде, но мы уже отошли от него в пользу плюсов, чтобы ограничения Python нам не мешали. И после выделения пула потоков с помощью для генерации игр я получил максимальное ускорение в целых 2 раза. При 12 логических ядрах процессора. Что‑то очевидно пошло не так. И, что пошло не так конкретно, очевидно по данному фрагменту кода.

batch_v = model.state_lists_to_batch(expand_states, expand_players, device)
logits_v, values_v = net(batch_v)

Несмотря на то, что код предусматривает одну критическую оптимизацию: группирует запросы к нейронной сети в рамках одного потока в батча, и отправляет их сразу батчем, потоков у нас внезапно стало несколько, и теперь каждый из них отправляет свой маленький батч. Раз уж мы уже поигрались с многопоточностью, то можно пойти дальше — выделить отдельный поток, который через promise/future будет получать запросы к нейронной сети со всех параллельных игр, и одним большим батчем их отправлять. И мы получаем первое значительное ускорение.

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

Почему? Потому что GPU не обращается к оперативной памяти напрямую, а вынужден стучаться к процессору, который уже берёт данные из RAM и отдаёт их GPU, из‑за чего процессор вместо построения дерева занимается огромную часть времени прогоном данных туда‑сюда. А почему она не может обращаться напрямую? Потому что физическая память процесса в RAM благодаря ОС живёт своей жизнью и постоянно меняет своё расположение. Может даже на файл подкачки улететь, что вообще производительность убивает.

Решение? Сказать ОС, что данную область памяти никуда двигать категорически нельзя. Что достигается с помощью pinned memory. Нам заранее известно количество потоков и количество поисков, объединённых в батч. Соответственно, мы можем заранее выделить память нужного размера, где будут хранится наши тензоры.

Да, за всё приходится платить, и выделение pinned memory стоит дороже, чем обычной, однако, нам заранее известно количество потоков и поисков в одном батче, а значит, размер памяти, который нужен одномоментно, так что динамическая аллокация нам не нужна, мы будем выделять pinned memory всего лишь один раз в начале работы программы, а потом просто записывать туда же новые тензоры. В результате мы получаем ещё один прирост в производительности. Победа?

Всё ещё нет. Прежде всего, хочу обратить ваше внимание на этот код и найти проблему.

gpu_input_buffer_ = torch::empty(
  { batch_size_, 4, GAME_ROWS, GAME_COLS },
  torch::TensorOptions().dtype(torch::kFloat32).device(device_)
);
auto input_view = gpu_input_buffer_.narrow(0, 0, static_cast(batch.size())); 

А проблема заключается в том, что мы отправляем на GPU данные динамического размера! А чем это грозит? Промахами кеша и сегментацией памяти. Я эту проблему обнаружил даже не сразу, а когда протестировал сеть размером 10 остаточных блоков по 128 фильтров. До этого момента видеокарта как‑то справлялась, а с по‑настоящему большой сетью сегментация памяти привела к тому, что на 28 итерации тренировки скорость обработки данных резко упала. Решение? Округлить размер пересылаемых данных до ближайшей степени двойки.

Ну а теперь переходим к главной второй оптимизации. В чём проблема текущего выделения потоков? В том, что они все активно используют кучу! Казалось бы, в чём проблема, ну используют и используют. А проблема в принципе работе malloc. Постоянно запрашивать новые данные в кучу у ОС для каждой переменной дорого, поэтому malloc поступает хитрее: он запрашивает сразу большую страницу, сам отслеживает, насколько она заполнена, и запрашивает новую только тогда, когда в этой больше данных не помещается.

Это всё хорошо работает, когда у нас один поток, а когда у нас их несколько? Правильно, malloc обрабатывает один в очереди и ставит mutex, а остальные в это время ждут. Производительность внезапно проседает Решение? Сделать собственный аллокатор памяти для каждого потока, для чего у нас есть std::pmr. В результате мы выжимаем последние крохи производительность и ощущаем полную выгоду от всех 12 потоков.

График средней скорости генерации игровых ходов во время генерации 512 игр с 32 поисками и 3 residual blocks по 64 фильтра в зависимости от количества одновременных игр представлен ниже:

Зависимость, конечно, далеко не линейная, но 6 раз разницы тоже неплохо: 150 позиций в секунду на одном ядре против 950 на 12. И это не говоря о Python, на котором скорость была вообще 65 позиций в секунду.

Ну и, естественно, куда же без результата профилировщиков. Начинаем с NSight Systems

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

Nsight Compute жалуется, что нагрузка на SMSP и SM неравномерная, и говорит, что ряд ядер по перемножению матриц можно ускорить значительно в теории, но что поделать, входные данные у нас не идеально ровные. Так что 20% потенциальной выгоды вообще попытки их исправить не стоят, накладные расходы запросто могут выше оказаться.

Как я пытался подобрать гиперпараметры

Итак, основная проблема — генерация данных, решена, теперь осталось собственно обучить сеть. Всё просто? Несовсем. Гиперпараметров много, начиная от того, сколько игр за раз генерировать и насколько глубокий делать поиск, и заканчивая тем, какая мощность сети и какой learning rate. Плюс ещё в попытке обучения периодически обновлял саму программную реализацию, исправляя баги или просто пробуя поменять какую‑то мелочь, так что чёткой хроники что и как я менял у меня не сохранилось. Ну ладно, пройдёмся, что я решил с основными параметрами.

Оригинальный алгоритм использовал SGD, от которого я в итоге отказался в пользу AdamW, потому что меня задолбало пытаться подбирать расписание learning rate вручную, и решил переложить это на сам алгоритм. После этого обучение вроде стало более гладким, и я решил оставить, но по‑моему, попытки контролировать его в Connect4 вообще мало на что влияют. Единственное, что я обнаружил, что SGD лучше играет при learning rate 0.1 или 0.01, а AdamW это слишком много, ему это оставил стандартный 0.001.

Идём дальше, нейронная сеть. После того, как моя нейронка продолжала упираться в достаточно высокое, но плато в качестве игры, мне надоело пытаться угадать, чего ей не хватает, поисков или ёмкости, и решил проверить это экспериментально: использовать солвер для генерации идеальных данных для обучения и смотреть, какая нейронка их сможет зазубрить.

План выглядел надёжным, однако, почему‑то все мои архитектуры, которые я до этого тестировал в alphazero, и где они показывали неплохой результат, играли при обучении на идеальных данных заметно хуже. Что пошло не так? Казалось бы, предсказываем единицы для лучших ходов и 0 для всех остальных. Это идеал. Нейросеть не будет тратить поиски на неперспективные ходы в принципе и результат что с MCTS, что без него будет идентичным. А вот и нет!

Если мы будем использовать такой one hot encoding, то для нейросети не будет разницы между долгим выигрышем, немедленным проигрышем, и невозможным ходом. И как это прикажите обобщать? Неудивительно, что нейронка стала играть хуже после обучения с учителем, от такого всё или ничего градиент просто взрывается и нейронка пытается зазубрить конкретные ситуации, а не сделать обобщающие выводы.

Поэтому в идеале вероятности надо ранжировать, имитируя поведение реального неуверенного MCTS. Солвер Понса возвращает большее значение для лучшего хода, так что просто прибавляем минимальное число, чтобы не было отрицательных, складываем и делим каждое значение на полученную сумму.

Всё бы хорошо, но нейронка всё равно играет хуже. Да что ж такое. Даёшь ей идеальные таргеты, а она на неидеальных учится лучше. Алгоритм, конечно не оптимальный для обучения с учителем, вероятности мы предсказуем, потому что конвертировать посещения в значение мата в n ходов задача нереальная, в отличие от обратной, что мы только что попробовали, но с сомнительным результатом. В общем, искусственные значения для нейронки почему‑то не заходят, так что берём 3 остаточных блока по 64 фильтра (я пробовал 32, но с ними играет чуть хуже и разница в скорости генерации минимальная, так что 64 надёжнее).

ОК, с размером сети определились, количество игр и поисков я определял чисто эмпирически, чисто смотрел, падает ли loss и играет ли модель лучше. Размер батча для обучения, буфера позиций и количества эпох определялся из соображения, чтобы модель увидела каждую позицию хотя бы раз, прежде чем та пропадёт (имейте в виду, что за счёт того, что игра симметричная, количество позиций в буфере идёт в два раза больше, за счёт того, что мы каждую сгенерированную позицию зеркалим. Больше данных — лучше обучение). Остановился на числе 512 игр (1024 после отзеркаливания), 512 размер батча для тренировки, 160 000 позиций в буфере, 64 поиска.

И вот после всего этого, моя модель достигает потрясающего результата всего за 30 минут, но не хочет играть лучше. Я попробовал улучшить, скрестив MCTS с минимаксом, помечая решённые ветки и не тратя поиски на них, но это дуристика, суть алгоритма в том, что у него мало поисков, и решённых веток у нас почти не будет. Поэтому, я лишь сделал, чтобы алгоритм сразу отметал зевки в один ход вместе с невозможными ходами (худший возможный ход, нулевая вероятность, конфликтов ни с чем нет).

Для тестирования итоговой силы игры я также добавил возможность MCTS вызывать неглубокий солвер при оценке позиции, чтобы тот отметал зевки на небольшой глубине, повышая тем самым точность. Правда, ненамного, MCTS и так хорошо работает, неглубоких зевков мало, а если ставить слишком большую глубину, то он уже чисто за счёт подсказок солвера играет, а не нейронки, что нарушает всю идею.

Результаты

В общем, итоговый график точности представлен после тренировки длинной полчаса представлен внизу (помощь неглубокого минимакса НЕ использовалась):

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

Как видно, разница буквально в пределах процента. При этом за полчаса было сгенерировано 49 итераций по 1024(с учётом отражений) игр, или 50 176 игр всего, что даёт среднюю скорость 1672 игры в минуту. В общем, результатом я в целом доволен. Не 100%, конечно, но учитывая, что обучение заняло полчаса, 85% идеальных ходов и меньше 5% явных ошибок — более чем солидный результат. Больше времени и поисков, правда, дают парочку процентов максимум, зато весь вау эффект от быстроты убивают.

Увы, поскольку идеальный солвер работает мгновенно, хочется идеального результата и от нейронки, но что поделать, алгоритм разрабатывался для игр, где идеальное дерево не построишь, а тут это как гвозди микроскопом забивать. В общем, практической ценности 0, но выжать максимум из железа компьютера для Deep Reinforcement Learning было довольно интересной задачей.

Исходный код проекта находится тут, если кто хочет потыкать: https://github.com/Leo88p/cpp‑connect4-alpha‑zero