Привет, Хабр!
Ape-X представляет собой подход к обучению с подкреплением, разработанный для использования в масштабируемых распределенных системах.
Основная идея Ape-X заключается в разделении ролей на акторов, которые взаимодействуют с окружением и собирают данные, и учеников, которые используют эти данные для обучения модели. Такое разделение позволяет ускорить процесс обучения и предотвратить заучивание субоптимальных политик.
Архитектура Ape-X
Немного про акторов и учеников
Акторы являются агентами, которые взаимодействуют со своими экземплярами среды. Каждый актор выбирает действия, основываясь на общей нейронной сети, которая используется для предсказания действий. Акторы аккумулируют свой опыт в общей памяти повторного воспроизведения. Этот опыт включает состояния, действия, вознаграждения и следующие состояния, полученные в результате взаимодействия со средой.
Ученики занимаются обучением модели. Они извлекают образцы опыта из общей памяти воспроизведения и обновляют параметры нейронной сети на основе этих данных. За счёт того, что обучение происходит вне среды, ученики могут параллельно работать с большим количеством данных.
Механизмы архитектуры
Приоритетная буферизация опыта (в дальнейшем PER) была предложена для решения проблемы равномерного распределения опыта в стандартных методах опыта с повторным воспроизведением. В классическом опыте с повторным воспроизведением все записи имеют одинаковую вероятность быть выбраны для обучения, что может привести к медленному обучению и субоптимальным результатам. PER решает эту проблему, присваивая приоритеты каждому опыту на основе их важности.
Если коротко, то PER работает так:
Оценка ошибки TD:
Каждому опыту присваивается приоритет на основе ошибки TD.
Опыт с большей ошибкой TD получает более высокий приоритет, т.к он содержит больше информации для обучения модели.
Обновление приоритетов:
При каждом добавлении нового опыта или обновлении модели, приоритеты пересчитываются.
Это позволяет системе адаптироваться к новым данным и изменяющимся условиям среды.
Выбор опыта для обучения:
Для выборки опытов используется стохастический процесс, основанный на приоритетах. Опыты с более высоким приоритетом имеют большую вероятность быть выбраны.
Тем не менее, чтобы избежать переобучения на небольшом наборе данных, в процесс добавляется некоторая степень случайности.
Коррекция смещения:
Использование приоритетной выборки может ввести смещение в процесс обучения. Для его коррекции используется важностное взвешивание, что позволяет уменьшить вероятность ошибки из-за неравномерного выбора данных.
PER позволяет быстрее выявлять и корректировать ошибки модели, ускоряя процесс сходимости. В Ape-X это мастхев, т.к множество параллельных акторов генерируют большие объемы данных, и важно по максимуму использовать эту информацию.
Реализация PER с помощью numpy
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. Записаться
