Создание масштабируемых RL систем с Ape-X

Привет, Хабр!

Ape-X представляет собой подход к обучению с подкреплением, разработанный для использования в масштабируемых распределенных системах.

Основная идея Ape-X заключается в разделении ролей на акторов, которые взаимодействуют с окружением и собирают данные, и учеников, которые используют эти данные для обучения модели. Такое разделение позволяет ускорить процесс обучения и предотвратить заучивание субоптимальных политик.

Архитектура Ape-X

Немного про акторов и учеников

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

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

Механизмы архитектуры

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

Если коротко, то PER работает так:

1. Оценка ошибки TD:

   • Каждому опыту присваивается приоритет на основе стохастический процесс, основанный на приоритетах. Опыты с более высоким приоритетом имеют большую вероятность быть выбраны.

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

2. Коррекция смещения:

   • Использование приоритетной выборки может ввести смещение в процесс обучения. Для его коррекции используется важностное взвешивание, что позволяет уменьшить вероятность ошибки из-за неравномерного выбора данных.

PER позволяет быстрее выявлять и корректировать ошибки модели, ускоряя процесс сходимости. В Ape-X это мастхев, т.к множество параллельных акторов генерируют большие объемы данных, и важно по максимуму использовать эту информацию.

import numpy as np  class PrioritizedReplayBuffer:     def __init__(self, capacity, alpha=0.6):         self.capacity = capacity         self.buffer = []         self.priorities = np.zeros((capacity,), dtype=np.float32)         self.pos = 0         self.alpha = alpha      def add(self, experience):         max_prio = self.priorities.max() if self.buffer else 1.0         if len(self.buffer) < self.capacity:             self.buffer.append(experience)         else:             self.buffer[self.pos] = experience         self.priorities[self.pos] = max_prio         self.pos = (self.pos + 1) % self.capacity      def sample(self, batch_size, beta=0.4):         if len(self.buffer) == self.capacity:             prios = self.priorities         else:             prios = self.priorities[:self.pos]         probs = prios ** self.alpha         probs /= probs.sum()          indices = np.random.choice(len(self.buffer), batch_size, p=probs)         samples = [self.buffer[idx] for idx in indices]          total = len(self.buffer)         weights = (total * probs[indices]) ** (-beta)         weights /= weights.max()          batch = zip(*samples)         return batch, indices, weights      def update_priorities(self, batch_indices, batch_priorities):         for idx, prio in zip(batch_indices, batch_priorities):             self.priorities[idx] = prio 

Идем к другим не менее важным механизмам.

Централизованная память воспроизведения:
Все акторы отправляют свой опыт в централизованную память воспроизведения. Это позволяет ученикам иметь доступ к большему количеству данных и использовать их для более качественного обучения модели. Коммуникация между акторами и памятью воспроизведения осуществляется пакетами.

Обучение вне политики:
Ape-X использует обучение вне политик. Каждый актор может иметь свою стратегию исследования среды, что расширяет разнообразие получаемого опыта.

Реализация Ape-X

Будем использовать PyTorch и Ray. PyTorch имееь средства для создания и обучения нейронных сетей, а Ray помогает управлять распределенными вычислениями и координировать работу множества акторов и учеников. Пошагово настроим окружения для Ape-X, используя AWS EC2 или локальные машины.

Будем использовать Python 3.8 и более новые версии библиотек PyTorch и Ray. Для начала, создадим виртуальное окружение:

python3 -m venv ape-x-env source ape-x-env/bin/activate

Далее, установим необходимые библиотеки:

pip install torch torchvision pip install ray[default] pip install gym pip install opencv-python pip install tensorboardX

Для масштабирования системы используем AWS EC2. Создаем несколько инстансов, которые будут работать в распределенной среде. Для этого заходим в консоль управления AWS, выбираем EC2 и создаем новый инстанс с нужными характеристиками (например, t2.large для акторов и t2.xlarge для учеников).

После создания инстансов, подключаемся к ним через SSH:

ssh -i "key.pem" ec2-user@your-ec2-instance.amazonaws.com

Устанавливаем необходимые пакеты на каждом инстансе:

sudo yum update -y sudo yum install python3 -y python3 -m venv ape-x-env source ape-x-env/bin/activate pip install torch torchvision ray[default] gym opencv-python tensorboardX

Окружение готово и мы готовы создавать конфигурационные файлы и запускать Ape-X. Начнем с создания файла config.py, который будет содержать основные параметры:

import ray from ray import tune  def get_config():     config = {         "env": "CartPole-v1",         "num_workers": 8,         "num_gpus": 1,         "framework": "torch",         "lr": 1e-4,         "train_batch_size": 500,         "rollout_fragment_length": 50,     }     return config

Затем создадим файл train.py для запуска процесса обучения:

import ray from ray.rllib.agents.dqn import ApexTrainer from config import get_config  if __name__ == "__main__":     ray.init(address="auto")  # Connect to the Ray cluster     config = get_config()     trainer = ApexTrainer(config=config)          for i in range(1000):         result = trainer.train()         print(f"Iteration: {i}, reward: {result['episode_reward_mean']}")         if i % 10 == 0:             checkpoint = trainer.save()             print(f"Checkpoint saved at {checkpoint}")

Для запуска на одном узле достаточно выполнить команду:

python train.py

Для многоузловой конфигурации нужно запустить Ray на каждом узле:

ray start --head --port=6379 # На всех других узлах: ray start --address='IP_HEAD_NODE:6379'

После этого можно запустить train.py, и все узлы будут участвовать в обучении.

С Packer можно автоматизировать создание образов машин. Создадим файл конфигурации packer.json:

{   "builders": [{     "type": "amazon-ebs",     "region": "us-west-2",     "source_ami": "ami-0c55b159cbfafe1f0",     "instance_type": "t2.micro",     "ssh_username": "ec2-user",     "ami_name": "ape-x-ami-{{timestamp}}"   }],   "provisioners": [{     "type": "shell",     "inline": [       "sudo yum update -y",       "sudo yum install python3 -y",       "python3 -m venv ape-x-env",       "source ape-x-env/bin/activate",       "pip install torch torchvision ray[default] gym opencv-python tensorboardX"     ]   }] }

Запускаем Packer для создания образа:

packer build packer.json

Теперь есть готовый образ, который можно использовать для создания новых инстансов EC2 с уже установленными необходимыми библиотеками.

Для запуска Ape-X на одной машине можно юзать этот скрипт:

import gym import numpy as np import torch from torch import nn, optim import torch.nn.functional as F from apex import ApeX, create_env  # define the neural network model class DQN(nn.Module):     def __init__(self, obs_space, action_space):         super(DQN, self).__init__()         self.fc1 = nn.Linear(obs_space, 128)         self.fc2 = nn.Linear(128, 128)         self.out = nn.Linear(128, action_space)      def forward(self, x):         x = F.relu(self.fc1(x))         x = F.relu(self.fc2(x))         return self.out(x)  # initialize environment and model env = gym.make('CartPole-v0') obs_space = env.observation_space.shape[0] action_space = env.action_space.n model = DQN(obs_space, action_space) optimizer = optim.Adam(model.parameters(), lr=0.001)  # initialize Ape-X apex = ApeX(env, model, optimizer)  # training loop for episode in range(1000):     state = env.reset()     done = False     while not done:         action = apex.select_action(state)         next_state, reward, done, _ = env.step(action)         apex.store_experience(state, action, reward, next_state, done)         apex.learn()         state = next_state      if episode % 10 == 0:         print(f"Episode {episode}, Score: {apex.get_score()}")

Для многоузловой конфигурации необходимо настроить несколько экземпляров и настроить взаимодействие между ними.

import ray from ray import tune from apex import ApeX  ray.init()  # define training function def train_apex(config):     env = create_env(config["env"])     model = DQN(env.observation_space.shape[0], env.action_space.n)     optimizer = optim.Adam(model.parameters(), lr=config["lr"])     apex = ApeX(env, model, optimizer)      for episode in range(config["num_episodes"]):         state = env.reset()         done = False         while not done:             action = apex.select_action(state)             next_state, reward, done, _ = env.step(action)             apex.store_experience(state, action, reward, next_state, done)             apex.learn()             state = next_state          if episode % 10 == 0:             tune.report(score=apex.get_score())  # run the distributed training tune.run(     train_apex,     config={         "env": "CartPole-v0",         "lr": 0.001,         "num_episodes": 1000,     } )

Теперь мы получили полностью настроенное и работоспособное окружение для реализации Ape-X.

В завершение предлагаю посетить открытые уроки по Reinforcement Learning от OTUS:

• 13 июня: Применение нейросетевых моделей в обучении с подкреплением. Записаться

• 17 июня: Построение торгового агента на базе фремворка FinRL. Записаться