Deep RL Course, юнит 8, практика: PPO с нуля и Doom

Перевод двух практических занятий восьмого юнита курса Hugging Face: PPO с нуля на PyTorch по образцу CleanRL (полный ppo.py с переведёнными комментариями) и асинхронный PPO в Sample Factory для ViZDoom.

Это перевод практических занятий юнита 8 курса Hugging Face Deep RL Course: первой части (units/en/unit8/hands-on-cleanrl.mdx, ноутбук unit8_part1.ipynb) и второй (hands-on-sf.mdx, ноутбук unit8_part2.ipynb). Теория PPO — на странице «Юнит 8: PPO», карта курса — в путеводителе.

В уроке 6.0.3 уже есть свой PPO на 70 строк для CartPole. Код CleanRL ниже делает то же самое «по-взрослому»: параллельные среды, отжиг скорости обучения, обрезка value loss, approx_kl, clipfrac и explained_variance в TensorBoard. Полезно положить две реализации рядом и найти в обеих GAE, клиппинг и полный loss. Вторая часть (Sample Factory и Doom) для языковых моделей не нужна, но хорошо показывает, как PPO масштабируют на много процессов, — та же идея разделения «сборщиков опыта» и «обучателя» есть в системах RL для LLM (verl, урок 6.4.5).

Ноутбук первой части написан для старого gym 0.22 (env.seed, reset() возвращает только наблюдение, step() — четыре значения) и среды LunarLander-v2. В свежих Gymnasium этот код не запустится без правок. Проще взять актуальный ppo.py из репозитория CleanRL: он уже на Gymnasium, аргументы разбирает библиотека tyro, а алгоритм тот же. Запуск, например: python ppo.py --env-id CartPole-v1 --total-timesteps 50000.
По умолчанию в коде курса 50 000 шагов: на CartPole это минуты на CPU. Для LunarLander шагов нужно на порядок больше — увеличьте --total-timesteps.
Sample Factory во второй части работает только в Linux и macOS. Обучение на 4 млн шагов, по оценке автора, — около 20 минут; Deathmatch обучать долго, его достаточно скачать готовым.
Лидерборд курса выключен, публикация на Hub необязательна.

Часть 1. PPO с нуля по образцу CleanRL

Теперь, когда мы изучили теорию PPO, лучший способ понять, как он работает, — реализовать его с нуля.

Реализовать архитектуру с нуля — лучший способ её понять, и это хорошая привычка. Мы уже делали так для метода на основе ценности (Q-learning) и для метода на основе политики (REINFORCE).

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

руководством Косты Хуана. Коста — автор CleanRL, библиотеки глубокого обучения с подкреплением, в которой каждая реализация — один качественный файл с удобными для исследований возможностями;
в дополнение к руководству, чтобы копнуть глубже, можно прочитать 13 ключевых деталей реализации: The 37 Implementation Details of PPO.

Затем, чтобы проверить надёжность, обучим агента в среде:

LunarLander-v2.

В оригинале здесь видео: лунный модуль садится между флажками.

И наконец, выложим обученную модель на Hub, чтобы оценить агента и посмотреть, как он играет.

LunarLander-v2 — первая среда, с которой вы начинали этот курс. Тогда вы не знали, как она устроена, а теперь можете написать агента с нуля и обучить его. Разве не потрясающе 🤩.

Ноутбук Colab: unit8_part1.ipynb.

Юнит 8: Proximal Policy Gradient (PPO) на PyTorch 🤖

В этом ноутбуке вы научитесь писать агента PPO с нуля на PyTorch, взяв за образец реализацию CleanRL.

Чтобы проверить надёжность, обучим агента в среде:

LunarLander-v2 🚀.

Мы постоянно улучшаем наши руководства, поэтому, если найдёте проблемы в этом ноутбуке, откройте issue в репозитории на GitHub.

Цели ноутбука 🏆. К концу ноутбука вы сможете:

написать агента PPO с нуля на PyTorch;
выложить обученного агента и код на Hub с красивым видео игры и оценкой качества 🔥.

Что нужно заранее 🏗️: 🔲 📚 изучить PPO по юниту 8 🤗

Чтобы занятие засчиталось для сертификата, нужно выложить одну модель. Минимального результата мы не требуем, но советуем попробовать разные гиперпараметры, чтобы получить результат получше. Если модели не видно, прокрутите страницу вниз и нажмите кнопку обновления.

Включаем GPU 💪

Чтобы ускорить обучение агента, используем GPU. Для этого откройте Runtime > Change Runtime type и выберите Hardware Accelerator > GPU.

Создаём виртуальный дисплей 🔽

Во время работы нам понадобится записать видео игры. В Colab для этого нужен виртуальный экран, чтобы отрисовывать среду (и записывать кадры). Следующая ячейка устанавливает библиотеки, создаёт и запускает виртуальный экран 🖥

apt install python-opengl
apt install ffmpeg
apt install xvfb
pip install pyglet==1.5
pip install pyvirtualdisplay

# Виртуальный дисплей
from pyvirtualdisplay import Display

virtual_display = Display(visible=0, size=(1400, 900))
virtual_display.start()

Устанавливаем зависимости 🔽

В этом упражнении мы используем gym==0.21, потому что видео записывалось с Gym.

pip install gym==0.22
pip install imageio-ffmpeg
pip install huggingface_hub
pip install gym[box2d]==0.22

(Так в оригинале: в тексте упомянут gym==0.21, а ставится 0.22.)

Пишем PPO с нуля по руководству Косты Хуана

Основу реализации PPO возьмём из отличного руководства Косты Хуана.
В дополнение к руководству, чтобы копнуть глубже, можно прочитать 37 ключевых деталей реализации: iclr-blog-track.github.io/2022/03/25/ppo-implementation-details.

👉 Видеоруководство: youtu.be/MEt6rrxH8W4 (на английском). В ноутбуке оно встроено ячейкой с IPython.display.HTML.

Добавляем интеграцию с Hugging Face 🤗

Чтобы отправить модель на Hub, нужно определить функцию package_to_hub.
Добавим зависимости, нужные для отправки модели на Hub.
Добавим в функцию parse_args() новый аргумент — id репозитория, куда отправлять модель.
Затем добавим методы, нужные для отправки модели на Hub. Они:
- _evaluate_agent() — оценивают агента;
- _generate_model_card() — создают карточку модели агента;
- _record_video() — записывают видео игры агента.
Наконец, вызовем эту функцию в конце обучения PPO.

Вот как выглядит итоговый файл ppo.py. В оригинале сначала отдельно показаны функции интеграции с Hub, а потом весь файл, где они повторяются; здесь файл приведён один раз, в трёх частях. Склейте их по порядку в один файл ppo.py.

ppo.py, часть 1: импорты и аргументы

# Документация и результаты экспериментов: https://docs.cleanrl.dev/rl-algorithms/ppo/#ppopy

import argparse
import os
import random
import time
from distutils.util import strtobool

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical
from torch.utils.tensorboard import SummaryWriter

from huggingface_hub import HfApi, upload_folder
from huggingface_hub.repocard import metadata_eval_result, metadata_save

from pathlib import Path
import datetime
import tempfile
import json
import shutil
import imageio

from wasabi import Printer

msg = Printer()


def parse_args():
    # fmt: off
    parser = argparse.ArgumentParser()
    parser.add_argument("--exp-name", type=str, default=os.path.basename(__file__).rstrip(".py"),
        help="название эксперимента")
    parser.add_argument("--seed", type=int, default=1,
        help="seed эксперимента")
    parser.add_argument("--torch-deterministic", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="если выключить, `torch.backends.cudnn.deterministic=False`")
    parser.add_argument("--cuda", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="если включено, по умолчанию используется cuda")
    parser.add_argument("--track", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
        help="если включено, эксперимент отслеживается в Weights and Biases")
    parser.add_argument("--wandb-project-name", type=str, default="cleanRL",
        help="название проекта в wandb")
    parser.add_argument("--wandb-entity", type=str, default=None,
        help="сущность (команда) проекта в wandb")
    parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
        help="записывать ли видео игры агента (смотрите папку `videos`)")

    # Аргументы, относящиеся к алгоритму
    parser.add_argument("--env-id", type=str, default="CartPole-v1",
        help="id среды")
    parser.add_argument("--total-timesteps", type=int, default=50000,
        help="общее число шагов эксперимента")
    parser.add_argument("--learning-rate", type=float, default=2.5e-4,
        help="скорость обучения оптимизатора")
    parser.add_argument("--num-envs", type=int, default=4,
        help="число параллельных игровых сред")
    parser.add_argument("--num-steps", type=int, default=128,
        help="сколько шагов делать в каждой среде за один сбор опыта политикой")
    parser.add_argument("--anneal-lr", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="включить отжиг (линейное уменьшение) скорости обучения сетей политики и ценности")
    parser.add_argument("--gae", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="считать преимущество через GAE")
    parser.add_argument("--gamma", type=float, default=0.99,
        help="коэффициент дисконтирования gamma")
    parser.add_argument("--gae-lambda", type=float, default=0.95,
        help="lambda для обобщённой оценки преимущества (GAE)")
    parser.add_argument("--num-minibatches", type=int, default=4,
        help="число мини-батчей")
    parser.add_argument("--update-epochs", type=int, default=4,
        help="K эпох обновления политики")
    parser.add_argument("--norm-adv", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="включить нормализацию преимуществ")
    parser.add_argument("--clip-coef", type=float, default=0.2,
        help="коэффициент обрезки суррогатной цели")
    parser.add_argument("--clip-vloss", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
        help="использовать ли обрезанный loss функции ценности, как в статье")
    parser.add_argument("--ent-coef", type=float, default=0.01,
        help="коэффициент энтропии")
    parser.add_argument("--vf-coef", type=float, default=0.5,
        help="коэффициент функции ценности")
    parser.add_argument("--max-grad-norm", type=float, default=0.5,
        help="максимальная норма для обрезки градиента")
    parser.add_argument("--target-kl", type=float, default=None,
        help="целевой порог KL-дивергенции")

    # Добавляем аргумент для Hugging Face
    parser.add_argument("--repo-id", type=str, default="ThomasSimonini/ppo-CartPole-v1", help="id репозитория модели на Hugging Face Hub {username/repo_name}")

    args = parser.parse_args()
    args.batch_size = int(args.num_envs * args.num_steps)
    args.minibatch_size = int(args.batch_size // args.num_minibatches)
    # fmt: on
    return args

ppo.py, часть 2: отправка на Hub

def package_to_hub(
    repo_id,
    model,
    hyperparameters,
    eval_env,
    video_fps=30,
    commit_message="Push agent to the Hub",
    token=None,
    logs=None,
):
    """
    Оценить модель, записать видео и загрузить модель на Hugging Face Hub.
    Метод выполняет весь конвейер:
    - оценивает модель
    - создаёт карточку модели
    - записывает видео игры агента
    - отправляет всё на Hub
    :param repo_id: id репозитория модели на Hugging Face Hub
    :param model: обученная модель
    :param eval_env: среда для оценки агента
    :param fps: сколько кадров в секунду в видео
    :param commit_message: сообщение коммита
    :param logs: локальная папка с логами TensorBoard, которые вы хотите загрузить
    """
    msg.info(
        "This function will save, evaluate, generate a video of your agent, "
        "create a model card and push everything to the hub. "
        "It might take up to 1min. \n "
        "This is a work in progress: if you encounter a bug, please open an issue."
    )
    # Шаг 1: клонируем или создаём репозиторий
    repo_url = HfApi().create_repo(
        repo_id=repo_id,
        token=token,
        private=False,
        exist_ok=True,
    )

    with tempfile.TemporaryDirectory() as tmpdirname:
        tmpdirname = Path(tmpdirname)

        # Шаг 2: сохраняем модель
        torch.save(model.state_dict(), tmpdirname / "model.pt")

        # Шаг 3: оцениваем модель и собираем JSON
        mean_reward, std_reward = _evaluate_agent(eval_env, 10, model)

        # Сначала получаем дату и время
        eval_datetime = datetime.datetime.now()
        eval_form_datetime = eval_datetime.isoformat()

        evaluate_data = {
            "env_id": hyperparameters.env_id,
            "mean_reward": mean_reward,
            "std_reward": std_reward,
            "n_evaluation_episodes": 10,
            "eval_datetime": eval_form_datetime,
        }

        # Пишем файл JSON
        with open(tmpdirname / "results.json", "w") as outfile:
            json.dump(evaluate_data, outfile)

        # Шаг 4: записываем видео
        video_path = tmpdirname / "replay.mp4"
        record_video(eval_env, model, video_path, video_fps)

        # Шаг 5: создаём карточку модели
        generated_model_card, metadata = _generate_model_card(
            "PPO", hyperparameters.env_id, mean_reward, std_reward, hyperparameters
        )
        _save_model_card(tmpdirname, generated_model_card, metadata)

        # Шаг 6: добавляем логи, если нужно
        if logs:
            _add_logdir(tmpdirname, Path(logs))

        msg.info(f"Pushing repo {repo_id} to the Hugging Face Hub")

        repo_url = upload_folder(
            repo_id=repo_id,
            folder_path=tmpdirname,
            path_in_repo="",
            commit_message=commit_message,
            token=token,
        )

        msg.info(f"Your model is pushed to the Hub. You can view your model here: {repo_url}")
    return repo_url


def _evaluate_agent(env, n_eval_episodes, policy):
    """
    Оценить агента на ``n_eval_episodes`` эпизодах и вернуть среднюю награду и её стандартное отклонение.
    :param env: среда для оценки
    :param n_eval_episodes: число эпизодов для оценки агента
    :param policy: агент
    """
    episode_rewards = []
    for episode in range(n_eval_episodes):
        state = env.reset()
        step = 0
        done = False
        total_rewards_ep = 0

        while done is False:
            state = torch.Tensor(state).to(device)
            action, _, _, _ = policy.get_action_and_value(state)
            new_state, reward, done, info = env.step(action.cpu().numpy())
            total_rewards_ep += reward
            if done:
                break
            state = new_state
        episode_rewards.append(total_rewards_ep)
    mean_reward = np.mean(episode_rewards)
    std_reward = np.std(episode_rewards)

    return mean_reward, std_reward


def record_video(env, policy, out_directory, fps=30):
    images = []
    done = False
    state = env.reset()
    img = env.render(mode="rgb_array")
    images.append(img)
    while not done:
        state = torch.Tensor(state).to(device)
        # Выбираем действие политикой
        # (комментарий оригинала остался от Q-learning: «действие с наибольшей ожидаемой будущей наградой»)
        action, _, _, _ = policy.get_action_and_value(state)
        state, reward, done, info = env.step(
            action.cpu().numpy()
        )  # сразу пишем next_state в state — для записи этого достаточно
        img = env.render(mode="rgb_array")
        images.append(img)
    imageio.mimsave(out_directory, [np.array(img) for i, img in enumerate(images)], fps=fps)


def _generate_model_card(model_name, env_id, mean_reward, std_reward, hyperparameters):
    """
    Создать карточку модели для Hub
    :param model_name: название модели
    :env_id: название среды
    :mean_reward: средняя награда агента
    :std_reward: стандартное отклонение награды агента
    :hyperparameters: аргументы обучения
    """
    # Шаг 1: выбираем теги
    metadata = generate_metadata(model_name, env_id, mean_reward, std_reward)

    # Превращаем пространство имён гиперпараметров в строку
    converted_dict = vars(hyperparameters)
    converted_str = str(converted_dict)
    converted_str = converted_str.split(", ")
    converted_str = "\n".join(converted_str)

    # Шаг 2: создаём карточку модели
    model_card = f"""
  # PPO Agent Playing {env_id}

  This is a trained model of a PPO agent playing {env_id}.

  # Hyperparameters
  """
    return model_card, metadata


def generate_metadata(model_name, env_id, mean_reward, std_reward):
    """
    Задать теги для карточки модели
    :param model_name: название модели
    :param env_id: название среды
    :mean_reward: средняя награда агента
    :std_reward: стандартное отклонение награды агента
    """
    metadata = {}
    metadata["tags"] = [
        env_id,
        "ppo",
        "deep-reinforcement-learning",
        "reinforcement-learning",
        "custom-implementation",
        "deep-rl-course",
    ]

    # Добавляем метрики
    eval = metadata_eval_result(
        model_pretty_name=model_name,
        task_pretty_name="reinforcement-learning",
        task_id="reinforcement-learning",
        metrics_pretty_name="mean_reward",
        metrics_id="mean_reward",
        metrics_value=f"{mean_reward:.2f} +/- {std_reward:.2f}",
        dataset_pretty_name=env_id,
        dataset_id=env_id,
    )

    # Объединяем оба словаря
    metadata = {**metadata, **eval}

    return metadata


def _save_model_card(local_path, generated_model_card, metadata):
    """Сохранить карточку модели в репозиторий.
    :param local_path: папка репозитория
    :param generated_model_card: карточка, созданная _generate_model_card()
    :param metadata: метаданные
    """
    readme_path = local_path / "README.md"
    readme = ""
    if readme_path.exists():
        with readme_path.open("r", encoding="utf8") as f:
            readme = f.read()
    else:
        readme = generated_model_card

    with readme_path.open("w", encoding="utf-8") as f:
        f.write(readme)

    # Сохраняем метрики в метаданные README
    metadata_save(readme_path, metadata)


def _add_logdir(local_path: Path, logdir: Path):
    """Добавить папку логов в репозиторий.
    :param local_path: папка репозитория
    :param logdir: папка логов
    """
    if logdir.exists() and logdir.is_dir():
        # Кладём логи в репозиторий, в новую папку logs
        repo_logdir = local_path / "logs"

        # Удаляем старые логи, если они есть
        if repo_logdir.exists():
            shutil.rmtree(repo_logdir)

        # Копируем папку логов в репозиторий
        shutil.copytree(logdir, repo_logdir)

ppo.py, часть 3: агент и цикл обучения PPO

def make_env(env_id, seed, idx, capture_video, run_name):
    def thunk():
        env = gym.make(env_id)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        if capture_video:
            if idx == 0:
                env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
        env.seed(seed)
        env.action_space.seed(seed)
        env.observation_space.seed(seed)
        return env

    return thunk


def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer


class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0),
        )
        self.actor = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, envs.single_action_space.n), std=0.01),
        )

    def get_value(self, x):
        return self.critic(x)

    def get_action_and_value(self, x, action=None):
        logits = self.actor(x)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(x)


if __name__ == "__main__":
    args = parse_args()
    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    if args.track:
        import wandb

        wandb.init(
            project=args.wandb_project_name,
            entity=args.wandb_entity,
            sync_tensorboard=True,
            config=vars(args),
            name=run_name,
            monitor_gym=True,
            save_code=True,
        )
    writer = SummaryWriter(f"runs/{run_name}")
    writer.add_text(
        "hyperparameters",
        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
    )

    # ЛУЧШЕ НЕ МЕНЯТЬ: задаём seed
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = args.torch_deterministic

    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

    # Настройка сред
    envs = gym.vector.SyncVectorEnv(
        [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
    )
    assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"

    agent = Agent(envs).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)

    # ЛОГИКА АЛГОРИТМА: хранилище опыта
    obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape).to(device)
    actions = torch.zeros((args.num_steps, args.num_envs) + envs.single_action_space.shape).to(device)
    logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)
    rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)
    dones = torch.zeros((args.num_steps, args.num_envs)).to(device)
    values = torch.zeros((args.num_steps, args.num_envs)).to(device)

    # ЛУЧШЕ НЕ МЕНЯТЬ: запускаем игру
    global_step = 0
    start_time = time.time()
    next_obs = torch.Tensor(envs.reset()).to(device)
    next_done = torch.zeros(args.num_envs).to(device)
    num_updates = args.total_timesteps // args.batch_size

    for update in range(1, num_updates + 1):
        # Отжиг скорости обучения, если он включён
        if args.anneal_lr:
            frac = 1.0 - (update - 1.0) / num_updates
            lrnow = frac * args.learning_rate
            optimizer.param_groups[0]["lr"] = lrnow

        for step in range(0, args.num_steps):
            global_step += 1 * args.num_envs
            obs[step] = next_obs
            dones[step] = next_done

            # ЛОГИКА АЛГОРИТМА: выбор действия
            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(next_obs)
                values[step] = value.flatten()
            actions[step] = action
            logprobs[step] = logprob

            # ЛУЧШЕ НЕ МЕНЯТЬ: делаем шаг в игре и пишем логи
            next_obs, reward, done, info = envs.step(action.cpu().numpy())
            rewards[step] = torch.tensor(reward).to(device).view(-1)
            next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)

            for item in info:
                if "episode" in item.keys():
                    print(f"global_step={global_step}, episodic_return={item['episode']['r']}")
                    writer.add_scalar("charts/episodic_return", item["episode"]["r"], global_step)
                    writer.add_scalar("charts/episodic_length", item["episode"]["l"], global_step)
                    break

        # бутстрэп ценности, если эпизод не закончен
        with torch.no_grad():
            next_value = agent.get_value(next_obs).reshape(1, -1)
            if args.gae:
                advantages = torch.zeros_like(rewards).to(device)
                lastgaelam = 0
                for t in reversed(range(args.num_steps)):
                    if t == args.num_steps - 1:
                        nextnonterminal = 1.0 - next_done
                        nextvalues = next_value
                    else:
                        nextnonterminal = 1.0 - dones[t + 1]
                        nextvalues = values[t + 1]
                    delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]
                    advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
                returns = advantages + values
            else:
                returns = torch.zeros_like(rewards).to(device)
                for t in reversed(range(args.num_steps)):
                    if t == args.num_steps - 1:
                        nextnonterminal = 1.0 - next_done
                        next_return = next_value
                    else:
                        nextnonterminal = 1.0 - dones[t + 1]
                        next_return = returns[t + 1]
                    returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
                advantages = returns - values

        # выпрямляем батч в одну ось
        b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)
        b_logprobs = logprobs.reshape(-1)
        b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
        b_advantages = advantages.reshape(-1)
        b_returns = returns.reshape(-1)
        b_values = values.reshape(-1)

        # Оптимизируем сети политики и ценности
        b_inds = np.arange(args.batch_size)
        clipfracs = []
        for epoch in range(args.update_epochs):
            np.random.shuffle(b_inds)
            for start in range(0, args.batch_size, args.minibatch_size):
                end = start + args.minibatch_size
                mb_inds = b_inds[start:end]

                _, newlogprob, entropy, newvalue = agent.get_action_and_value(
                    b_obs[mb_inds], b_actions.long()[mb_inds]
                )
                logratio = newlogprob - b_logprobs[mb_inds]
                ratio = logratio.exp()

                with torch.no_grad():
                    # считаем approx_kl: http://joschu.net/blog/kl-approx.html
                    old_approx_kl = (-logratio).mean()
                    approx_kl = ((ratio - 1) - logratio).mean()
                    clipfracs += [((ratio - 1.0).abs() > args.clip_coef).float().mean().item()]

                mb_advantages = b_advantages[mb_inds]
                if args.norm_adv:
                    mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)

                # Loss политики
                pg_loss1 = -mb_advantages * ratio
                pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
                pg_loss = torch.max(pg_loss1, pg_loss2).mean()

                # Loss ценности
                newvalue = newvalue.view(-1)
                if args.clip_vloss:
                    v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
                    v_clipped = b_values[mb_inds] + torch.clamp(
                        newvalue - b_values[mb_inds],
                        -args.clip_coef,
                        args.clip_coef,
                    )
                    v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
                    v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)
                    v_loss = 0.5 * v_loss_max.mean()
                else:
                    v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()

                entropy_loss = entropy.mean()
                loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef

                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
                optimizer.step()

            if args.target_kl is not None:
                if approx_kl > args.target_kl:
                    break

        y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()
        var_y = np.var(y_true)
        explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y

        # ЛУЧШЕ НЕ МЕНЯТЬ: пишем награды для графиков
        writer.add_scalar("charts/learning_rate", optimizer.param_groups[0]["lr"], global_step)
        writer.add_scalar("losses/value_loss", v_loss.item(), global_step)
        writer.add_scalar("losses/policy_loss", pg_loss.item(), global_step)
        writer.add_scalar("losses/entropy", entropy_loss.item(), global_step)
        writer.add_scalar("losses/old_approx_kl", old_approx_kl.item(), global_step)
        writer.add_scalar("losses/approx_kl", approx_kl.item(), global_step)
        writer.add_scalar("losses/clipfrac", np.mean(clipfracs), global_step)
        writer.add_scalar("losses/explained_variance", explained_var, global_step)
        print("SPS:", int(global_step / (time.time() - start_time)))
        writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)

    envs.close()
    writer.close()

    # Создаём среду для оценки
    eval_env = gym.make(args.env_id)

    package_to_hub(
        repo_id=args.repo_id,
        model=agent,  # модель, которую хотим сохранить
        hyperparameters=args,
        eval_env=gym.make(args.env_id),
        logs=f"runs/{run_name}",
    )

Что в теории	Где в коде
Актор $\pi_\theta$ и критик $V$ — две сети	класс `Agent`: `self.actor` выдаёт логиты действий, `self.critic` — одно число
Сбор опыта старой политикой	цикл `for step in range(0, args.num_steps)`: 4 среды × 128 шагов = батч из 512 переходов
Преимущество через GAE	блок `if args.gae:` — `delta` это TD-ошибка, `lastgaelam` накапливает её с весом $\gamma\lambda$
Несколько эпох на одном батче	`for epoch in range(args.update_epochs)` и мини-батчи по `minibatch_size`
Отношение $r_t(\theta)$	`ratio = logratio.exp()`, где `logratio = newlogprob - b_logprobs`
Обрезанная суррогатная цель	`pg_loss1`, `pg_loss2` и `torch.max` (минимум цели = максимум потерь со знаком минус)
$c_1 L^{VF}$ и $c_2 S$	`v_loss * args.vf_coef` и `- args.ent_coef * entropy_loss`
Насколько политика ушла от старой	`approx_kl` и `clipfrac` — доля примеров, где сработала обрезка

Коэффициенты по умолчанию: $\epsilon = 0{,}2$ , $\gamma = 0{,}99$ , $\lambda = 0{,}95$ , $c_1 = 0{,}5$ , $c_2 = 0{,}01$ .

Вход в аккаунт Hugging Face

Чтобы поделиться моделью с сообществом, нужно сделать ещё три шага.

1️⃣ (Если ещё нет) создайте аккаунт на HF: huggingface.co/join.

2️⃣ Войдите и получите токен авторизации на сайте Hugging Face: создайте новый токен (huggingface.co/settings/tokens) с правом записи (write), скопируйте его, запустите ячейку ниже и вставьте токен.

from huggingface_hub import notebook_login
notebook_login()
!git config --global credential.helper store

Если вы не пользуетесь Google Colab или Jupyter Notebook, используйте вместо этого команду huggingface-cli login.

Запускаем обучение 🔥

⚠️ ⚠️ ⚠️ Не используйте тот же id репозитория, что в юните 1.

Мы написали PPO с нуля и добавили интеграцию с Hugging Face — всё готово к обучению 🔥
Сначала скопируйте весь код в файл, который назовёте ppo.py. В оригинале здесь два скриншота Colab: как создать файл в панели файлов и вставить в него код.
Теперь остаётся запустить этот скрипт командой python <имя-скрипта>.py с дополнительными параметрами, которые мы задали через argparse.
Стоит поменять и другие гиперпараметры, иначе обучение будет не очень стабильным.

python ppo.py --env-id="LunarLander-v2" --repo-id="YOUR_REPO_ID" --total-timesteps=50000

Дополнительные задания 🏆

Лучший способ учиться — пробовать самому! Почему бы не попробовать другую среду? Или не переделать реализацию под Gymnasium?

Увидимся во второй части юнита 8, где мы будем обучать агентов играть в Doom 🔥 Продолжайте учиться и оставайтесь крутыми 🤗

Часть 2. PPO с Sample Factory: играем в Doom по пикселям

Введение

Во второй части юнита 8 мы глубже займёмся оптимизацией PPO с помощью Sample-Factory — асинхронной реализации алгоритма PPO — и обучим агента играть в ViZDoom (открытую версию Doom).

В ноутбуке вы обучите агента уровню Health Gathering, где нужно собирать аптечки, чтобы не погибнуть. После этого можно обучить агента более сложным уровням, например Deathmatch.

Звучит увлекательно? Начнём! 🚀

Это практическое занятие подготовил Эдвард Бичинг — исследователь машинного обучения в Hugging Face. Он работал над Godot Reinforcement Learning Agents — открытым интерфейсом для разработки сред и агентов в игровом движке Godot.

Ноутбук Colab: unit8_part2.ipynb.

Юнит 8, часть 2: продвинутое глубокое обучение с подкреплением. Играем в Doom по пикселям с Sample Factory

В этом ноутбуке мы научимся обучать глубокую нейросеть собирать предметы в трёхмерной среде на основе игры Doom; в оригинале показано видео получившейся политики. Политику мы обучаем с помощью Sample Factory — асинхронной реализации алгоритма PPO.

Обратите внимание:

Sample Factory — продвинутый фреймворк RL, и он работает только в Linux и Mac (не в Windows).
Лучше всего фреймворк работает на машине с GPU и множеством ядер CPU, где достигает скорости 100 тысяч взаимодействий в секунду. Ресурсы обычного ноутбука Colab ограничивают производительность библиотеки, поэтому скорость в этих условиях не отражает реальной производительности.
Результаты Sample Factory в разных условиях опубликованы; посмотрите примеры, если хотите узнать больше.

Чтобы занятие засчиталось для сертификата, нужно выложить одну модель:

doom_health_gathering_supreme — результат ≥ 5.

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

Включаем GPU 💪

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

Sample Factory

Sample Factory — одна из самых быстрых библиотек RL; её конёк — очень эффективные синхронные и асинхронные реализации градиента политики (PPO).

Sample Factory тщательно протестирована, ею пользуются многие исследователи и практики, и она активно поддерживается. Наша реализация, как известно, достигает уровня лучших результатов (SOTA) в разных областях, сокращая время обучения RL-экспериментов и требования к железу.

Ключевые возможности

Сильно оптимизированная архитектура алгоритма для максимальной пропускной способности обучения.
Синхронный и асинхронный режимы обучения.
Последовательный (однопроцессный) режим для удобной отладки.
Оптимальная производительность и в средах на CPU, и в средах с ускорением на GPU.
Обучение одного и нескольких агентов, self-play, обучение нескольких политик одновременно на одной или нескольких GPU.
Популяционное обучение (PBT, Population-Based Training).
Дискретные, непрерывные и гибридные пространства действий.
Пространства наблюдений из векторов, изображений и словарей.
Автоматически строит архитектуру модели по описанию пространств действий и наблюдений. Поддерживает свои архитектуры моделей.
Сделана так, чтобы её импортировали в другие проекты; свои среды — полноправные участники.
Подробные сводки в WandB и TensorBoard, свои метрики.
Интеграция с Hugging Face 🤗 (загрузка обученных моделей и метрик на Hub).
Несколько примеров интеграций сред с настроенными параметрами и обученными моделями.

Все перечисленные политики доступны на 🤗 Hub. Ищите по тегу sample-factory.

Как устроена Sample Factory

Sample Factory — одна из самых оптимизированных реализаций RL, доступных сообществу.

Она запускает несколько процессов: воркеры сбора опыта (rollout workers), воркеры инференса (inference workers) и воркер-обучатель (learner worker).

Воркеры общаются через общую память, что снижает стоимость обмена данными между процессами.

Воркеры сбора опыта взаимодействуют со средой и отправляют наблюдения воркерам инференса.

Воркеры инференса обращаются к зафиксированной версии политики и отправляют действия обратно воркерам сбора опыта.

Через k шагов воркеры сбора опыта отправляют траекторию опыта воркеру-обучателю, который обновляет по ней сеть политики агента.

flowchart LR
R["Воркеры сбора опыта: среды"] -->|"наблюдения"| I["Воркеры инференса: политика"]
I -->|"действия"| R
R -->|"траектории через k шагов"| L["Обучатель: обновляет политику"]
L -->|"новые веса"| I

Процессы Sample Factory общаются через общую память (в оригинале — схема)

Модели actor-critic в Sample Factory

Модели actor-critic в Sample Factory состоят из трёх частей:

кодировщик (encoder) — обрабатывает входные наблюдения (изображения, векторы) и отображает их в вектор. Эту часть модели вам, скорее всего, и захочется настраивать;
ядро (core) — объединяет векторы от одного или нескольких кодировщиков; в агенте с памятью может включать однослойную или многослойную LSTM/GRU;
декодер (decoder) — применяет дополнительные слои к выходу ядра модели перед вычислением выходов политики и ценности.

Библиотека спроектирована так, чтобы автоматически поддерживать любые пространства наблюдений и действий. Пользователи легко могут добавлять свои модели. Подробнее — в документации.

ViZDoom

ViZDoom — открытый интерфейс на Python к движку Doom.

Библиотеку создали в 2016 году Марек Видмух и Михал Кемпка в Институте информатики Познанского технологического университета (Польша).

Библиотека позволяет обучать агентов прямо по пикселям экрана в разных сценариях, включая командный deathmatch (в оригинале — видео). Поскольку ViZDoom основана на игре 90-х годов, на современном железе она работает с ускорением, и сложное поведение ИИ можно выучить довольно быстро.

Возможности библиотеки:

работает на разных платформах (Linux, macOS, Windows);
API для Python и C++;
обёртки сред OpenAI Gym;
легко создавать свои сценарии (есть визуальные редакторы, язык сценариев и примеры);
асинхронный и синхронный режимы для одного и нескольких игроков;
лёгкая (несколько мегабайт) и быстрая (до 7000 кадров в секунду в синхронном режиме в одном потоке);
настраиваемое разрешение и параметры отрисовки;
доступ к буферу глубины (3D-зрение);
автоматическая разметка игровых объектов, видимых в кадре;
доступ к аудиобуферу;
доступ к списку акторов и объектов и к геометрии карты;
отрисовка вне экрана и запись эпизодов;
масштабирование времени в асинхронном режиме.

Сначала ставим зависимости, нужные среде ViZDoom

Среда выполнения Colab настроена — начнём с установки зависимостей, нужных для запуска ViZDoom в Linux.

Если вы работаете на своём Mac, следуйте инструкциям по установке на странице GitHub.

# Ставим зависимости ViZDoom по инструкции
# https://github.com/mwydmuch/ViZDoom/blob/master/doc/Building.md#-linux

apt-get install build-essential zlib1g-dev libsdl2-dev libjpeg-dev \
nasm tar libbz2-dev libgtk2.0-dev cmake git libfluidsynth-dev libgme-dev \
libopenal-dev timidity libwildmidi-dev unzip ffmpeg

# Библиотеки Boost
apt-get install libboost-all-dev

# Зависимости для привязок к Lua
apt-get install liblua5.1-dev

Затем ставим Sample Factory и ViZDoom

Это может занять 7 минут.

pip install sample-factory
pip install vizdoom

Настраиваем среду Doom в Sample Factory

import functools

from sample_factory.algo.utils.context import global_model_factory
from sample_factory.cfg.arguments import parse_full_cfg, parse_sf_args
from sample_factory.envs.env_utils import register_env
from sample_factory.train import run_rl

from sf_examples.vizdoom.doom.doom_model import make_vizdoom_encoder
from sf_examples.vizdoom.doom.doom_params import add_doom_env_args, doom_override_defaults
from sf_examples.vizdoom.doom.doom_utils import DOOM_ENVS, make_doom_env_from_spec


# Регистрируем все среды ViZDoom
def register_vizdoom_envs():
    for env_spec in DOOM_ENVS:
        make_env_func = functools.partial(make_doom_env_from_spec, env_spec)
        register_env(env_spec.name, make_env_func)


# Sample Factory позволяет зарегистрировать свою архитектуру нейросети
# Подробнее: https://github.com/alex-petrenko/sample-factory/blob/master/sf_examples/vizdoom/doom/doom_model.py
def register_vizdoom_models():
    global_model_factory().register_encoder_factory(make_vizdoom_encoder)


def register_vizdoom_components():
    register_vizdoom_envs()
    register_vizdoom_models()


# разбираем аргументы командной строки и создаём конфигурацию
def parse_vizdoom_cfg(argv=None, evaluation=False):
    parser, _ = parse_sf_args(argv=argv, evaluation=evaluation)
    # параметры, специфичные для сред Doom
    add_doom_env_args(parser)
    # переопределяем значения параметров алгоритма по умолчанию для Doom
    doom_override_defaults(parser)
    # второй проход разбора даёт итоговую конфигурацию
    final_cfg = parse_full_cfg(parser, argv)
    return final_cfg

Настройка закончена — можно обучать агента. Мы выбрали задачу ViZDoom под названием Health Gathering Supreme.

Сценарий Health Gathering Supreme

В оригинале здесь кадр из игры: зелёный кислотный пол и аптечки.

Цель сценария — научить агента выживать, не зная, что помогает ему выжить. Агент знает только, что жизнь драгоценна, а смерть — плохо, поэтому он должен понять, что продлевает его существование и что здоровье связано с выживанием.

Карта — прямоугольник со стенами и зелёным кислотным полом, который периодически наносит игроку урон. Сначала по карте равномерно разбросано несколько аптечек. Время от времени с неба падает новая аптечка. Аптечки восстанавливают часть здоровья, и чтобы выжить, агенту нужно их подбирать. Эпизод заканчивается после смерти игрока или по тайм-ауту.

Остальные настройки:

награда за то, что жив (Living_reward) = 1;
3 доступные кнопки: повернуть налево, повернуть направо, идти вперёд;
1 доступная игровая переменная: HEALTH (здоровье);
штраф за смерть = 100.

Подробнее о сценариях ViZDoom — здесь.

Для ViZDoom создано и несколько более сложных сценариев, например описанные на этой странице GitHub.

Обучаем агента

Обучим агента на 4 000 000 шагов. Это займёт около 20 минут.

## Запускаем обучение, это займёт около 15 минут
register_vizdoom_components()

# Сегодня обучаемся на сценарии health gathering
# другие сценарии: "doom_basic", "doom_two_colors_easy", "doom_dm", "doom_dwango5", "doom_my_way_home", "doom_deadly_corridor", "doom_defend_the_center", "doom_defend_the_line"
env = "doom_health_gathering_supreme"
cfg = parse_vizdoom_cfg(
    argv=[f"--env={env}", "--num_workers=8", "--num_envs_per_worker=4", "--train_for_env_steps=4000000"]
)

status = run_rl(cfg)

Посмотрим, как играет обученная политика, и запишем видео агента

from sample_factory.enjoy import enjoy

cfg = parse_vizdoom_cfg(
    argv=[f"--env={env}", "--num_workers=1", "--save_video", "--no_render", "--max_num_episodes=10"], evaluation=True
)
status = enjoy(cfg)

Теперь посмотрим видео

from base64 import b64encode
from IPython.display import HTML

mp4 = open("/content/train_dir/default_experiment/replay.mp4", "rb").read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML(
    """
<video width=640 controls>
      <source src="%s" type="video/mp4">
</video>
"""
    % data_url
)

Агент чему-то научился, но играет так себе. Очевидно, учить нужно дольше. Но давайте выложим эту модель на Hub.

Загружаем чекпоинт и видео на Hugging Face Hub

Чтобы поделиться моделью с сообществом, нужно сделать ещё три шага.

1️⃣ (Если ещё нет) создайте аккаунт на HF: huggingface.co/join.

Если вы не пользуетесь Google Colab или Jupyter Notebook, используйте вместо этого команду huggingface-cli login.

from huggingface_hub import notebook_login
notebook_login()
!git config --global credential.helper store

from sample_factory.enjoy import enjoy

hf_username = "ThomasSimonini"  # впишите сюда своё имя пользователя на Hugging Face

cfg = parse_vizdoom_cfg(
    argv=[
        f"--env={env}",
        "--num_workers=1",
        "--save_video",
        "--no_render",
        "--max_num_episodes=10",
        "--max_num_frames=100000",
        "--push_to_hub",
        f"--hf_repository={hf_username}/rl_course_vizdoom_health_gathering_supreme",
    ],
    evaluation=True,
)
status = enjoy(cfg)

Загрузим другую модель

Этот агент играл неплохо, но можно лучше! Скачаем с Hub и посмотрим агента, обученного на 10 млрд шагов.

# скачиваем агента с Hub
python -m sample_factory.huggingface.load_from_hub -r edbeeching/doom_health_gathering_supreme_2222 -d ./train_dir

ls train_dir/doom_health_gathering_supreme_2222

env = "doom_health_gathering_supreme"
cfg = parse_vizdoom_cfg(
    argv=[
        f"--env={env}",
        "--num_workers=1",
        "--save_video",
        "--no_render",
        "--max_num_episodes=10",
        "--experiment=doom_health_gathering_supreme_2222",
        "--train_dir=train_dir",
    ],
    evaluation=True,
)
status = enjoy(cfg)

mp4 = open("/content/train_dir/doom_health_gathering_supreme_2222/replay.mp4", "rb").read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML(
    """
<video width=640 controls>
      <source src="%s" type="video/mp4">
</video>
"""
    % data_url
)

Дополнительные задания 🏆: Doom Deathmatch

Обучение агента для deathmatch в Doom занимает много часов на машине мощнее, чем доступна в Colab.

К счастью, мы уже обучили агента в этом сценарии, и он есть на 🤗 Hub! Скачаем модель и посмотрим, как играет агент.

# Скачиваем агента с Hub
python -m sample_factory.huggingface.load_from_hub -r edbeeching/doom_deathmatch_bots_2222 -d ./train_dir

Агент играет долго, поэтому запись видео может занять 10 минут.

from sample_factory.enjoy import enjoy

register_vizdoom_components()
env = "doom_deathmatch_bots"
cfg = parse_vizdoom_cfg(
    argv=[
        f"--env={env}",
        "--num_workers=1",
        "--save_video",
        "--no_render",
        "--max_num_episodes=1",
        "--experiment=doom_deathmatch_bots_2222",
        "--train_dir=train_dir",
    ],
    evaluation=True,
)
status = enjoy(cfg)
mp4 = open("/content/train_dir/doom_deathmatch_bots_2222/replay.mp4", "rb").read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML(
    """
<video width=640 controls>
      <source src="%s" type="video/mp4">
</video>
"""
    % data_url
)

Вы можете попробовать обучить агента в этой среде по коду выше, но не в Colab. Удачи 🤞

Если хотите сценарий попроще, попробуйте обучиться в другом сценарии ViZDoom, например doom_deadly_corridor или doom_defend_the_center.

На этом последний юнит заканчивается. Но это ещё не всё! 🤗 В следующем бонусном разделе собраны одни из самых интересных, продвинутых и передовых работ в глубоком обучении с подкреплением.

Заключение второй части

На сегодня всё. Поздравляем с окончанием юнита и практического занятия! ⭐️

Вы успешно обучили агента для Doom — почему бы не попробовать deathmatch? Помните, это гораздо более сложный уровень, чем тот, на котором вы только что обучались, но это хороший эксперимент, и я советую его попробовать.

Если попробуете, поделитесь моделью в канале #rl-i-made-this на нашем Discord-сервере.

На этом последний юнит заканчивается, но это ещё не всё! 🤗 В следующем бонусном юните собраны одни из самых интересных, продвинутых и передовых работ в глубоком обучении с подкреплением — главное из него переведено в путеводителе по юнитам 5, 7 и бонусам.

До встречи 🔥 Продолжайте учиться и оставайтесь крутыми 🤗

Словарик

English	Русский
single-file implementation	реализация в одном файле
rollout, rollout worker	сбор опыта, воркер сбора опыта
inference worker / learner worker	воркер инференса / воркер-обучатель
shared memory	общая память
vectorized environments (`SyncVectorEnv`)	векторизованные среды (несколько копий параллельно)
learning rate annealing	отжиг (постепенное уменьшение) скорости обучения
advantage normalization	нормализация преимуществ
value loss clipping	обрезка loss функции ценности
gradient clipping (`max_grad_norm`)	обрезка градиента по норме
approximate KL, clip fraction	приближённая KL-дивергенция, доля обрезанных примеров
explained variance	объяснённая дисперсия (качество критика)
encoder / core / decoder	кодировщик / ядро / декодер
Population-Based Training (PBT)	популяционное обучение
medkit, living reward, death penalty	аптечка, награда за жизнь, штраф за смерть

Перевод выполнен по лицензии оригинала (Apache-2.0); автор оригинала — Hugging Face (Томас Симонини; часть 2 — Эдвард Бичинг; реализация PPO — Коста Хуан, CleanRL). Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04