эйай ньюз

XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX

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

Но на CPU достичь больших скоростей трудно - даже несмотря на то что у современных процов часто сотни ядер, они просто не могут угнаться за современными видяхами, мощность которых часто измеряется в петафлопсах. И вот поэтому среды стали переносить на GPU, но писать CUDA код, особенно быстрый не каждый ресёрчер сможет

Поэтому не так давно пошла мода писать среды на Jax - фреймворке от Google, основном конкуренте PyTorch. Создавали его в том числе под DeepMind, поэтому тамошние ресёрчеры убедились чтобы в нём было достаточно гибкости для создания таких сред. Но хоть Jax и в опенсорсе, а свои среды Google никому не даёт.

XLand-MiniGrid, созданный чуваками из T-Bank AI Research и AIRI, как раз и есть опенсорс репродукция Xland, закрытой среды от Google. Это grid-world среда, где агент перемещается по сетке и взаимодействует с объектами по определённым правилам. Такие среды можно очень просто и быстро симулировать, при этом задачи остаются нетривиальными, а результаты часто переносятся на более сложные домены. В Xland-MiniGrid агент решает задачи на основе системы правил и целей. Например:

➖ Агент видит на поле синюю пирамиду и фиолетовый квадрат. Его задача - поднять пирамиду и положить рядом с квадратом. Когда он это делает, срабатывает правило NEAR и оба объекта превращаются в красный круг.
➖ Появляется новая цель - поместить красный круг рядом с зелёным. Но если агент поместит фиолетовый квадрат рядом с жёлтым кругом, задача становится нерешаемой.

Такие правила можно комбинировать, создавая деревья задач разной глубины. В простых бенчмарках всего пара правил, в сложных - до 18. При этом позиции объектов рандомизируются при каждом сбросе среды, а правила и цели скрыты от агента. Чтобы решить задачу, ему нужно экспериментировать и запоминать, какие действия к чему приводят.

Работает это всё на бешенных скоростях - на одной RTX 4090 может симулироваться до 800к действий в секунду, а ведь можно использовать далеко не одну GPU. А для того чтобы не генерить с нуля, авторы уже сгенерили и выложили в опенсорс датасет на 100 миллиардов взаимодействий, о котором я уже писал.

Пейпер
Код

@ai_newz

www.group-telegram.com/es/ai_newz.com/3510

17.3K viewsedited  Dec 9 at 16:14

XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX

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

Но на CPU достичь больших скоростей трудно - даже несмотря на то что у современных процов часто сотни ядер, они просто не могут угнаться за современными видяхами, мощность которых часто измеряется в петафлопсах. И вот поэтому среды стали переносить на GPU, но писать CUDA код, особенно быстрый не каждый ресёрчер сможет

Поэтому не так давно пошла мода писать среды на Jax - фреймворке от Google, основном конкуренте PyTorch. Создавали его в том числе под DeepMind, поэтому тамошние ресёрчеры убедились чтобы в нём было достаточно гибкости для создания таких сред. Но хоть Jax и в опенсорсе, а свои среды Google никому не даёт.

XLand-MiniGrid, созданный чуваками из T-Bank AI Research и AIRI, как раз и есть опенсорс репродукция Xland, закрытой среды от Google. Это grid-world среда, где агент перемещается по сетке и взаимодействует с объектами по определённым правилам. Такие среды можно очень просто и быстро симулировать, при этом задачи остаются нетривиальными, а результаты часто переносятся на более сложные домены. В Xland-MiniGrid агент решает задачи на основе системы правил и целей. Например:

➖ Агент видит на поле синюю пирамиду и фиолетовый квадрат. Его задача - поднять пирамиду и положить рядом с квадратом. Когда он это делает, срабатывает правило NEAR и оба объекта превращаются в красный круг.
➖ Появляется новая цель - поместить красный круг рядом с зелёным. Но если агент поместит фиолетовый квадрат рядом с жёлтым кругом, задача становится нерешаемой.

Такие правила можно комбинировать, создавая деревья задач разной глубины. В простых бенчмарках всего пара правил, в сложных - до 18. При этом позиции объектов рандомизируются при каждом сбросе среды, а правила и цели скрыты от агента. Чтобы решить задачу, ему нужно экспериментировать и запоминать, какие действия к чему приводят.

Работает это всё на бешенных скоростях - на одной RTX 4090 может симулироваться до 800к действий в секунду, а ведь можно использовать далеко не одну GPU. А для того чтобы не генерить с нуля, авторы уже сгенерили и выложили в опенсорс датасет на 100 миллиардов взаимодействий, о котором я уже писал.

Пейпер
Код

@ai_newz

BY эйай ньюз