Deep RL Course, юнит 4, практика: REINFORCE на PyTorch

Это перевод практического занятия юнита 4 курса Hugging Face Deep RL Course (ноутбук notebooks/unit4/unit4.ipynb). Теория — на странице «Юнит 4: градиент политики», карта курса — в путеводителе.

Как читать вместе с курсом ChinaAI

Это главная практика курса Hugging Face для нашего LLM-трека: первый агент, написанный своими руками на PyTorch. Делайте её после теории юнита 4 и раздела «REINFORCE и reward-to-go» урока 6.0.3. Сравните код с REINFORCE для «игрушечной LLM» в уроке 6.3.1: там действие — токен, а в остальном те же три строки — логарифм вероятности, отдача, минус их произведения в потерях.

Примечание переводчика: версии библиотек

Ноутбук написан для старого пакета gym, а не Gymnasium: среда Pixelcopter из gym-games не обновлена. В старом gym env.reset() возвращает только состояние, а env.step() — четыре значения state, reward, done, info. Если вы ставите свежие пакеты, CartPole удобнее запускать в Gymnasium — в конце страницы есть версия цикла обучения для Gymnasium. Обучение на CartPole идёт на CPU за минуты; GPU не нужна. Публикация на Hub необязательна: лидерборд курса выключен.

Теперь, когда мы изучили теорию REINFORCE, вы готовы написать агента REINFORCE на PyTorch. И проверить его надёжность на CartPole-v1 и PixelCopter.

Потом вы сможете дорабатывать и улучшать эту реализацию для более сложных сред. В оригинале здесь анимация двух сред.

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

• получить результат ≥ 350 в Cartpole-v1;
• получить результат ≥ 5 в PixelCopter.

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

Чтобы начать, нажмите кнопку «Open In Colab»: ноутбук unit4.ipynb в Colab.

Мы настоятельно рекомендуем выполнять практические задания в Google Colab, а не на личном компьютере. С Colab вы сосредоточитесь на учёбе и экспериментах и не будете отвлекаться на техническую настройку окружения.

Юнит 4: первый алгоритм глубокого обучения с подкреплением на PyTorch — REINFORCE. И проверка его надёжности 💪

В этом ноутбуке вы с нуля напишете свой первый алгоритм глубокого обучения с подкреплением — REINFORCE (его называют также Монте-Карло градиентом политики).

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

Точнее, REINFORCE — метод градиента политики, подкласс методов на основе политики, который оптимизирует политику напрямую, оценивая веса оптимальной политики градиентным подъёмом.

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

• CartPole-v1;
• PixelcopterEnv.

⬇️ Вот пример того, чего вы добьётесь к концу ноутбука (в оригинале — анимация). ⬇️

🎮 Среды:

• CartPole-v1;
• PixelCopter.

📚 Библиотеки RL: Python и PyTorch.

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

Цели ноутбука 🏆

К концу ноутбука вы сможете:

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

Что нужно заранее 🏗️

🔲 📚 Изучить градиент политики по юниту 4.

Напишем алгоритм REINFORCE с нуля 🔥

Советы 💡

Лучше запускать копию ноутбука из своего Google Drive: тогда, если сессия оборвётся по тайм-ауту, сохранённый ноутбук останется у вас и не придётся заполнять всё заново. Для этого нажмите Ctrl + S или File > Save a copy in Google Drive.

Включаем GPU 💪

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

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

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

%%capture
!apt install python-opengl
!apt install ffmpeg
!apt install xvfb
!pip install pyvirtualdisplay
!pip install pyglet==1.5.1

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

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

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

Первый шаг — установить зависимости. Их несколько:

• gym;
• gym-games — дополнительные среды gym, сделанные на PyGame;
• huggingface_hub — Hub служит центральным местом, где любой может делиться моделями и датасетами и искать их. В нём есть версии, метрики, визуализации и другие возможности для совместной работы.

Вы можете спросить, почему мы ставим gym, а не gymnasium — более новую версию gym. Потому что gym-games, которые мы используем, ещё не обновлены под gymnasium.

Отличия, с которыми вы здесь столкнётесь:

• в gym нет terminated и truncated, есть только done;
• в gym вызов env.step() возвращает state, reward, done, info.

Подробнее об отличиях Gym и Gymnasium — в руководстве по миграции.

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

pip install -r https://raw.githubusercontent.com/huggingface/deep-rl-class/main/notebooks/unit4/requirements-unit4.txt

Примечание переводчика: что внутри requirements-unit4.txt

Файл ставит PyGame Learning Environment и gym-games прямо с GitHub (git+https://github.com/ntasfi/PyGame-Learning-Environment.git и git+https://github.com/simoninithomas/gym-games), а также huggingface_hub, imageio-ffmpeg и pyyaml==6.0. Для одного CartPole всё это не нужно: хватит pip install gymnasium torch matplotlib.

Импортируем пакеты 📦

Кроме установленных библиотек, импортируем:

• imageio — библиотеку, которая поможет записать видео игры.

import numpy as np

from collections import deque

import matplotlib.pyplot as plt
%matplotlib inline

# PyTorch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

# Gym
import gym
import gym_pygame

# Hugging Face Hub
from huggingface_hub import notebook_login  # чтобы войти в аккаунт Hugging Face и загружать модели на Hub
import imageio

Проверяем, есть ли GPU

• Проверим, есть ли GPU.
• Если есть, вы увидите device:cuda0.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print(device)

Теперь всё готово, чтобы реализовать алгоритм REINFORCE 🔥

Первый агент: играем в CartPole-v1 🤖

Создаём среду CartPole и разбираемся, как она устроена

Зачем брать простую среду вроде CartPole-v1?

Как объясняется в Reinforcement Learning Tips and Tricks, когда вы реализуете агента с нуля, нужно убедиться, что он работает правильно, и найти ошибки на простых средах, прежде чем идти дальше: в простых средах искать ошибки гораздо легче.

Постарайтесь получить хоть какие-то «признаки жизни» на игрушечных задачах.

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

Среда CartPole-v1

Шест прикреплён неуправляемым шарниром к тележке, которая движется по рельсу без трения. Маятник ставят на тележку вертикально, и цель — удерживать шест в равновесии, прикладывая к тележке силы влево и вправо.

Итак, начинаем с CartPole-v1. Цель — толкать тележку влево или вправо так, чтобы шест оставался в равновесии.

Эпизод заканчивается, если:

• угол шеста больше ±12°;
• положение тележки больше ±2,4;
• длина эпизода больше 500.

За каждый шаг, пока шест в равновесии, мы получаем награду 💰 +1.

env_id = "CartPole-v1"
# Создаём среду
env = gym.make(env_id)

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

# Получаем размеры пространства состояний и пространства действий
s_size = env.observation_space.shape[0]
a_size = env.action_space.n

print("_____OBSERVATION SPACE_____ \n")
print("The State Space is: ", s_size)
print("Sample observation", env.observation_space.sample())  # случайное наблюдение

print("\n _____ACTION SPACE_____ \n")
print("The Action Space is: ", a_size)
print("Action Space Sample", env.action_space.sample())  # случайное действие

Примечание переводчика: что вы увидите

Пространство состояний — 4 числа (положение и скорость тележки, угол и угловая скорость шеста), пространство действий — 2 действия (толкнуть влево и толкнуть вправо). Значит, s_size = 4, a_size = 2.

Строим архитектуру REINFORCE

Эта реализация основана на трёх реализациях:

• официальном примере RL в PyTorch;
• REINFORCE от Udacity;
• улучшении интеграции от Chris1nexus.

flowchart LR
S["Состояние: 4 числа"] --> FC1["fc1: Linear 4 → h_size, ReLU"]
FC1 --> FC2["fc2: Linear h_size → 2"]
FC2 --> SM["Softmax"]
SM --> P["Вероятности: влево, вправо"]

Сеть политики REINFORCE для CartPole (в оригинале — картинка)

Нам нужны:

• два полносвязных слоя (fc1 и fc2);
• ReLU как функция активации после fc1;
• Softmax на выходе, чтобы получить распределение вероятностей по действиям.

class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super(Policy, self).__init__()
        # Создайте два полносвязных слоя



    def forward(self, x):
        # Опишите прямой проход:
        # состояние идёт в fc1, затем применяем функцию активации ReLU

        # выход fc1 идёт в fc2

        # на выходе — softmax

    def act(self, state):
        """
        По состоянию выбрать действие
        """
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = self.forward(state).cpu()
        m = Categorical(probs)
        action = np.argmax(m)
        return action.item(), m.log_prob(action)

Решение

class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(s_size, h_size)
        self.fc2 = nn.Linear(h_size, a_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.softmax(x, dim=1)

    def act(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = self.forward(state).cpu()
        m = Categorical(probs)
        action = np.argmax(m)
        return action.item(), m.log_prob(action)

Я допустил ошибку — догадаетесь, где?

• Чтобы выяснить, сделаем прямой проход:

debug_policy = Policy(s_size, a_size, 64).to(device)
debug_policy.act(env.reset())

• Сообщение об ошибке гласит: ValueError: The value argument to log_prob must be a Tensor.
• Значит, action в m.log_prob(action) должен быть тензором, а он не тензор.
• Понимаете почему? Посмотрите на функцию act и попробуйте понять, что в ней не так.

Совет 💡: в реализации что-то не так. Вспомните, что в функции act мы хотим сэмплировать действие из распределения вероятностей по действиям.

(Настоящее) решение

class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(s_size, h_size)
        self.fc2 = nn.Linear(h_size, a_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.softmax(x, dim=1)

    def act(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = self.forward(state).cpu()
        m = Categorical(probs)
        action = m.sample()
        return action.item(), m.log_prob(action)

На CartPole отлаживаться было проще: мы знаем, что ошибка в нашей реализации, а не в простой среде.

• Раз мы хотим сэмплировать действие из распределения вероятностей по действиям, нельзя писать action = np.argmax(m): так всегда выбиралось бы действие с наибольшей вероятностью.
• Это нужно заменить на action = m.sample(), который сэмплирует действие из распределения $P(\cdot \mid s)$.

Примечание переводчика: почему argmax — это ещё и ошибка в теории

Кроме падения с ValueError (у np.argmax от объекта Categorical нет тензора), argmax сломал бы и саму идею: политика стала бы детерминированной, агент перестал бы исследовать среду, а градиент политики выводится для сэмплированных действий — математическое ожидание в формуле берётся по траекториям текущей политики.

Строим алгоритм обучения REINFORCE

В оригинале здесь картинка с псевдокодом REINFORCE. Комментарии в коде ниже ссылаются на её строки. Схема алгоритма в словах (реконструкция по этим комментариям):

для каждого эпизода:                                    # строка 3
    сыграть эпизод политикой π_θ, запоминая log π_θ(a_t|s_t) и награды   # строка 4
    для каждого шага t эпизода посчитать отдачу
        G_t = r_{t+1} + γ·r_{t+2} + γ²·r_{t+3} + …       # строка 6
    потери = −Σ_t log π_θ(a_t|s_t) · G_t                 # строка 7
    шаг градиентного спуска по θ                         # строка 8

• Когда мы считаем отдачу $G_t$ (строка 6), мы суммируем дисконтированные награды, начиная с шага $t$.
• Почему? Потому что политика должна подкреплять действия по их последствиям: награды, полученные до действия, бесполезны (они получены не из-за этого действия), важны только те, что пришли после него.
• Прежде чем писать код, прочитайте раздел Spinning Up «Don’t let the past distract you» («Не отвлекайтесь на прошлое»): там объясняется, почему мы используем градиент политики с отдачей «от текущего шага» (reward-to-go). Перевод этого раздела — в нашей странице Spinning Up, часть 3.

Мы используем интересный приём, который реализовал Chris1nexus, чтобы эффективно считать отдачу на каждом шаге. Процедура объяснена в комментариях. Посмотрите также объяснение в пул-реквесте. Но общая идея — эффективно посчитать отдачу на каждом шаге.

Второй вопрос, который может у вас возникнуть: почему мы минимизируем потери? Разве мы не говорили о градиентном подъёме, а не спуске?

• Мы хотим максимизировать функцию полезности $J(\theta)$, но в PyTorch и TensorFlow лучше минимизировать целевую функцию.
  • Допустим, мы хотим подкрепить действие 3 на каком-то шаге. До обучения вероятность этого действия — 0,25.
  • Значит, мы хотим изменить $\theta$ так, чтобы $\pi_\theta(a_3 \mid s; \theta) > 0{,}25$.
  • Поскольку все вероятности в сумме дают 1, максимизация $\pi_\theta(a_3 \mid s; \theta)$ уменьшит вероятности остальных действий.
  • Значит, PyTorch нужно сказать: минимизируй $1 - \pi_\theta(a_3 \mid s; \theta)$.
  • Эти потери стремятся к 0, когда $\pi_\theta(a_3 \mid s; \theta)$ приближается к 1.
  • Так мы подталкиваем градиент к максимизации $\pi_\theta(a_3 \mid s; \theta)$.

Примечание переводчика: что на самом деле минимизирует код

Пример с $1 - \pi_\theta$ — интуиция авторов. В коде минимизируется $-\sum_t \log \pi_\theta(a_t \mid s_t) \cdot G_t$: это градиентный подъём по $\sum_t \log \pi_\theta(a_t \mid s_t) \cdot G_t$, записанный как спуск. Знак $G_t$ решает, поднимать вероятность действия или опускать. Перед этим отдачи стандартизируются (вычитается среднее, делится на стандартное отклонение): примерно половина действий эпизода получает отрицательный вес. Это грубый вариант baseline, о котором рассказывает юнит 6.

def reinforce(policy, optimizer, n_training_episodes, max_t, gamma, print_every):
    # Помогает считать счёт во время обучения
    scores_deque = deque(maxlen=100)
    scores = []
    # Строка 3 псевдокода
    for i_episode in range(1, n_training_episodes+1):
        saved_log_probs = []
        rewards = []
        state = # TODO: сбросьте среду
        # Строка 4 псевдокода
        for t in range(max_t):
            action, log_prob = # TODO: получите действие
            saved_log_probs.append(log_prob)
            state, reward, done, _ = # TODO: сделайте шаг в среде
            rewards.append(reward)
            if done:
                break
        scores_deque.append(sum(rewards))
        scores.append(sum(rewards))

        # Строка 6 псевдокода: считаем отдачу
        returns = deque(maxlen=max_t)
        n_steps = len(rewards)
        # Считаем дисконтированные отдачи на каждом шаге
        # как сумму дисконтированной отдачи с момента t (G_t) и награды в момент t

        # За время O(N), где N — число шагов
        # (это определение дисконтированной отдачи G_t взято со страницы 44
        # черновика 2-го издания Саттона и Барто, 2017)
        # G_t = r_(t+1) + r_(t+2) + ...

        # При такой формулировке отдачу на каждом шаге t можно посчитать,
        # переиспользуя уже посчитанную будущую отдачу G_(t+1):
        # G_t = r_(t+1) + gamma*G_(t+1)
        # G_(t-1) = r_t + gamma* G_t
        # (это подход динамического программирования: мы запоминаем решения,
        # чтобы не вычислять их многократно)

        # Это верно, потому что выражение выше равносильно (см. также стр. 46
        # черновика 2-го издания Саттона и Барто, 2017)
        # G_(t-1) = r_t + gamma*r_(t+1) + gamma*gamma*r_(t+2) + ...


        ## Исходя из этого, отдачу на шаге t считаем как
        #               gamma[t] * return[t] + reward[t]
        #
        ## Считаем с последнего шага к первому, чтобы применять формулу выше
        ## и избежать лишних вычислений, которые понадобились бы,
        ## если бы мы шли от первого шага к последнему.

        ## Поэтому очередь returns будет хранить отдачи в хронологическом порядке, от t=0 до t=n_steps,
        ## благодаря функции appendleft(), которая добавляет элемент в позицию 0 за постоянное время O(1);
        ## обычному списку Python для этого понадобилось бы O(N).
        for t in range(n_steps)[::-1]:
            disc_return_t = (returns[0] if len(returns)>0 else 0)
            returns.appendleft(    ) # TODO: допишите здесь

        ## стандартизация отдач делает обучение стабильнее
        eps = np.finfo(np.float32).eps.item()

        ## eps — наименьшее представимое число float; его прибавляют
        # к стандартному отклонению отдач, чтобы избежать численной неустойчивости
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + eps)

        # Строка 7:
        policy_loss = []
        for log_prob, disc_return in zip(saved_log_probs, returns):
            policy_loss.append(-log_prob * disc_return)
        policy_loss = torch.cat(policy_loss).sum()

        # Строка 8: PyTorch предпочитает градиентный спуск
        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        if i_episode % print_every == 0:
            print('Episode {}\tAverage Score: {:.2f}'.format(i_episode, np.mean(scores_deque)))

    return scores

Решение

def reinforce(policy, optimizer, n_training_episodes, max_t, gamma, print_every):
    # Помогает считать счёт во время обучения
    scores_deque = deque(maxlen=100)
    scores = []
    # Строка 3 псевдокода
    for i_episode in range(1, n_training_episodes + 1):
        saved_log_probs = []
        rewards = []
        state = env.reset()
        # Строка 4 псевдокода
        for t in range(max_t):
            action, log_prob = policy.act(state)
            saved_log_probs.append(log_prob)
            state, reward, done, _ = env.step(action)
            rewards.append(reward)
            if done:
                break
        scores_deque.append(sum(rewards))
        scores.append(sum(rewards))

        # Строка 6 псевдокода: считаем отдачу
        returns = deque(maxlen=max_t)
        n_steps = len(rewards)
        # (подробные комментарии о подсчёте отдач — в заготовке выше)
        # Идём с конца эпизода: G_t = r_(t+1) + gamma * G_(t+1)
        for t in range(n_steps)[::-1]:
            disc_return_t = returns[0] if len(returns) > 0 else 0
            returns.appendleft(gamma * disc_return_t + rewards[t])

        ## стандартизация отдач делает обучение стабильнее
        eps = np.finfo(np.float32).eps.item()
        ## eps — наименьшее представимое число float; его прибавляют
        # к стандартному отклонению отдач, чтобы избежать численной неустойчивости
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + eps)

        # Строка 7:
        policy_loss = []
        for log_prob, disc_return in zip(saved_log_probs, returns):
            policy_loss.append(-log_prob * disc_return)
        policy_loss = torch.cat(policy_loss).sum()

        # Строка 8: PyTorch предпочитает градиентный спуск
        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        if i_episode % print_every == 0:
            print("Episode {}\tAverage Score: {:.2f}".format(i_episode, np.mean(scores_deque)))

    return scores

Примечание переводчика: подсчёт отдач на числах

Пусть эпизод дал награды $[1, 1, 1]$ и $\gamma = 0{,}9$. Идём с конца: $G_2 = 1$; $G_1 = 1 + 0{,}9 \cdot 1 = 1{,}9$; $G_0 = 1 + 0{,}9 \cdot 1{,}9 = 2{,}71$. Очередь returns после appendleft — $[2{,}71;\ 1{,}9;\ 1]$, уже в порядке времени. В юните 4 для CartPole взято $\gamma = 1{,}0$: тогда $G_t$ — просто число оставшихся шагов эпизода.

Обучаем

• Всё готово, чтобы обучить агента.
• Но сначала заведём переменную со всеми гиперпараметрами обучения.
• Параметры обучения можно (и нужно 😉) менять.

cartpole_hyperparameters = {
    "h_size": 16,
    "n_training_episodes": 1000,
    "n_evaluation_episodes": 10,
    "max_t": 1000,
    "gamma": 1.0,
    "lr": 1e-2,
    "env_id": env_id,
    "state_space": s_size,
    "action_space": a_size,
}

# Создаём политику и переносим её на устройство
cartpole_policy = Policy(
    cartpole_hyperparameters["state_space"],
    cartpole_hyperparameters["action_space"],
    cartpole_hyperparameters["h_size"],
).to(device)
cartpole_optimizer = optim.Adam(cartpole_policy.parameters(), lr=cartpole_hyperparameters["lr"])

scores = reinforce(
    cartpole_policy,
    cartpole_optimizer,
    cartpole_hyperparameters["n_training_episodes"],
    cartpole_hyperparameters["max_t"],
    cartpole_hyperparameters["gamma"],
    100,
)

Определяем метод оценки 📝

Здесь мы определяем метод оценки, которым будем проверять агента REINFORCE.

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

        for step in range(max_steps):
            action, _ = policy.act(state)
            new_state, reward, done, info = env.step(action)
            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

Оцениваем агента 📈

evaluate_agent(
    eval_env, cartpole_hyperparameters["max_t"], cartpole_hyperparameters["n_evaluation_episodes"], cartpole_policy
)

Публикуем обученную модель на Hub 🔥

Мы убедились, что после обучения результаты хорошие, и теперь можем одной строкой кода опубликовать модель на Hub 🤗. В оригинале здесь пример карточки модели.

Отправка на Hub. Этот код не меняйте

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

from pathlib import Path
import datetime
import json
import imageio

import tempfile

import os

def record_video(env, policy, out_directory, fps=30):
    """
    Записать видео игры агента
    :param env
    :param Qtable: Q-таблица агента (докстринг остался от юнита 2; здесь передаётся политика)
    :param out_directory
    :param fps: сколько кадров в секунду (для taxi-v3 и frozenlake-v1 мы берём 1)
    """
    images = []
    done = False
    state = env.reset()
    img = env.render(mode="rgb_array")
    images.append(img)
    while not done:
        # Выбираем действие политикой
        # (комментарий в оригинале остался от юнита 2: «действие с наибольшей ожидаемой будущей наградой»)
        action, _ = policy.act(state)
        state, reward, done, info = env.step(action)  # сразу пишем 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 push_to_hub(repo_id,
                model,
                hyperparameters,
                eval_env,
                video_fps=30
                ):
  """
  Оценить модель, записать видео и загрузить модель на Hugging Face Hub.
  Метод выполняет весь конвейер:
  - оценивает модель
  - создаёт карточку модели
  - записывает видео игры агента
  - отправляет всё на Hub

  :param repo_id: id репозитория модели на Hugging Face Hub
  :param model: модель PyTorch, которую хотим сохранить
  :param hyperparameters: гиперпараметры обучения
  :param eval_env: среда для оценки
  :param video_fps: сколько кадров в секунду записывать в видео
  """

  _, repo_name = repo_id.split("/")
  api = HfApi()

  # Шаг 1: создаём репозиторий
  repo_url = api.create_repo(
        repo_id=repo_id,
        exist_ok=True,
  )

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

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

    # Шаг 3: сохраняем гиперпараметры в JSON
    with open(local_directory / "hyperparameters.json", "w") as outfile:
      json.dump(hyperparameters, outfile)

    # Шаг 4: оцениваем модель и собираем JSON
    mean_reward, std_reward = evaluate_agent(eval_env,
                                            hyperparameters["max_t"],
                                            hyperparameters["n_evaluation_episodes"],
                                            model)
    # Получаем дату и время
    eval_datetime = datetime.datetime.now()
    eval_form_datetime = eval_datetime.isoformat()

    evaluate_data = {
          "env_id": hyperparameters["env_id"],
          "mean_reward": mean_reward,
          "n_evaluation_episodes": hyperparameters["n_evaluation_episodes"],
          "eval_datetime": eval_form_datetime,
    }

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

    # Шаг 5: создаём карточку модели
    env_name = hyperparameters["env_id"]

    metadata = {}
    metadata["tags"] = [
          env_name,
          "reinforce",
          "reinforcement-learning",
          "custom-implementation",
          "deep-rl-class"
      ]

    # Добавляем метрики
    eval = metadata_eval_result(
        model_pretty_name=repo_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_name,
        dataset_id=env_name,
      )

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

    model_card = f"""
  # **Reinforce** Agent playing **{env_id}**
  This is a trained model of a **Reinforce** agent playing **{env_id}** .
  To learn to use this model and train yours check Unit 4 of the Deep Reinforcement Learning Course: https://huggingface.co/deep-rl-course/unit4/introduction
  """

    readme_path = local_directory / "README.md"
    readme = ""
    if readme_path.exists():
        with readme_path.open("r", encoding="utf8") as f:
          readme = f.read()
    else:
      readme = model_card

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

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

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

    # Шаг 7: отправляем всё на Hub
    api.upload_folder(
          repo_id=repo_id,
          folder_path=local_directory,
          path_in_repo=".",
    )

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

Текст карточки модели (model_card) оставлен по-английски: он попадает на Hub. Перевод: «Агент REINFORCE играет в {env_id}. Это обученная модель агента REINFORCE для {env_id}. Как пользоваться моделью и обучить свою — в юните 4 курса Deep Reinforcement Learning».

С push_to_hub вы оцениваете агента, записываете видео игры, создаёте карточку модели и отправляете всё на Hub. Так вы можете:

• показать свою работу 🔥;
• посмотреть, как играет агент 👀;
• поделиться с сообществом агентом, которым могут пользоваться другие 💾;
• попасть на лидерборд 🏆 и сравнить агента с однокурсниками: Deep RL Leaderboard.

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

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

2️⃣ Войдите и сохраните токен авторизации с сайта Hugging Face: создайте новый токен (huggingface.co/settings/tokens) с правом записи (write).

notebook_login()

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

3️⃣ Теперь всё готово, чтобы выложить обученного агента на Hub 🤗 🔥 функцией push_to_hub() (в тексте оригинала ошибочно написано package_to_hub()).

repo_id = ""  # TODO: задайте id репозитория {username/Reinforce-{model-id}}
push_to_hub(
    repo_id,
    cartpole_policy,  # модель, которую хотим сохранить
    cartpole_hyperparameters,  # гиперпараметры
    eval_env,  # среда для оценки
    video_fps=30
)

Мы проверили надёжность реализации — попробуем среду посложнее: PixelCopter 🚁

Второй агент: PixelCopter 🚁

Изучаем среду PixelCopter 👀

Документация среды.

env_id = "Pixelcopter-PLE-v0"
env = gym.make(env_id)
eval_env = gym.make(env_id)
s_size = env.observation_space.shape[0]
a_size = env.action_space.n

print("_____OBSERVATION SPACE_____ \n")
print("The State Space is: ", s_size)
print("Sample observation", env.observation_space.sample())  # случайное наблюдение

print("\n _____ACTION SPACE_____ \n")
print("The Action Space is: ", a_size)
print("Action Space Sample", env.action_space.sample())  # случайное действие

Пространство наблюдений (7 чисел) 👀:

• положение игрока по вертикали (y);
• скорость игрока;
• расстояние от игрока до пола;
• расстояние от игрока до потолка;
• расстояние по x от игрока до следующего блока;
• координата y верха следующего блока;
• координата y низа следующего блока.

Пространство действий (2) 🎮:

• вверх (нажать газ);
• ничего не делать (не нажимать газ).

Функция награды 💰: за каждый пройденный вертикальный блок — награда +1. Каждый раз при попадании в терминальное состояние — награда −1.

Определяем новую политику 🧠

Среда сложнее, поэтому нужна более глубокая нейросеть.

class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super(Policy, self).__init__()
        # Опишите здесь три слоя

    def forward(self, x):
        # Опишите здесь прямой проход
        return F.softmax(x, dim=1)

    def act(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = self.forward(state).cpu()
        m = Categorical(probs)
        action = m.sample()
        return action.item(), m.log_prob(action)

Решение

class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(s_size, h_size)
        self.fc2 = nn.Linear(h_size, h_size * 2)
        self.fc3 = nn.Linear(h_size * 2, a_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.softmax(x, dim=1)

    def act(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = self.forward(state).cpu()
        m = Categorical(probs)
        action = m.sample()
        return action.item(), m.log_prob(action)

Задаём гиперпараметры ⚙️

• Потому что эта среда сложнее.
• Особенно скрытый слой: нужно больше нейронов.

pixelcopter_hyperparameters = {
    "h_size": 64,
    "n_training_episodes": 50000,
    "n_evaluation_episodes": 10,
    "max_t": 10000,
    "gamma": 0.99,
    "lr": 1e-4,
    "env_id": env_id,
    "state_space": s_size,
    "action_space": a_size,
}

Обучаем

Всё готово, чтобы обучить агента 🔥.

# Создаём политику и переносим её на устройство
# torch.manual_seed(50)
pixelcopter_policy = Policy(
    pixelcopter_hyperparameters["state_space"],
    pixelcopter_hyperparameters["action_space"],
    pixelcopter_hyperparameters["h_size"],
).to(device)
pixelcopter_optimizer = optim.Adam(pixelcopter_policy.parameters(), lr=pixelcopter_hyperparameters["lr"])

scores = reinforce(
    pixelcopter_policy,
    pixelcopter_optimizer,
    pixelcopter_hyperparameters["n_training_episodes"],
    pixelcopter_hyperparameters["max_t"],
    pixelcopter_hyperparameters["gamma"],
    1000,
)

Примечание переводчика: 50 000 эпизодов — это долго

Для Pixelcopter в ноутбуке стоит 50 000 эпизодов: это намного дольше CartPole. Начните с 2000–5000 эпизодов и смотрите, растёт ли средний счёт. Обратите внимание: функция reinforce использует глобальную переменную env, поэтому перед обучением Pixelcopter env должна указывать на новую среду (ячейка выше это делает).

Публикуем обученную модель на Hub 🔥

repo_id = ""  # TODO: задайте id репозитория {username/Reinforce-{model-id}}
push_to_hub(
    repo_id,
    pixelcopter_policy,  # модель, которую хотим сохранить
    pixelcopter_hyperparameters,  # гиперпараметры
    eval_env,  # среда для оценки
    video_fps=30
)

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

Лучший способ учиться — пробовать самому! Как вы видели, нынешний агент играет так себе. Для начала можно учить дольше. А ещё — попробовать найти параметры получше.

На лидерборде есть ваши агенты. Сможете пробиться наверх?

Несколько идей, как подняться в таблице:

• учить дольше;
• попробовать другие гиперпараметры, посмотрев, что сделали однокурсники: модели REINFORCE на Hub;
• выложить новую обученную модель на Hub 🔥;
• улучшить реализацию для более сложных сред (например, как насчёт того, чтобы заменить сеть на свёрточную и подавать на вход кадры?).

---

Поздравляем с окончанием юнита! Информации было много. И поздравляем с окончанием практического занятия: вы только что написали с нуля на PyTorch своего первого агента глубокого обучения с подкреплением и поделились им на Hub 🥳.

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

В следующем юните мы подробнее познакомимся с Unity ML-Agents и будем обучать агентов в средах Unity. Так вы подготовитесь к состязаниям «ИИ против ИИ», где ваши агенты будут соревноваться с другими агентами в игре в снежки и в футболе.

Звучит здорово? До встречи!

Авторы курса будут рады узнать ваше мнение о курсе и о том, как его улучшить — для этого есть форма обратной связи. Увидимся в юните 5! 🔥 Продолжайте учиться и оставайтесь крутыми 🤗

Версия для Gymnasium

Примечание переводчика: не часть оригинала

Этот раздел добавлен при переводе. Если вы ставите свежие пакеты (pip install gymnasium torch), ячейки ноутбука для CartPole нужно поправить в трёх местах: reset() возвращает пару (state, info), step() возвращает пять значений, а эпизод заканчивается, когда terminated or truncated. Остальное — классы Policy и формулы — не меняется. Код ниже — тот же REINFORCE из ноутбука; 1000 эпизодов CartPole на CPU — несколько минут.

from collections import deque

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

torch.manual_seed(0)
device = torch.device("cpu")              # для CartPole хватает процессора
env = gym.make("CartPole-v1")
s_size = env.observation_space.shape[0]   # 4 числа состояния
a_size = env.action_space.n               # 2 действия


class Policy(nn.Module):
    def __init__(self, s_size, a_size, h_size):
        super().__init__()
        self.fc1 = nn.Linear(s_size, h_size)
        self.fc2 = nn.Linear(h_size, a_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1)

    def act(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        m = Categorical(self.forward(state).cpu())
        action = m.sample()                   # сэмплируем, а не argmax
        return action.item(), m.log_prob(action)


def reinforce(policy, optimizer, n_training_episodes, max_t, gamma, print_every):
    scores_deque = deque(maxlen=100)
    for i_episode in range(1, n_training_episodes + 1):
        saved_log_probs, rewards = [], []
        state, _ = env.reset(seed=i_episode)  # Gymnasium: reset возвращает (state, info)
        for t in range(max_t):
            action, log_prob = policy.act(state)
            saved_log_probs.append(log_prob)
            state, reward, terminated, truncated, _ = env.step(action)  # пять значений
            rewards.append(reward)
            if terminated or truncated:       # вместо done
                break
        scores_deque.append(sum(rewards))

        returns = deque(maxlen=max_t)         # отдачи с конца эпизода: G_t = r + γ·G_(t+1)
        for t in reversed(range(len(rewards))):
            disc_return_t = returns[0] if len(returns) > 0 else 0
            returns.appendleft(gamma * disc_return_t + rewards[t])
        eps = np.finfo(np.float32).eps.item()
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + eps)

        policy_loss = torch.cat([-lp * g for lp, g in zip(saved_log_probs, returns)]).sum()
        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        if i_episode % print_every == 0:
            print(f"Episode {i_episode}\tAverage Score: {np.mean(scores_deque):.2f}")


policy = Policy(s_size, a_size, 16).to(device)
optimizer = optim.Adam(policy.parameters(), lr=1e-2)
reinforce(policy, optimizer, n_training_episodes=1000, max_t=1000, gamma=1.0, print_every=100)

Средний счёт (за последние 100 эпизодов) должен расти от двадцати с небольшим к сотням. В нашем проверочном прогоне этого кода (Gymnasium 1.3, PyTorch 2, CPU, 200 эпизодов вместо 1000 — около 30 секунд) он вырос с 27,7 на 50-м эпизоде до 355,8 на 200-м; у вас числа будут другими — они зависят от seed и версий библиотек. REINFORCE шумный: счёт может проседать и снова расти — это та самая высокая дисперсия, с которой начинается юнит 6.

Словарик

English	Русский
hands-on, notebook	практическое занятие, ноутбук
fully connected layer (fc1, fc2)	полносвязный слой
sample an action	сэмплировать (выбрать случайно по распределению) действие
log probability, log_prob	логарифм вероятности
discounted return $G_t$	дисконтированная отдача
reward-to-go	отдача от текущего шага
standardization of returns	стандартизация отдач
policy loss	потери политики
evaluation episodes	эпизоды оценки
model card	карточка модели
sign of life on toy problems	«признаки жизни» на игрушечных задачах