Mahjax: GPU-ускоренный симулятор маджонга для обучения с подкреплением достиг 2 млн шагов в секунду
Исследователи из сообщества машинного обучения представили Mahjax — новый симулятор для игры в Риичи-маджонг, полностью реализованный на фреймворке JAX и оптимизированный для работы на графических процессорах. Разработка нацелена на ускорение исследований в области обучения с подкреплением (RL) в условиях неполной информации и высокой размерности состояний.
Риичи-маджонг — сложная многопользовательская игра со стохастическими элементами, которая сочетает в себе вызовы, характерные для реальных задач принятия решений. Ранее для обучения агентов часто использовалось обучение с учителем на основе записей игр людей, однако методы, способные обучаться с нуля (tabula rasa), как в семействе AlphaZero, обладают большей универсальностью.
Mahjax представляет собой полностью векторизованную среду, что позволяет параллельно проводить миллионы симуляций на GPU. По данным авторов, на восьми ускорителях NVIDIA A100 симулятор достигает пропускной способности до 2 миллионов шагов в секунду для версии без правил красной пятёрки (no-red) и до 1 миллиона шагов в секунду с правилами red.
В состав инструмента также входит качественный визуализатор для отладки и взаимодействия с обученными агентами. Это упрощает процесс разработки и анализа поведения моделей.
Экспериментальные результаты подтверждают, что среда пригодна для полноценного обучения с подкреплением: агенты, обученные в Mahjax, успешно улучшают свой ранг при игре против базовых стратегий. Разработчики отмечают, что высокая производительность симулятора позволит ускорить исследования в области multi-agent RL и сложных игр.
Публикация на arXiv (ID 2605.20577) привлекла внимание сообщества, так как подобные GPU-ускоренные среды становятся всё более востребованными для масштабных экспериментов. Ожидается, что Mahjax может стать стандартным инструментом для исследований в области маджонга и родственных задач.
