RLtools: быстрая портативная библиотека глубокого обучения с подкреплением для непрерывного управления

7.5/10

Ключевые слова: reinforcement learning, continuous control, deep learning, TinyRL.

Аннотация

Глубокое обучение с подкреплением (RL) способно давать эффективных агентов и управляющие политики в ряде предметных областей, но часто страдает от неприемлемо долгого обучения. Кроме того, в задачах непрерывного управления применимость выученных политик на реальных встраиваемых устройствах ограничена отсутствием real-time гарантий и портируемости существующих библиотек. Чтобы решить эти проблемы, мы представляем RLtools — библиотеку глубокого обучения с подкреплением и с учителем на чистом C++ без зависимостей, в формате header-only. Новая архитектура позволяет использовать RLtools на широком спектре платформ: от HPC-кластеров и рабочих станций до смартфонов, смарт-часов и микроконтроллеров. Благодаря плотной интеграции RL-алгоритмов со средой симуляции RLtools решает популярные RL-задачи до 76 раз быстрее других распространённых фреймворков. Мы также бенчмаркуем inference на разнообразном наборе микроконтроллеров и показываем, что в большинстве случаев наша оптимизированная реализация — самая быстрая. Наконец, RLtools демонстрирует первое в истории обучение алгоритма глубокого RL непосредственно на микроконтроллере, открывая область Tiny Reinforcement Learning (TinyRL). Исходный код, документация и live-демо доступны на странице проекта rl.tools.

1. Введение

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

Долгое обучение мешает быстрой итерации в пространстве задач (дизайн функции вознаграждения, подбор гиперпараметров и т.д.), а развёртывание на вычислительно ограниченных встраиваемых платформах необходимо для управления большинством физических систем: роботами, автомобильными компонентами, медицинскими приборами, smart grid и др. В нефизических системах, таких как финансовые рынки, потребность в высокочастотном принятии решений ведёт к схожим real-time требованиям, которые текущие deep RL библиотеки не выполняют.

Чтобы решить эти вызовы, мы представляем RLtools — header-only библиотеку на чистом C++ без зависимостей для глубокого обучения с учителем и с подкреплением, объединяющую следующие вклады:

• Новая архитектура. Мы описываем инновации в дизайне ПО, обеспечивающие беспрецедентную скорость обучения и inference на широком спектре устройств — от HPC-кластеров до смартфонов, смарт-часов и микроконтроллеров.
• Реализация. Модульная, высокопортируемая и эффективная open-source реализация архитектуры с документацией и тестами.
• Самое быстрое обучение. Крупные ускорения wall-clock времени обучения.
• Самый быстрый inference. Крупные ускорения времени inference обученных политик на типичных микроконтроллерах.
• TinyRL. С помощью RLtools мы впервые демонстрируем обучение алгоритма deep RL для непрерывного управления непосредственно на микроконтроллере.

2. Связанные работы

Предложено множество deep RL фреймворков и библиотек. Многие ориентированы на алгоритмические исследования — с абстракциями или без (Acme (Hoffman et al., 2020), skrl (Serrano-Muñoz et al., 2023) и CleanRL (Huang et al., 2022) соответственно). Другие фокусируются на полноте набора алгоритмов: RLlib (Liang et al., 2018), ReinforcementLearning.jl (Tian, 2020), MushroomRL (D'Eramo et al., 2021), Stable-Baselines3 (Raffin et al., 2021), ChainerRL (Fujita et al., 2021), Tianshou (Weng et al., 2022), TorchRL (Bou et al., 2024).

В отличие от них, RLtools нацелен на быструю итерацию в пространстве задач — например, дизайн функции вознаграждения (Eschmann, 2021) и оптимизацию гиперпараметров. В пространстве задач алгоритмические тонкости и разнообразие алгоритмов менее важны, чем робастность, скорость обучения, итоговая производительность и понимание того, как обучать их надёжно. Из перечисленных фреймворков RLlib (Liang et al., 2018) наиболее близок по миссии — быстрая итерация и развёртывание (ср. бенчмарки в разделе 4). Сосредоточившись на итерации в пространстве задач и последующем развёртывании на real-time платформах, мы проводим параллель между RLtools и ПО ACADOS (Verschueren et al., 2022) для синтеза MPC: RLtools стремится стать его RL-эквивалентом.

3. Подход

За последние годы в RL для непрерывного управления наиболее распространённые аппроксиматоры функций остаются относительно небольшими полносвязными нейросетями. В приложении A мы анализируем архитектуры из ключевых публикаций и обосновываем фокус RLtools на (небольших) MLP. Из этого следует, что гибкость фреймворков автоматического дифференцирования вроде TensorFlow или PyTorch может быть избыточной для многих задач непрерывного управления. Мы считаем выгодным обменять гибкость архитектуры аппроксиматора на общую скорость обучения. Сокращение времени обучения экономит энергию, упрощает воспроизводимость и демократизирует доступ к state-of-the-art RL. Быстрое обучение также облегчает систематический поиск гиперпараметров и улучшает сопоставимость результатов.

Архитектура

Архитектура ПО направлена на максимизацию эффективности обучения без потери returns. ПО должно работать на разных ускорителях и устройствах (CPU, GPU, микроконтроллеры и др.), чтобы обученные политики можно было напрямую развернуть на микроконтроллерах с device-specific инструкциями, обеспечивающими высокую частоту и жёсткие real-time гарантии. Это также означает отсутствие внешних зависимостей — они могут быть недоступны на целевых микроконтроллерах.

Для максимальной производительности компоненты интегрируются настолько плотно, насколько нужно, сохраняя модульность. Мы активно используем систему шаблонов C++. Template meta-programming даёт компилятору максимум информации о структуре кода для агрессивной оптимизации. Размеры всех циклов известны на этапе компиляции — компилятор может инлайнить и развернуть циклы (см. приложения B и F).

На чистом C++ без зависимостей реализованы: глубокое обучение (MLP, backpropagation, Adam и др.), RL (GAE, PPO, TD3, SAC) и симуляция (Pendulum, Acrobot, квадрокоптер, гоночный автомобиль, интерфейс MuJoCo). RLtools модульна благодаря новой парадигме статического multiple dispatch, вдохновлённой (динамическим) multiple dispatch из Julia (Bezanson et al., 2012). Рекомендуем пример кода в приложении B и ablation study в приложении F.

4. Результаты

Горизонтальный бенчмарк

На рисунках 1 и 2 — среднее время обучения PPO и SAC за десять запусков на ноутбуке на Intel (детали в таблице 6). RLtools опережает существующие библиотеки с большим отрывом: PPO в среднем 0,54 с (SAC — 2,59 с).

Рисунок 1 — смотреть в оригинале (PDF)

Рисунок 2 — смотреть в оригинале (PDF)

Вертикальный бенчмарк

На рисунке 3 — результаты обучения RLtools на широком спектре устройств, несовместимых с другими RL-библиотеками. Важно, мы демонстрируем первое обучение deep RL агента для непрерывного управления на микроконтроллере Teensy 4.1.

Рисунок 3 — смотреть в оригинале (PDF)

Inference на микроконтроллерах

Таблица 1 показывает время inference на микроконтроллерах разной вычислительной мощности (Crazyflie — 27-граммовый квадрокоптер с очень ограниченными ресурсами, см. приложение E). Generic-реализация уже даёт приемлемое время, но dispatch к DSP-библиотеке производителя улучшает производительность. Дополнительная оптимизация (например, fusion операторов активации) даёт значительное ускорение даже относительно DSP-библиотек производителя.

5. Заключение

Мы считаем, что RLtools заполняет пробел: быстрая итерация в пространстве задач и последующее real-time развёртывание политик. Кроме того, RLtools впервые делает возможным обучение deep RL на микроконтроллере. Мы признаём более крутую кривую обучения C++ (по сравнению с Python), но по нашему опыту более быстрая итерация за счёт короткого обучения может перевесить начальные затраты. Сейчас RLtools ограничен плотными наблюдениями; планируем добавить vision. Снижая требования к вычислениям и будучи полностью open-source, RLtools демократизирует обучение state-of-the-art RL и ускоряет прогресс в непрерывном управлении.

Благодарности и раскрытие финансирования

Работа поддержана Technology Innovation Institute, NSF CAREER Award 2145277 и DARPA YFA Grant D22AP00156-00. Giuseppe Loianno консультирует Technology Innovation Institute. Эта договорённость проверена и одобрена New York University в соответствии с политикой объективности исследований.

Приложение A. Анализ ландшафта deep RL

В этом разделе мы анализируем модели аппроксиматоров функций из ключевых публикаций по deep RL для непрерывного управления (таблица 2 и таблица 7). Главное наблюдение: за все годы архитектура (небольшие полносвязные сети) не изменилась. В непрерывном управлении наблюдения — обычно плотные состояния системы без пространственной или временной структуры, как у изображений или временных рядов, что не требует CNN или RNN. Эта регулярность мотивирует фокус RLtools на оптимизации и плотной интеграции MLP как первый шаг; в будущем планируем recurrent и, возможно, свёрточные слои.

Приложение B. Парадигма программирования

Для максимальной производительности мы избегаем lookup в таблице виртуальных методов C++ (VMT), не используя классическую ОО-парадигму, а опираясь на функциональный стиль с шаблонами и перегрузкой — статическую compile-time интерпретацию multiple dispatch.

// file: implementation_generic.h
template <typename DEVICE, auto M, auto N, auto K>
void multiply(DEVICE device, Matrix<M, K> a, Matrix<K, N> b, Matrix<M, N> result){
  // Generic code for matrix multiplication
  ...
}
// file: implementation_microcontroller.h
template <auto M, auto N, auto K>
void multiply(MICROCONTROLLER device, Matrix<M, K> a, Matrix<K, N> b, Matrix<M, N> result){
  // Optimized code for matrix multiplication on a particular microcontroller
  ...
}
// file: implementation_gpu.h
template <auto M, auto N, auto K>
void multiply(GPU device, Matrix<M, K> a, Matrix<K, N> b, Matrix<M, N> result){
  // Optimized GPU code for matrix multiplication
  ...
}
template<typename DEVICE, typename OBJECT_A, typename OBJECT_B, typename OBJECT_C>
void algorithm(DEVICE device, OBJECT_A a, OBJECT_B b, OBJECT_C c){
  ...
  multiply(device, a, b, c);
  ...
}

Рисунок 4 — смотреть в оригинале (PDF)

Multiple dispatch популяризован Julia (Bezanson et al., 2012) и основан на продвинутой перегрузке функций. Высокоуровневые функции (forward/backward pass MLP) задают действия на подкомпонентах, а конкретная реализация зависит от типов аргументов. Так легко делить код между GPU и CPU: достаточно реализовать низкоуровневые примитивы для каждого типа устройства и сигнализировать тип через tag dispatch. Если специализированной реализации нет, компилятор использует generic-версию (например, вложенные циклы) — чистый C++ без зависимостей. Размеры операндов проверяются на этапе компиляции. Для сложного dispatch используются static_assert и enable_if через SFINAE.

В Julia multiple dispatch разрешается в runtime, что даёт высокую композируемость (Flux (Innes, 2018) конкурирует с PyTorch при меньших ресурсах). Мы же разрешаем все вызовы на этапе компиляции. Julia не подходит для микроконтроллеров из-за размера runtime и недетерминированного GC-поведения. Тем не менее в наших бенчмарках Julia — один из ближайших конкурентов по скорости обучения.

Мы строим библиотеку, а не фреймворк. Фреймворки ограничивают свободу пользователя ради упрощения узкого набора задач; накладные расходы на крутую кривую обучения и обход ограничений не всегда оправданы. Фреймворки вызывают код пользователя из своего контекста; библиотеки — пользователь полностью контролирует вызов компонентов. Контексты фреймворков обычно несовместимы; библиотеки — совместимы по умолчанию.

В RLtools любой тип аппроксиматора функции задаётся пользователем на этапе компиляции через шаблоны — при условии, что предоставлены forward и backward. Установлена конвенция передавать device-контекст в каждую функцию через tag dispatch для упрощения работы с разными вычислительными устройствами.

Приложение C. Детали бенчмарков

Приложение D. Фреймворки и библиотеки deep RL

Приложение E. Встраиваемые платформы

1. Crazyflie — небольшой open-source квадрокоптер массой 27 г с батареей. Основной процессор STM32F405 (ARM Cortex-M4), 192 KB RAM, 168 MHz.
2. Pixhawk 6C — используем Pixracer Pro (FCU семейства Pixhawk, стандарт Pixhawk 6C), поддерживает прошивку PX4 (Meier et al., 2015) для разных типов аппаратов. Процессор STM32H743 (ARM Cortex-M7), 460 MHz, 1024 KB RAM.
3. Teensy 4.1 — универсальное встраиваемое устройство на i.MX RT1060 (ARM Cortex-M7), 1024 KB on-chip и 16 MB off-chip RAM, 600 MHz.
4. ESP32 — один из самых распространённых IoT/edge микроконтроллеров со встроенными Wi-Fi и Bluetooth; продано около миллиарда устройств на этой и предыдущей версии чипа (~$5 за dev kit). Процессор Xtensa LX7; также оцениваем ESP32-C3 на RISC-V.

Приложение F. Ablation study

Мы проводим ablation study, чтобы оценить вклад компонентов и оптимизаций в быстрое wall-clock обучение RLtools. На рисунке 5 — время обучения после удаления разных компонент. «Baseline» — та же конфигурация, что в эксперименте SAC Pendulum-v1 в разделе 4.

Мы симулируем медлительность Python-среды, замедляя C++-реализацию на среднее время шага Python-версии. C++-реализация динамики Pendulum-v1 даёт измеримый, но не доминирующий вклад. Аблируем уровни оптимизации компилятора -O0 … -O3 (baseline использует -O3): при -O0 RLtools примерно между ACME и CleanRL (рисунок 2). Конфигурация «No Fast Math» убирает -ffast-math; «BLAS» — Intel oneMKL для умножения матриц; «AVX/AVX2» отключает векторные расширения.

Благодаря дизайну RLtools (приложение B), когда размеры циклов и структур данных известны на этапе компиляции, компилятор активно использует SIMD. Из 11 243 машинных инструкций 2276 + 1430 = 3706 (AVX + SSE) относятся к векторным регистрам. При отключении AVX компилятор заменяет инструкции на SSE (5406 из 11 243), что демонстрирует «дружелюбность» архитектуры к компилятору.

Рисунок 5 — смотреть в оригинале (PDF)

Приложение G. Исследование сходимости

Чтобы убедиться в корректности реализаций PPO, TD3 и SAC, мы сравниваем кривые обучения с другими реализациями на одинаковых гиперпараметрах. Для Pendulum-v1 — 100 seeds, для Hopper-v4 — 30; на каждом шаге оценки — 100 эпизодов со случайными начальными состояниями.

Даже при большом числе seeds выбросы существенно влияют на средний финальный return. По рекомендации Agarwal et al. (2021) мы отчитываем Inter Quantile Mean (IQM), отбрасывая нижний и верхний квантили (по 5 %). Для идентичности сред вместо переимплементации в C++ построен Python-обёртка RLtools для сред Gymnasium (Towers et al.). Обёртка упрощает использование, но снижает производительность, если среда на Python (см. приложение F).

Рисунок 6 — смотреть в оригинале (PDF)

Рисунок 7 — смотреть в оригинале (PDF)

Рисунок 8 — смотреть в оригинале (PDF)

Рисунок 9 — смотреть в оригинале (PDF)

Рисунок 10 — смотреть в оригинале (PDF)

Рисунок 11 — смотреть в оригинале (PDF)

Рисунок 12 — смотреть в оригинале (PDF)

Рисунок 13 — смотреть в оригинале (PDF)

Рисунок 14 — смотреть в оригинале (PDF)

Литература

• Agarwal, R., Schwarzer, M., Castro, P.S., Courville, A.C., Bellemare, M. Deep reinforcement learning at the edge of the statistical precipice. NeurIPS 34, 2021.
• Bach, N., Melnik, A., Schilling, M., Korthals, T., Ritter, H. Learn to Move Through a Combination of Policy Gradient Algorithms: DDPG, D4PG, and TD3. LOD 2020.
• Bezanson, J., Karpinski, S., Shah, V.B., Edelman, A. Julia: A Fast Dynamic Language for Technical Computing. arXiv:1209.5145, 2012.
• Bou, A., Bettini, M., Dittert, S., Kumar, V., Sodhani, S., Yang, X., De Fabritiis, G., Moens, V. TorchRL: A data-driven decision-making library for PyTorch. ICLR 2024.
• D'Eramo, C., Tateo, D., Bonarini, A., Restelli, M., Peters, J. MushroomRL: Simplifying reinforcement learning research. JMLR 22(131), 2021.
• Eschmann, J. Reward function design in reinforcement learning. Reinforcement Learning Algorithms: Analysis and Applications, 25–33, 2021.
• Fujimoto, S., van Hoof, H., Meger, D. Addressing function approximation error in actor-critic methods. ICML 2018.
• Fujita, Y., Nagarajan, P., Kataoka, T., Ishikawa, T. ChainerRL: A deep reinforcement learning library. JMLR 22(77), 2021.
• Haarnoja, T., Zhou, A., Abbeel, P., Levine, S. Soft actor-critic: Off-policy maximum entropy deep RL with a stochastic actor. ICML 2018.
• Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., Levine, S. Soft Actor-Critic Algorithms and Applications. arXiv:1812.05905, 2019.
• Hoffman, M.W., Shahriari, B., Aslanides, J., et al. Acme: A research framework for distributed reinforcement learning. arXiv:2006.00979, 2020.
• Huang, S., Dossa, R.F.J., Ye, C., Braga, J., Chakraborty, D., Mehta, K., Araujo, J.G.M. CleanRL: High-quality single-file implementations of deep RL algorithms. JMLR 23(274), 2022.
• Innes, M. Flux: Elegant machine learning with Julia. JOSS 3(25):602, 2018.
• Kumar, A., Fu, Z., Pathak, D., Malik, J. RMA: Rapid Motor Adaptation for Legged Robots. RSS 2021.
• Kuznetsov, A., Shvechikov, P., Grishin, A., Vetrov, D. Controlling overestimation bias with truncated mixture of continuous distributional quantile critics. ICML 2020.
• Liang, E., Liaw, R., Nishihara, P., Moritz, P., Fox, R., Goldberg, K., Gonzalez, J., Jordan, M., Stoica, I. RLlib: Abstractions for distributed reinforcement learning. ICML 2018.
• Lillicrap, T.P., Hunt, J.J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., Wierstra, D. Continuous control with deep reinforcement learning. arXiv:1509.02971, 2016.
• Meier, L., Honegger, D., Pollefeys, M. PX4: A node-based multithreaded open source robotics framework. ICRA 2015.
• Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., Dormann, N. Stable-baselines3: Reliable reinforcement learning implementations. JMLR 22(268), 2021.
• Schulman, J., Levine, S., Abbeel, P., Jordan, M., Moritz, P. Trust region policy optimization. ICML 2015.
• Schulman, J., Moritz, P., Levine, S., Jordan, M., Abbeel, P. High-dimensional continuous control using generalized advantage estimation. ICLR 2016.
• Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O. Proximal Policy Optimization Algorithms. arXiv:1707.06347, 2017.
• Serrano-Muñoz, A., Chrysostomou, D., Bøgh, S., Arana-Arexolaleiba, N. skrl: Modular and flexible library for reinforcement learning. JMLR 24(254), 2023.
• Tian, J. ReinforcementLearning.jl: A reinforcement learning package for the Julia programming language, 2020.
• Towers, J., Terry, J.K., Kwiatkowski, A., et al. Gymnasium. github.com/Farama-Foundation/Gymnasium
• Verschueren, R., Frison, G., Kouzoupis, D., et al. acados—a modular open-source framework for fast embedded optimal control. Mathematical Programming Computation 14(1), 147–183, 2022.
• Weng, J., Chen, H., Yan, D., You, K., Duburcq, A., Zhang, M., Su, Y., Su, H., Zhu, J. Tianshou: A highly modularized deep reinforcement learning library. JMLR 23(267), 2022.

Оригинал статьи: Eschmann, J., Albani, D., Loianno, G., «RLtools: A Fast, Portable Deep Reinforcement Learning Library for Continuous Control», JMLR 25(301), 2024 · код: rl-tools/rl-tools · лицензия CC BY 4.0