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

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. Записаться

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