JAX-LOB: симулятор лимитного стакана с ускорением на GPU, открывающий путь к крупномасштабному обучению с подкреплением для торговли

9/10

Ключевые слова: лимитные стаканы, обучение с подкреплением, высокочастотная торговля, исполнение сделок, воспроизведение рынка, симулятор стакана.

Аннотация

Финансовые биржи по всему миру используют лимитные стаканы (limit order books, LOB) для обработки заявок и сведения сделок. Для исследовательских целей важно иметь крупномасштабные эффективные симуляторы динамики стакана. Симуляторы стакана ранее реализовывались в контексте агентных моделей (ABM), сред обучения с подкреплением (RL) и генеративных моделей, обрабатывая потоки заявок как из исторических наборов данных, так и от сконструированных вручную агентов. Для многих приложений требуется обработка нескольких книг — либо для калибровки агентных моделей, либо для обучения RL-агентов. Мы представляем первый симулятор стакана с поддержкой GPU, спроектированный для параллельной обработки тысяч книг с заметно сниженным временем обработки одного сообщения. Реализация нашего симулятора — JAX-LOB — основана на проектных решениях, нацеленных на наилучшее использование возможностей JAX без ущерба реалистичности механизмов, связанных со стаканом. Мы интегрируем JAX-LOB с другими пакетами JAX, чтобы показать пример решения задачи оптимального исполнения с помощью обучения с подкреплением и поделиться предварительными результатами сквозного RL-обучения на GPU.

Категории CCS: Вычислительные методологии → Дискретно-событийная симуляция; Обучение с подкреплением; Массово-параллельные и высокопроизводительные симуляции.

1. Введение

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

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

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

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

В разделе 2 мы даём краткий обзор стакана и задачи исполнения сделок, а в разделе 3 рассматриваем связанные работы как по построению симуляторов стакана, так и по решению задачи исполнения с помощью RL. В разделе 4 мы представляем первый симулятор стакана с ускорением на GPU на основе JAX и обсуждаем детальное устройство стакана — как общее, так и специфичное для нашей реализации. В разделе 5 мы закладываем основу для использования нашего симулятора в RL — области применения, значительно выигрывающей от параллелизма, который обеспечивает наша реализация. В разделе 5.2 мы вносим вклад в решение задачи оптимального исполнения с помощью RL, обернув симулятор JAX-LOB в нативную для JAX среду исполнения gymnax. В нашей среде и сбор опыта (то есть взаимодействие агентов с миром для сбора данных), и обновления обучения (то есть обучение агентов на данных) выполняются на одном и том же GPU, что устраняет узкое место в обмене данными между GPU и CPU. Наконец, в разделе 6 мы интегрируем среду с PureJaxRL и используем рекуррентную оптимизацию проксимальной политики (PPO) для обучения агента исполнения в качестве примера применения. В завершение мы сравниваем вычислительную скорость с эквивалентными реализациями на CPU и показываем, что максимальное ускорение достигается при применении JAX-LOB к задаче RL, где полностью используется параллелизм.

Наши основные вклады таковы:

1. Первый симулятор стакана с ускорением на GPU — JAX-LOB.
2. Ускорение обучения RL-агента как минимум в 7 раз по сравнению с эквивалентными симуляторами стакана на CPU: 550 против 74 шагов в секунду во время обучения на одном и том же оборудовании (раздел 6.2).

Мы обязуемся открыть исходный код JAX-LOB, включая интеграцию с RL, и твёрдо убеждены, что это откроет сообществу ряд новых направлений исследований.

2. Предпосылки

Стакан — лежащая в основе структура данных и связанные с ней механизмы, на которых работают современные электронные биржи, — изучен весьма подробно. Стакан представляет собой совокупность всех заявок на покупку или продажу ценной бумаги, которые ещё не были сведены. В стакан в любой момент могут подаваться заявки двух типов: лимитные заявки, требующие указания цены и количества, и рыночные заявки, требующие только количества. Трейдеры вправе подавать заявки в любой момент, после чего заявка либо сводится с совместимой противоположной заявкой в книге, либо — в случае лимитной заявки — добавляется в книгу. Дальнейшие детали устройства стакана обсуждаются в разделе 4 вместе с деталями, специфичными для нашей реализации. У симуляторов стакана есть ряд областей применения, включая генеративные модели; в этой работе мы сосредоточены на их использовании для обучения торговых агентов методами RL для таких задач, как маркетмейкинг или исполнение сделок.

JAX — независимый от ускорителя фреймворк, обеспечивающий JIT-компиляцию с использованием ускоренной линейной алгебры (XLA), автоматическое дифференцирование и автоматическую векторизацию, которые легко исполняются на GPU. Фреймворк спроектирован для высокопроизводительных исследований в области машинного обучения и лежит в основе фреймворков gymnax и PureJaxRL, которые мы используем, взяв оптимальное исполнение сделок в качестве примера задачи. JAX применялся и в других приложениях, где выгоден массовый параллелизм, — например, в гидродинамике или молекулярной динамике. Хотя эта работа сосредоточена на JIT-компиляции и векторизации на GPU, дифференцирование моделей стакана для калибровки остаётся темой, недостаточно исследованной в литературе.

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

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

Существует ряд реализаций, воспроизводящих механику стакана. Их можно далее разделить по области применения: агентные модели (ABM) и воспроизведение рынка (market replay). Фреймворки, построенные для ABM, такие как ABIDES (agent-based interactive discrete event simulation) или MAXE (multi-agent exchange environment), реализуют разнородный набор агентов, подающих в книгу вымышленные заявки. ABIDES реализован на обычном Python и имеет расширение ABIDES-gym, создающее gym-подобный интерфейс для обучения с подкреплением. MAXE реализован на C++ и выигрывает в скорости за счёт компиляции. В моделях воспроизведения рынка, таких как фреймворк RL4MM (reinforcement learning for market making), подаются исторические заявки на основе данных сообщений, например из набора данных LOBSTER («limit order book system — the efficient reconstructor»).

Основное ограничение агентных моделей — трудность создания достаточно реалистичного поведения агентов. Несмотря на значительные недавние усилия, это остаётся ограничением, особенно в сравнении с воспроизводимыми рыночными данными, которые напрямую воспроизводят распределение данных. Однако сильная сторона агентных моделей одновременно является ключевым ограничением воспроизведения рынка. Реализация стратегических агентов приводит к динамически реагирующему потоку заявок при внешнем вмешательстве (например, со стороны RL-агента), тогда как воспроизводимый исторический поток заявок останется неизменным. JAX-LOB, изначально спроектированный для обработки данных LOBSTER, может обрабатывать сообщения из любого источника, включая стратегических агентов. В отличие от обсуждённых реализаций, JAX-LOB работает на GPU, что позволяет масштабно распараллеливать вычисления.

Ning и соавторы используют двойные глубокие Q-сети (DDQN) для решения задачи оптимального исполнения, применяя только рыночные заявки. Такие заявки, однако, в целом субоптимальны, поскольку сделки могут «идти по книге», то есть потреблять весь доступный объём на одном или нескольких уровнях стакана, и лишают возможного выигрыша от спреда, который даёт использование пассивных лимитных заявок. Кроме того, рыночное воздействие там стилизовано через штрафную функцию для крупных заявок. В такой постановке они превосходят бенчмарк TWAP (средневзвешенная по времени цена) на 7 из 9 рассмотренных акций. Dabérius и соавторы сравнивают использование DDQN и PPO для оптимального исполнения с применением только рыночных заявок в стохастической модели ценовой динамики без ценового воздействия. Создавая симулятор стакана, мы даём инструменты, позволяющие уйти от необходимости стилизовать рыночное воздействие и создать среду, допускающую подачу лимитных заявок по разным ценам, тем самым избегая риска непреднамеренно потребить объём на нескольких уровнях книги. Тем не менее исключительное использование исторических данных стакана означает, что учитывается только прямое рыночное воздействие: у таких исторических сообщений нет стратегического поведения, которое реагировало бы на действия RL-агента.

Устраняя этот недостаток, Karpe и соавторы используют DDQN в агентной симуляционной среде на основе симулятора ABIDES, которая воспроизводит реальные данные сообщений стакана, добавляя при этом дополнительных эвристических агентов, реагирующих на изменения состояния стакана, вызванные RL-агентом. Их постановка симуляции позволяет решать задачи оптимального исполнения и оптимального размещения по уровням совместно. Они обнаруживают, что их агент сходится к стратегии TWAP. Хотя дополнительные агенты в принципе позволяют получить более реалистичное рыночное воздействие, эвристики ABIDES используют довольно упрощённую моментум-стратегию на основе агрегированных сводных статистик вместо выученной рациональной стратегии. Аналогично, Fang и соавторы используют PPO для решения совмещённой задачи оптимального исполнения и размещения, применяя подход воспроизведения исторических данных в симуляторе ABIDES. Их агент обучается размещению одновременно до трёх уровней стакана. Они используют большое число сконструированных признаков, включая технические индикаторы, архитектуру сети на основе внимания под названием Dual-Window Denoise PPO и формулировку награды на основе издержек исполнения и имитационного члена, направляющего политику в сторону TWAP. Результаты указывают на потенциальное улучшение по сравнению с TWAP, но страдают от больших стандартных ошибок. Среда, которую мы создаём, нацелена на отказ от таких сконструированных пространств признаков и функций награды и сосредоточена на использовании исключительно рекуррентных нейронных сетей для автоматического извлечения признаков и учёта более длительной памяти агента.

4. Симулятор JAX-LOB

Некоторые из ключевых сложностей применения глубокого обучения с подкреплением к исполнению сделок и другим высокочастотным задачам — низкое отношение сигнал/шум, риск переобучения на конкретные обучающие дни и построение симулятора с реалистичным рыночным воздействием. Прямой способ компенсировать первые две проблемы — увеличить число доступных для обучения переходов «состояние — действие». Чтобы ускорить генерацию наблюдений точных представлений стакана на высокочастотных данных, мы используем JAX.

Мы принимаем ряд проектных решений, чтобы учесть некоторые ограничения JIT-компилируемого кода на JAX:

• чистые функции без побочных эффектов (например, без глобальных переменных);
• фиксированный размер и тип массивов;
• поток управления, который может быть эффективно скомпилирован и распараллелен (см. раздел 4.3).

Большинство реализаций стакана на CPU основаны на хеш-таблицах, очередях, двусвязных списках и отсортированных словарях, которые обеспечивают быстрый доступ к данным и поддерживают книгу отсортированной всё время. Учитывая требование JAX о массивах фиксированного размера на этапе компиляции, реализация аналогичной структуры означает, что память должна быть предварительно выделена под все ценовые уровни и заявки. Использование массивов означает, что переупорядочивание при удалении элементов обходится гораздо дороже, чем со связными списками. Поэтому мы выбираем архитектуру, которая не использует древовидную структуру и не держит заявки отсортированными всё время. Вместо этого мы определяем два массива $A$ и $B$, представляющие стороны стакана и содержащие соответственно все активные заявки ask и bid.

Каждая сторона книги имеет фиксированную ёмкость на $N$ заявок, где каждая заявка $\boldsymbol{o}$ имеет шесть признаков (2): цена $P$, количество $Q$, идентификатор заявки $OID$, идентификатор трейдера $TID$, время в секундах $Ts$ и время в наносекундах $Tns$. Пустые позиции в $A$ или $B$ обозначаются установкой всех признаков в $-1$. Размер $N$ должен быть выбран так, чтобы книга не переполнялась в ходе конкретного эксперимента.

4.1. Базовые операции

Есть три операции, которые могут быть применены к любой стороне стакана: добавление новой заявки, отмена существующей заявки и сведение существующей заявки со входящей заявкой на противоположной стороне книги с последующим её удалением из книги.

1. Добавление заявки требует нахождения пустой позиции ($\boldsymbol{o_i} = -1$) в массиве и вставки данных заявки в соответствующие поля.
2. Отмена требует нахождения идентификатора заявки ($OID$), подлежащей отмене, и удаления соответствующего количества из книги.
3. При операции сведения входящая агрессивная заявка, обозначаемая $\boldsymbol{o_a}$, сводится с существующей заявкой на противоположной стороне книги — стоящей заявкой $\boldsymbol{o_s}$. Логика сведения находит остаточные количества для агрессивной ($Q_a \in \boldsymbol{o_a}$) и стоящей ($Q_s \in \boldsymbol{o_s}$) заявок с помощью операций \[ Q'_s = \max(0, Q_s - Q_a), \qquad Q'_a = Q_a - Q_s. \]

Когда две заявки сводятся, записывается сделка $\boldsymbol{t_j}$ (3).

При этом:

До $N$ сделок записывается в массив фиксированного размера $T$ из-за ограничений компилятора XLA (4).

Здесь также $\boldsymbol{t_j} = -1$ обозначает пустой слот. Для всех операций по завершении и $Asks$, и $Bids$ проверяются на наличие заявок $\boldsymbol{o_i}$, для которых $Q_i \le 0$; в этом случае мы полагаем $\boldsymbol{o_i} = -1$.

Одна входящая заявка может свестись более чем с одной стоящей заявкой. Логика сведения содержит цикл while, который многократно пытается свести входящую заявку со следующей наилучшей стоящей заявкой на противоположной стороне книги. Наилучшая стоящая заявка $Best(\boldsymbol{o_s})$ определяется алгоритмом приоритета «цена — время», наиболее распространённым алгоритмом сведения в стакане. Наилучшая заявка — та, у которой цена ask или цена bid (5). Если несколько заявок имеют эту цену, рассматривается заявка с наиболее ранним временем прихода.

Примечание переводчика: в оригинале для $P_{bid}$ также записан минимум; по смыслу приоритета «цена — время» лучшей ценой bid является максимум.

Цикл while продолжается, пока книга непуста, $Q_a > 0$, и цены перекрываются следующим образом:

Исходя из вышеописанного, интуитивно ожидаема разная вычислительная сложность этих базовых операций. Чтобы проверить эту интуицию, мы измеряем время базовых операций в стаканах различной максимальной ёмкости $N$ (таблица 1). Операции занимают больше времени с ростом ёмкости книги $N$, и, как и ожидалось, самая медленная операция (сведение) выполняется более чем вдвое дольше самой быстрой (отмена).

Относительный размер входящей заявки $Q_a \in \boldsymbol{o_a}$ существенно влияет на время вычисления операции сведения. При прочих равных большее количество требует рассмотрения большего числа стоящих заявок для сведения приходящей заявки, а значит, большего числа итераций логики сведения. Простой способ это проиллюстрировать — рассмотреть стакан ёмкостью $N = 100$, заполненный на треть, и подавать рыночные заявки разного размера (таблица 2). Рост времени не драматичен, пока не подаются особенно крупные заявки, но показывает важность ограничения пространства действий разумными количествами для обеспечения быстрого исполнения.

4.2. Типы сообщений

То, какая из трёх базовых операций (раздел 4.1) вызывается, зависит от типа $T$ сообщения $\boldsymbol{m}$ (6), передаваемого в стакан.

Структура данных, содержащихся в сообщениях $\boldsymbol{m}$:

1. $T$ — тип сообщения: лимитные заявки ($T = 1$), отмены ($T = 2$), удаления ($T = 3$) и рыночные заявки ($T = 4$).
2. $S$ — сторона заявки: bid ($S = 1$) или ask ($S = -1$).
3. $P$ — цена, по которой должна быть подана заявка, либо цена заявки, подлежащей отмене или удалению. Для рыночных заявок игнорируется.
4. $Q$ — размер заявки, то есть количество к покупке или продаже, либо количество, снимаемое с существующей заявки для заявок отмены и удаления.
5. $OID$ — уникальный идентификатор заявки.
6. $TID$ — идентификатор трейдера, метка источника заявки. Используется для идентификации заявок отдельных агентов.
7. $Ts, Tns$ — время получения заявки, разделённое на два поля: полные секунды и дробные разряды в наносекундах, представленные 32-битными целыми.

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

1. Лимитные заявки сначала сводятся с противоположной стороной книги, оставшееся количество затем добавляется как одна новая заявка на соответствующую сторону.
2. Заявки отмены и удаления обрабатываются одинаково — вызовом отмены (раздел 4.1).
3. Рыночные заявки сводятся с противоположной стороной с ценой $P_m = 0$, если это заявка ask, или $P_m = \max\_int$, если это заявка bid. Любое несведённое количество отбрасывается.

Рассмотрение каждого типа и каждой стороны как восьми отдельных случаев позволяет определить явные функции для каждого случая, что даёт возможность использовать единственный условный оператор при получении сообщения вместо нескольких ветвлений внутри логики сведения; мы обнаруживаем, что это улучшает производительность под vmap (раздел 4.3). Время вычисления (таблица 3) для обработки каждого из трёх типов сообщений различается и связано с требуемыми базовыми операциями. По сравнению с результатами таблицы 1 время лишь незначительно увеличивается за счёт оператора ветвления, и по-прежнему наблюдается заметная разница между разными типами заявок.

4.3. Векторизующее отображение — vmap

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

Одна из особенностей оператора vmap состоит в том, что он преобразует ряд операторов потока управления в операторы выбора, которые во время выполнения исполняют все условные ветви. В случае лимитных стаканов это означает, что вычисляется каждый из восьми возможных случаев, и время выполнения, следовательно, ограничено самой медленной ветвью. Это можно наблюдать в результатах измерений (таблица 4) при обработке одного и того же сообщения на $N\_books = 1000$ идентичных стаканах параллельно. Время больше не варьируется между разными типами сообщений, а зависит только от ёмкости. Причина, по которой мы избегаем непрерывной сортировки, в том, что при удалении заявок в представлении на основе массива требуется пересортировать весь массив. Такая сортировка сравнительно затратна, что не было бы проблемой, если бы она вызывалась редко. Однако упомянутое поведение оператора vmap означает, что она вызывалась бы для каждого сообщения, даже если ни одну заявку удалять не нужно. Хотя мы провели обширные тесты, возможно, что узкое место такого рода всё ещё существует, что объясняло бы худшее время обработки сообщения свыше 2 секунд в таблице 4.

Полученные времена по типам сообщений сравниваются с двумя реализациями на CPU (таблица 5): одна реализация стакана на красно-чёрных деревьях и связных списках, а также схожая реализация под названием RL4MM на numpy. Среды на CPU описаны подробнее в разделах 5 и 6. При распараллеливании более чем на 1000 книг наш JAX-LOB даёт меньшее время обработки одного сообщения.

5. Среды gymnax для стакана

Наш JAX-LOB хорошо подходит для фреймворка сред обучения с подкреплением gymnax, спроектированного вокруг JAX для параллельных прогонов среды на GPU. В следующих разделах мы кратко описываем построение общей базовой среды, которая может использоваться для порождения сред под различные задачи, такие как маркетмейкинг, исполнение сделок или другие внутридневные задачи. Далее мы описываем нашу примерную среду исполнения заявок, которую используем в разделе 6 для начала обучения рекуррентного агента на основе PPO.

5.1. Базовая среда

Основные роли базовой среды — загружать данные LOBSTER для воспроизведения, предоставлять интерфейс к стакану на JAX и задавать каркас торговой среды на основе JAX-LOB.

5.1.1. Загрузка данных. Мы предпочитаем предварительно загружать данные сообщений фиксированными непересекающимися временными окнами (рисунок 1), чтобы сэкономить время при выполнении. Число сообщений на шаг постоянно, а длительность шага переменна. Наоборот, длительность одного временного окна фиксирована, что подразумевает переменный горизонт, то есть максимальное число шагов. Это позволяет исполнять сделки при строгом ограничении по времени. Чтобы удовлетворить ограничению на массивы фиксированного размера, добавляются дополнительные шаги (обозначенные серыми рамками на рисунке 1) со всеми значениями, установленными в ноль. Эти шаги никогда не обрабатываются, поскольку условие завершения эпизода выполняется раньше.

Далее, состояние стакана (уровень 2) первых десяти уровней книги в начале каждого временного окна используется для инициализации книги. В отличие от данных сообщений, это не даёт понимания того, содержит ли ценовой уровень одну или несколько заявок, и не предоставляет идентификаторов заявок. Мы обходим это, предполагая, что на каждом ценовом уровне есть ровно одна заявка, содержащая сумму перечисленных объёмов. Для этих начальных заявок мы используем $OID_i \in [-\inf, -9000]$, начиная с $-9000$ и по убыванию. Мы допускаем отмену этих заявок, если сообщение отмены удовлетворяет $P_i = P_{cancel}$ и $OID_i < -9000$.

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

5.1.2. Основные требования к пространству состояний. Базовая среда определяет основные признаки, необходимые в пространстве состояний; для производных сред могут добавляться дополнительные элементы:

• Стакан:
  • сторона ask: Array($N_{orders} \times N_{features}$);
  • сторона bid: Array($N_{orders} \times N_{features}$).
• Сделки: Array($N_{trades} \times N_{features}$).
• Начальное время: $[T_s, T_{ns}]$.
• Текущее время: $[T_s, T_{ns}]$.
• Счётчик идентификаторов.
• Индекс окна.
• Счётчик шагов.

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

5.1.3. Функциональность. Мы оставляем описание базовой среды кратким, отнеся рассмотрение пространств действий и наблюдений, а также функции награды к среде исполнения в разделе 5.2. Основная функциональность среды в течение шага такова:

1. Полученные действия (зависящие от среды) преобразуются в сообщения для обработки в JAX-LOB.
2. Сообщения данных берутся из загруженных данных в параметрах среды с использованием индекса окна и счётчика шагов (раздел 5.1.1) для индексации.
3. Текущее время обновляется на основе последнего сообщения данных.
4. Все сообщения обрабатываются JAX-LOB.

Эпизод завершается, когда прошедшее время превышает время эпизода.

5.2. Среда исполнения

Среда исполнения расширяет пространство состояний и определяет пространства наблюдений и действий, функцию награды и новое условие завершения.

5.2.1. Расширенное пространство состояний. Пространство состояний из 5.1.2 дополняется внутришаговыми данными и информацией, специфичной для агента:

• Данные уровня 1 для каждого обработанного сообщения из данных с момента последнего шага:
  • лучшая цена ask: Array($N_{messages}$);
  • лучший объём ask: Array($N_{messages}$);
  • лучшая цена bid: Array($N_{messages}$);
  • лучший объём bid: Array($N_{messages}$).
• Начальная средняя цена $\frac{P_{ask} + P_{bid}}{2}$ в начале эпизода.
• Размер задачи исполнения.
• Исполненное к текущему моменту количество.
• Суммарная выручка от исполнения к текущему моменту.

5.2.2. Пространство наблюдений. Пространство состояний является внутренним для среды. Пространство наблюдений доступно RL-агенту и спроектировано так, чтобы содержать информацию, доступную в реальной жизни, и избегать избыточности. Структура представляет собой единый одномерный массив размера $N_{messages} \times 6 + 2 \times 2 + 4$, содержащий следующие метрики:

• лучшие цены bid между шагами: Array($N_{messages}$);
• лучшие цены ask между шагами: Array($N_{messages}$);
• средние цены между шагами: Array($N_{messages}$);
• цены на $n$ тиков вглубь на стороне книги, соответствующей задаче, между шагами: \[ P_{passive} = \begin{cases} P_{ask} + n \times ticksize, & \text{если задача — продажа} \\ P_{bid} - n \times ticksize, & \text{если задача — покупка} \end{cases} \tag{7}\]
• спреды между шагами: $spread = P_{ask} - P_{bid}$: Array($N_{messages}$);
• текущее время: $[T_s, T_{ns}]$;
• прошедшее время в эпизоде: $[T_s, T_{ns}]$;
• начальная средняя цена на старте эпизода: $P_{init}$;
• ценовой снос: $P_{mid} - P_{init}$;
• размер задачи исполнения: $Q_{task}$;
• исполненное к текущему моменту количество: $Q_{exec}$;
• дисбалансы уровня 1 между шагами: $Q_{ask} - Q_{bid}$: Array($N_{messages}$).

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

1. Цена «дальнего касания» (far touch) — лучшая цена на противоположной стороне книги: \[ P_{far\_touch} = \begin{cases} P_{bid}, & \text{если задача — продажа} \\ P_{ask}, & \text{если задача — покупка} \end{cases} \tag{8}\]
2. Средняя цена (раздел 5.2.2).
3. Цена «ближнего касания» (near touch) — лучшая цена на стороне книги, соответствующей задаче.
4. «Пассивная» цена, как определено в разделе 5.2.2.

5.2.4. Функция награды. Мы предлагаем функцию награды, объединяющую способность агента превзойти базовую стратегию средневзвешенной по объёму цены (VWAP) и способность использовать предсказанные ценовые тренды. Она состоит из двух частей — преимущества и сноса. Первая представляет преимущество, полученное над средней ценой исполнения, взвешенной по объёму, в течение шага. Вторая представляет эффект ценовых движений средней цены сделок по отношению к средней цене в начале эпизода. Параметр $\lambda$ может настраиваться для взвешивания эффекта этого сноса. В уравнениях (9) и (10) индексы $i$ и $j$ относятся соответственно к множеству сделок, исполненных в течение шага, и множеству сделок, исполненных RL-агентом.

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

5.3. Сравнение стоимости прогонов

Прежде чем рассматривать задачу RL в разделе 6, рассмотрим улучшения времени выполнения среды исполнения на основе JAX-LOB по сравнению с сопоставимыми gym-средами, работающими на CPU. Мы сравниваем нашу среду исполнения со средой RL4MM, спроектированной для маркетмейкинга, но с сопоставимой основной логикой. Все стаканы работают на данных от 2 января 2015 года по десяти уровням для акции TSLA. В средах на CPU и на JAX-LOB за шаг обрабатывается 100 сообщений. RL4MM измеряет длительность шага во времени, но мы обнаруживаем, что 100 сообщений обрабатываются примерно каждые 2 минуты.

В случае среды JAX-LOB мы запускаем разное число сред параллельно на GPU Nvidia 2080 Ti, тогда как среды RL4MM и на CPU запускаются последовательно на CPU Apple M1. Улучшение (таблица 6) составляет примерно фактор 10 по сравнению со средой RL4MM и фактор 5 по сравнению с нашей реализацией на CPU при 1000 параллельных сред. Дальнейшее увеличение числа сред возможно, особенно на GPU с большей памятью, но такая экономия на масштабе в конечном счёте ограничена памятью.

6. Использование среды gymnax для обучения с подкреплением

6.1. Цикл обучения

Мы переписываем рекуррентный PPO, чтобы допустить непрерывное пространство действий. Полная архитектура сетей актора и критика приведена на рисунке 2. Для каждого шага обновления сети собирается 10 шагов по 1000 сред, что даёт батч размера 10 000 переходов. Для каждой из четырёх эпох на шаг обновления этот батч подразделяется на четыре мини-батча, которые используются для вычисления функций потерь, взятия градиентов и обновления сети градиентным спуском с оптимизатором Adam, реализованным в optax.

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

6.2. Ускорение обучения благодаря JAX

Хотя у нас нет бенчмарков для точно такой же постановки задачи в других подходах, мы тем не менее сравниваем скорости обучения для другой среды, работающей на CPU. Это дополняет данные раздела 5.3, поскольку передача данных с CPU на GPU для обновления сети добавляет значительные накладные расходы ко времени выполнения. При работе gymnax на GPU этого не требуется, и мы видим дальнейшие улучшения на рисунке 3. Пакет RL4MM, который мы используем для сравнения в разделах 4.3 и 5.3, в этом сравнении опущен, так как нам не удалось воспроизвести справедливые экспериментальные условия из-за проблем, возникших при использовании GPU.

Чтобы уточнить сравнение, дадим очень краткий обзор стакана на CPU, который мы используем для сравнения. Он основан на симуляторе стакана с традиционной древовидной структурой и использует реализацию рекуррентного PPO из Stable-Baselines3 для решения задачи исполнения. Есть ряд ключевых переменных, общих для обоих экспериментов, которые мы стремимся контролировать настолько, насколько возможно, чтобы сравнение было равноправным:

• обрабатываемых сообщений на шаг — 100 сообщений;
• рассматриваемых уровней данных LOBSTER — 10 уровней;
• длительность эпизода — 30 минут;
• оборудование — GPU: Nvidia A40, CPU: 32-ядерный AMD EPYC 7513; в обоих экспериментах обновления сети выполняются на GPU;
• распараллеливание — 1000 параллельных сред для JAX-LOB, 1000 векторизованных сред для CPU, но собираемых последовательно;
• размер батча — 10 000, 10 шагов по 1000 сред;
• размер мини-батча — 2500, по 4 на батч;
• число эпох — 4 эпохи;
• данные — высокочастотные данные LOBSTER за апрель 2021 года по AMZN.

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

Мы наблюдаем 7-кратное ускорение цикла обучения на основе симулятора JAX-LOB по сравнению с использованием реализации стакана на CPU. Это больше, чем 5-кратное ускорение, приведённое в таблице 2. Мы приписываем это отсутствию передачи данных между вычислительными устройствами. Здесь ещё не учитывается подбор гиперпараметров, который дополнительно выиграл бы от распараллеливания и тривиально достижим в нашей полностью JAX-основанной реализации.

6.3. Обучение задаче исполнения

Мы начинаем обучать агента исполнения на предложенной постановке на данных одного дня (1 апреля 2021 года), выбранных столь узко, поскольку цель этой работы — прежде всего продемонстрировать функциональность представленного фреймворка. Кривая обучения со скользящим средним по окну в 30 шагов показана на рисунке 4. Мы полагаем $\lambda = 0$, чтобы упростить обучение, так как часть награды, связанная со сносом, выучивается гораздо труднее, и сравниваем кривую обучения с бенчмарковой стратегией TWAP, которая состоит в линейном исполнении желаемого объёма в течение эпизода. Хотя показано, что стратегия TWAP превзойдена в обучении, мы не можем делать каких-либо заявлений об успешности этого агента без проведения полных тестов вне выборки. Тем не менее это указывает на перспективное направление исследований на основе представленного фреймворка.

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

7. Выводы и дальнейшая работа

Мы представляем первую реализацию симулятора лимитного стакана для GPU под названием JAX-LOB. Используя её как часть среды gymnax для обучения с подкреплением, мы получаем как минимум 5-кратное ускорение по сравнению с сопоставимой реализацией на CPU. При использовании среды для обучения RL-агента с PureJaxRL мы наблюдаем даже 7-кратное ускорение относительно нашей сопоставимой реализации на CPU. Ожидается, что это ускорение за счёт распараллеливания будет способствовать исследованиям в области применения RL к высокочастотной торговле и задачам исполнения, требующим реагирующего симулятора стакана.

В рамках этой работы мы приводим пример использования — обучение RL-агента задаче исполнения на данных сообщений за один день. Мы планируем расширить эту работу, следуя строгому конвейеру RL-обучения для задачи исполнения и других задач, включая тестирование вне выборки и подбор гиперпараметров. Ожидается, что последнее дополнительно выиграет от параллельной природы реализации JAX-LOB. Для обучения устойчивого и экономичного RL-агента исполнения потребуются детальные эксперименты в отношении пространства действий, пространства признаков, формирования награды и архитектуры сети.

Литература

• [1] Almgren, R., Chriss, N. Optimal execution of portfolio transactions. The Journal of Risk, 2001, 3(2), 5–39.
• [2] Amrouni, S., Moulin, A., Vann, J., Vyetrenko, S., Balch, T., Veloso, M. ABIDES-gym: gym environments for multi-agent discrete event simulation and application to financial markets. ICAIF '21, 2022, 1–9.
• [3] Anonymous Author(s). Generative AI for End-to-End Limit Order Book Modelling (готовится к публикации, 2023).
• [4] Babuschkin, I. и др. The DeepMind JAX Ecosystem, 2020. github.com/deepmind
• [5] Belcak, P., Calliess, J.-P., Zohren, S. Fast Agent-Based Simulation Framework with Applications to Reinforcement Learning and the Study of Trading Latency Effects. Multi-Agent-Based Simulation XXII, Springer, 2022, 42–56.
• [6] Bertsimas, D., Lo, A. W. Optimal control of execution costs. Journal of Financial Markets, 1998, 1(1), 1–50.
• [7] Beysolow II, T. Market making via reinforcement learning. Applied Reinforcement Learning with Python, 2019, 77–94.
• [8] Bouchaud, J.-P., Bonart, J., Donier, J., Gould, M. Trades, Quotes and Prices: Financial Markets Under the Microscope. Cambridge University Press, 2018.
• [9] Bradbury, J. и др. JAX: composable transformations of Python+NumPy programs, 2018. github.com/google/jax
• [10] Byrd, D., Hybinette, M., Balch, T. H. ABIDES: Towards High-Fidelity Multi-Agent Market Simulation. SIGSIM-PADS '20, 2020, 11–22.
• [11] Casgrain, P. LimitOrderBook.jl, 2020. github.com/p-casgrain/LimitOrderBook.jl
• [12] Cont, R., Stoikov, S., Talreja, R. A Stochastic Model for Order Book Dynamics. Operations Research, 2010, 58(3), 549–563.
• [13] Dabérius, K., Granat, E., Karlsson, P. Deep execution-value and policy based reinforcement learning for trading and beating market benchmarks. SSRN 3374766, 2019.
• [14] Fang, J., Weng, J., Xiang, Y., Zhang, X. Imitate then Transcend: Multi-Agent Optimal Execution with Dual-Window Denoise PPO. arXiv:2206.10736, 2022.
• [15] Ganesh, S., Vadori, N., Xu, M., Zheng, H., Reddy, P., Veloso, M. Reinforcement learning for market making in a multi-agent dealer market. arXiv:1911.05892, 2019.
• [16] Google. XLA: Optimizing Compiler for machine learning. tensorflow.org/xla
• [17] Gould, M. D., Porter, M. A., Williams, S., McDonald, M., Fenn, D. J., Howison, S. D. Limit order books. Quantitative Finance, 2013, 13(11), 1709–1742.
• [18] Huang, R., Polak, T. LOBSTER: Limit Order Book Reconstruction System, 2011.
• [19] Jerome, J., Palmer, G., Savani, R. Market Making with Scaled Beta Policies. ICAIF '22, 2022, 214–222.
• [20] Jerome, J., Sanchez-Betancourt, L., Savani, R., Herdegen, M. Model-based gym environments for limit order book trading. arXiv:2209.07823, 2022.
• [21] Karpe, M., Fang, J., Ma, Z., Wang, C. Multi-agent reinforcement learning in a realistic limit order book market simulation. ICAIF '20, 2020.
• [22] Kingma, D., Ba, J. Adam: A Method for Stochastic Optimization. ICLR, 2015.
• [23] Kochkov, D., Smith, J. A., Alieva, A., Wang, Q., Brenner, M. P., Hoyer, S. Machine learning–accelerated computational fluid dynamics. PNAS, 2021, 118(21), e2101784118.
• [24] Kyle, A. S. Continuous Auctions and Insider Trading. Econometrica, 1985, 53(6), 1315–1335.
• [25] Lange, R. T. gymnax: A JAX-based Reinforcement Learning Environment Library, 2022. github.com/RobertTLange/gymnax
• [26] Lu, C., Kuba, J., Letcher, A., Metz, L., Schroeder de Witt, C., Foerster, J. Discovered policy optimisation. NeurIPS, 2022, 35, 16455–16468.
• [27] Nevmyvaka, Y., Feng, Y., Kearns, M. Reinforcement learning for optimized trade execution. ICML, 2006, 673–680.
• [28] Ning, B., Lin, F. H. T., Jaimungal, S. Double deep Q-learning for optimal execution. Applied Mathematical Finance, 2021, 28(4), 361–380.
• [29] Obizhaeva, A. A., Wang, J. Optimal trading strategy and supply/demand dynamics. Journal of Financial Markets, 2013, 16(1), 1–32.
• [30] Paulin, J. Understanding flash crash contagion and systemic risk: a calibrated agent-based approach, 2019.
• [31] Quera-Bofarull, A., Dyer, J., Calinescu, A., Wooldridge, M. Some challenges of calibrating differentiable agent-based models. arXiv:2307.01085, 2023.
• [32] Schoenholz, S. S., Cubuk, E. D. JAX M.D. A Framework for Differentiable Physics. NeurIPS, 2020, т. 33.
• [33] Scholl, M. P., Calinescu, A., Farmer, J. D. How market ecology explains market malfunction. PNAS, 2021, 118(26), e2015574118.
• [34] Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O. Proximal Policy Optimization Algorithms. arXiv:1707.06347, 2017.
• [35] Van Hasselt, H., Guez, A., Silver, D. Deep reinforcement learning with double Q-learning. AAAI, 2016, т. 30.
• [36] Vyetrenko, S., Byrd, D., Petosa, N., Mahfouz, M., Dervovic, D., Veloso, M., Balch, T. Get real: realism metrics for robust limit order book market simulations. ICAIF '20, 2020, 1–8.

Оригинал статьи: Frey, S. и др., «JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading», arXiv:2308.13289 · лицензия CC BY 4.0