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

Kate

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

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

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

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

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

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

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

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

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

  1. Оценка ошибки TD:
    • Каждому опыту присваивается приоритет на основе ошибки TD.
    • Опыт с большей ошибкой TD получает более высокий приоритет, т.к он содержит больше информации для обучения модели.
  2. Обновление приоритетов:
    • При каждом добавлении нового опыта или обновлении модели, приоритеты пересчитываются.
    • Это позволяет системе адаптироваться к новым данным и изменяющимся условиям среды.
  3. Выбор опыта для обучения:
    • Для выборки опытов используется стохастический процесс, основанный на приоритетах. Опыты с более высоким приоритетом имеют большую вероятность быть выбраны.
    • Тем не менее, чтобы избежать переобучения на небольшом наборе данных, в процесс добавляется некоторая степень случайности.
  4. Коррекция смещения:
    • Использование приоритетной выборки может ввести смещение в процесс обучения. Для его коррекции используется важностное взвешивание, что позволяет уменьшить вероятность ошибки из-за неравномерного выбора данных.
PER позволяет быстрее выявлять и корректировать ошибки модели, ускоряя процесс сходимости. В Ape-X это мастхев, т.к множество параллельных акторов генерируют большие объемы данных, и важно по максимуму использовать эту информацию.

Реализация PER с помощью numpy









































Идем к другим не менее важным механизмам.
Централизованная память воспроизведения:
Все акторы отправляют свой опыт в централизованную память воспроизведения. Это позволяет ученикам иметь доступ к большему количеству данных и использовать их для более качественного обучения модели. Коммуникация между акторами и памятью воспроизведения осуществляется пакетами.
Обучение вне политики:
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.

 
Сверху