Spinning Up: упражнения

Как это связано с курсом

Упражнения 1.1 и 1.2 закрепляют гауссову политику из части 1 и PPO из урока 6.0.3. Упражнение 2.2 — лучший короткий урок о том, как RL-код ломается молча; та же беда с формами тензоров подстерегает и в практикуме по GRPO, где награды, лог-вероятности и маски имеют разные размерности. Ошибку бродкастингаbroadcasting (растягивание форм)EN broadcastingПравило, по которому NumPy и PyTorch выполняют операции над массивами разных форм без копирования данных. из этого упражнения полезно один раз увидеть своими глазами.

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

Упражнения рассчитаны на пакет spinup 2020 года: старый gym, среды MuJoCo версии -v2 (InvertedPendulum-v2, HalfCheetah-v2, Hopper-v2) через устаревший mujoco-py, MPI. Поставить это сегодня трудно. Что реально сделать:

• 1.1 — без всякого окружения, только PyTorch: ниже есть самопроверка без spinup.
• 1.2 и 1.3 — читать заготовки и решения в репозитории, а запускать на современном стеке: gymnasium и среды InvertedPendulum-v4/HalfCheetah-v4 (или новее), заменив gym.make(...) и учтя, что env.reset() и env.step() в Gymnasium возвращают больше значений. Это само по себе полезное упражнение по чтению чужого RL-кода.
• 2.1 и 2.2 — главное в них понять механизм; решения OpenAI переведены ниже.

Набор задач 1: основы реализации

Упражнение 1.1: логарифм правдоподобия гауссова распределения

Путь к упражнению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_1/exercise1_1.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_1/exercise1_1.py

Путь к решению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_1_solutions/exercise1_1_soln.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_1_solutions/exercise1_1_soln.py

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

Возможно, полезно освежить формулу из этого раздела введения в RL (раздел о стохастических политиках, диагональная гауссова политика).

Решение напишите в exercise1_1.py и запустите этот файл, чтобы автоматически проверить работу.

Критерии оценки. Ваше решение проверят, сравнив его выход с заведомо верной реализацией на батче случайных входов.

Самопроверка без пакета spinup (добавлено переводчиком)

Заготовка из репозитория импортирует эталон из spinup. Вот та же задача, которую можно запустить с одним PyTorch: эталоном служит torch.distributions.Normal. Напомним формулу: для $k$-мерного диагонального гауссова распределения со средним $\mu$ и стандартными отклонениями $\sigma = \exp(\log\sigma)$

$$\log \pi(a) = -\frac{1}{2}\left(\sum_{i=1}^k \left(\frac{(a_i - \mu_i)^2}{\sigma_i^2} + 2 \log \sigma_i \right) + k \log 2\pi \right).$$

import torch
import numpy as np

def gaussian_likelihood(x, mu, log_std):
    """
    Args:
        x: тензор формы [batch, dim]
        mu: тензор формы [batch, dim]
        log_std: тензор формы [batch, dim] или [dim]

    Returns:
        тензор формы [batch]
    """
    # YOUR CODE HERE — ваш код здесь
    return torch.zeros(1)


if __name__ == '__main__':
    # Самопроверка без пакета spinup: сравниваем с torch.distributions.Normal
    torch.manual_seed(0)
    batch_size, dim = 32, 10
    x = torch.rand(batch_size, dim)
    mu = torch.rand(batch_size, dim)
    log_std = torch.rand(dim)

    yours = gaussian_likelihood(x, mu, log_std)
    true = torch.distributions.Normal(mu, torch.exp(log_std)).log_prob(x).sum(dim=-1)
    ok = yours.shape == true.shape and torch.allclose(yours, true, atol=1e-5)
    print('Верно!' if ok else 'Пока неверно: сравните форму и значения с эталоном.')

С заготовкой файл печатает «Пока неверно…», с правильным решением — «Верно!». Эталонное решение OpenAI — пять строк в exercise1_1_soln.py; не подглядывайте, пока не попробуете сами.

Упражнение 1.2: политика для PPO

Путь к упражнению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_1/exercise1_2.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_1/exercise1_2.py

Путь к решению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_1_solutions/exercise1_2_soln.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_1_solutions/exercise1_2_soln.py

Задание. Реализуйте для PPO диагональную гауссову политику на основе MLP.

Решение напишите в exercise1_2.py и запустите этот файл, чтобы автоматически проверить работу.

Критерии оценки. Решение оценят запуском на 20 эпох в среде Gym InvertedPendulum-v2; это должно занять порядка 3–5 минут (в зависимости от вашей машины и других процессов в фоне). Планка успеха — средний результат больше 500 за последние 5 эпох или результат 1000 (максимально возможный) в последних 5 эпохах.

Что именно дописать (по заготовке, добавлено переводчиком)

В PyTorch-заготовке три места с YOUR CODE HERE: функция mlp(sizes, activation, output_activation), которая строит nn.Sequential из линейных слоёв и активаций; метод sample() класса DiagonalGaussianDistribution (выборка $\mu + \sigma \odot z$, $z \sim \mathcal{N}(0, I)$); конструктор MLPGaussianActor, где нужно завести сеть для среднего mu_net и обучаемый параметр log_std формы действия, не зависящий от наблюдения и равный $[-0{,}5, \dots, -0{,}5]$ в начале. Логарифм правдоподобия берётся из вашего решения упражнения 1.1, поэтому начните с него.

Упражнение 1.3: граф вычислений для TD3

Путь к упражнению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_1/exercise1_3.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_1/exercise1_3.py

Путь к решению:

• версия на PyTorch: spinup/algos/pytorch/td3/td3.py
• версия на TensorFlow: spinup/algos/tf1/td3/td3.py

Задание. Реализуйте основную математическую логику алгоритма TD3.

В качестве заготовки вам дан весь алгоритм TD3, кроме основной математической логики (по сути, функций потерь и промежуточных вычислений, которые для них нужны). Начните с поиска строк «YOUR CODE HERE».

Возможно, полезно освежить псевдокод на нашей странице о TD3.

Решение напишите в exercise1_3.py и запустите этот файл, чтобы увидеть результаты своей работы. Автоматической проверки в этом упражнении нет.

Критерии оценки. Проверьте свой код, запустив exercise1_3.py на HalfCheetah-v2, InvertedPendulum-v2 и ещё одной среде Gym MuJoCo на ваш выбор (задаётся флагом --env). Скрипт настроен на нейросети меньше обычных для TD3 (скрытые слои [128, 128]), максимальную длину эпизода 150 и всего 10 эпох. Цель — увидеть заметный прогресс обучения относительно быстро (по реальному времени). Эксперименты, скорее всего, займут порядка 10 минут.

Используйте флаг --use_soln, чтобы вместо вашей реализации запустить TD3 из Spinning Up. По опыту, за 10 эпох результат в HalfCheetah должен превысить 300, а в InvertedPendulum — дойти до максимума 150.

Что именно дописать (по заготовке, добавлено переводчиком)

В PyTorch-заготовке пустые места помечены комментариями: Q-значения двух критиков, сглаживание целевой политики (обрезанный шум к целевому действию), целевые Q-значения с минимумом из двух целевых критиков, среднеквадратичная ошибка относительно цели Беллмана и функция потерь политики compute_loss_pi. Все формулы — в разделе TD3 на странице TRPO, DDPG, TD3, SAC.

Набор задач 2: как алгоритмы ломаются

Упражнение 2.1: подгонка функции ценности в TRPO

Путь к упражнению. (Не применимо: кода в этом упражнении нет.)

Путь к решению. Решение — ниже на этой странице.

На качество алгоритмов градиента политики влияет множество факторов, но мало какой влияет так сильно, как качество выученной функции ценности, которая используется для оценки преимущества.

В этом упражнении вы сравните результаты запусков TRPO, в которых вы вкладываете много усилий в подгонку функции ценности (train_v_iters=80), с запусками, где вы почти не подгоняете функцию ценности (train_v_iters=0).

Задание. Выполните команду

python -m spinup.run trpo --env Hopper-v2 --train_v_iters[v] 0 80 --exp_name ex2-1 --epochs 250 --steps_per_epoch 4000 --seed 0 10 20 --dt

и постройте графики результатов. (Каждый эксперимент может занять около 10 минут, а эта команда запускает шесть.) Что вы обнаружили?

Упражнение 2.2: «тихий» баг в DDPG

Путь к упражнению:

• версия на PyTorch: spinup/exercises/pytorch/problem_set_2/exercise2_2.py
• версия на TensorFlow: spinup/exercises/tf1/problem_set_2/exercise2_2.py

Путь к решению. Решение — ниже на этой странице.

Самое трудное в написании RL-кода — борьба с багами, потому что сбои часто «тихие». Код выглядит так, будто работает правильно, но качество агента ниже, чем у реализации без бага, — иногда настолько, что он вообще ничему не учится.

В этом упражнении вы понаблюдаете за багом «в живой природе» и сравните результаты с правильным кодом. Баг один и тот же (по сути, если не в точной реализации) в версиях этого упражнения на PyTorch и на TensorFlow.

Задание. Запустите exercise2_2.py: он запустит эксперименты DDPG с багом и без него. Версия без бага запускает стандартную реализацию DDPG из Spinning Up со стандартным способом создания сетей актора и критика. Версия с багом запускает тот же код DDPG, но создаёт сети способом с ошибкой.

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

Не заглядывая в правильный код актора-критика (то есть не смотрите в файл core.py у DDPG), попробуйте понять, в чём баг, и объясните, как он всё ломает.

Подсказка. Чтобы понять, что идёт не так, подумайте о том, как код DDPG реализует граф вычислений DDPG. В версии на TensorFlow посмотрите на этот фрагмент:

# Цель Беллмана для Q-функции
backup = tf.stop_gradient(r_ph + gamma*(1-d_ph)*q_pi_targ)

# Функции потерь DDPG
pi_loss = -tf.reduce_mean(q_pi)
q_loss = tf.reduce_mean((q-backup)**2)

Как баг в коде актора-критика может повлиять на это место?

Бонус. Есть ли такие значения гиперпараметров, при которых эффект бага был бы незаметен?

Задачи посложнее

Пишите код с нуля. Как мы советуем в эссе, попробуйте заново реализовать с нуля разные алгоритмы глубокого RL.

Запросы на исследования. Если вы уверенно пишете код для глубокого обучения и глубокого RL, попробуйте продвинуться в какой-нибудь из открытых задач из списков OpenAI «Requests for Research»:

• Requests for Research 1
• Requests for Research 2

---

Решения

Сначала попробуйте сами

Ниже — решения OpenAI для упражнений 2.1 и 2.2. Решения упражнений набора 1 лежат в репозитории по путям, указанным выше. Графики из оригинала мы не копируем — они описаны словами, ссылки на картинки даны в тексте.

Решение упражнения 2.1 (нажмите, чтобы раскрыть)

Решение упражнения 2.1

Рисунок в оригинале: кривые обучения TRPO в Hopper-v2 с разными значениями train_v_iters, усреднённые по трём случайным сидам.

Разница весьма существенная: с обученной функцией ценности агент быстро продвигается. С необученной функцией ценности агент застревает в самом начале.

Почему (пояснение переводчика). При train_v_iters=0 функция ценности остаётся случайной, и оценка преимущества $\hat{A}_t$ по GAE складывается из TD-ошибок со случайными $V(s)$. Базовая линия, которая должна уменьшать разброс, превращается в источник шума, и сигнал «это действие лучше среднего» тонет в нём.

Решение упражнения 2.2 (нажмите, чтобы раскрыть)

Решение упражнения 2.2

Рисунки в оригинале: кривые обучения DDPG в HalfCheetah-v2 для реализаций актора-критика с багом и без него, усреднённые по трём случайным сидам (отдельно для TensorFlow и для PyTorch).

Полезно знать

На этой странице решение дано в основном через подробный разбор версии упражнения на TensorFlow. Но проблема в версии на PyTorch по сути та же, и решение у неё такое же.

Баг в коде: версия на TensorFlow

Единственная разница между правильным кодом актора-критика

"""
Актор-критик
"""
def mlp_actor_critic(x, a, hidden_sizes=(400,300), activation=tf.nn.relu,
                     output_activation=tf.tanh, action_space=None):
    act_dim = a.shape.as_list()[-1]
    act_limit = action_space.high[0]
    with tf.variable_scope('pi'):
        pi = act_limit * mlp(x, list(hidden_sizes)+[act_dim], activation, output_activation)
    with tf.variable_scope('q'):
        q = tf.squeeze(mlp(tf.concat([x,a], axis=-1), list(hidden_sizes)+[1], activation, None), axis=1)
    with tf.variable_scope('q', reuse=True):
        q_pi = tf.squeeze(mlp(tf.concat([x,pi], axis=-1), list(hidden_sizes)+[1], activation, None), axis=1)
    return pi, q, q_pi

и кодом актора-критика с багом

"""
Актор-критик с багом
"""
def bugged_mlp_actor_critic(x, a, hidden_sizes=(400,300), activation=tf.nn.relu,
                            output_activation=tf.tanh, action_space=None):
    act_dim = a.shape.as_list()[-1]
    act_limit = action_space.high[0]
    with tf.variable_scope('pi'):
        pi = act_limit * mlp(x, list(hidden_sizes)+[act_dim], activation, output_activation)
    with tf.variable_scope('q'):
        q = mlp(tf.concat([x,a], axis=-1), list(hidden_sizes)+[1], activation, None)
    with tf.variable_scope('q', reuse=True):
        q_pi = mlp(tf.concat([x,pi], axis=-1), list(hidden_sizes)+[1], activation, None)
    return pi, q, q_pi

— это форма тензоров Q-функций. Правильная версия сжимает (squeeze) выходы, чтобы их форма была [batch size], а версия с багом этого не делает, и Q-функции получают форму [batch size, 1].

Баг в коде: версия на PyTorch

В версии упражнения на PyTorch разница практически та же. Правильный код актора-критика при прямом проходе Q-функции сжимает её выход:

"""
Правильная Q-функция
"""
class MLPQFunction(nn.Module):

    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):
        super().__init__()
        self.q = mlp([obs_dim + act_dim] + list(hidden_sizes) + [1], activation)

    def forward(self, obs, act):
        q = self.q(torch.cat([obs, act], dim=-1))
        return torch.squeeze(q, -1) # Критично для правильной формы q.

а версия с багом — нет:

"""
Q-функция с багом
"""
class BuggedMLPQFunction(nn.Module):

    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):
        super().__init__()
        self.q = mlp([obs_dim + act_dim] + list(hidden_sizes) + [1], activation)

    def forward(self, obs, act):
        return self.q(torch.cat([obs, act], dim=-1))

Как это всё портит: версия на TensorFlow

Рассмотрим фрагмент той части кода, которая строит граф вычислений DDPG:

# Цель Беллмана для Q-функции
backup = tf.stop_gradient(r_ph + gamma*(1-d_ph)*q_pi_targ)

# Функции потерь DDPG
pi_loss = -tf.reduce_mean(q_pi)
q_loss = tf.reduce_mean((q-backup)**2)

Именно здесь в игру вступает проблема с формой тензоров. Важно знать, что r_ph и d_ph имеют форму [batch size].

Строка, которая вычисляет цель Беллмана, написана в предположении, что она складывает тензоры одинаковой формы. Но эта строка может сложить и тензоры разной формы, если они совместимы для бродкастинга.

Тензоры форм [batch size] и [batch size, 1] совместимы для бродкастинга, но результат совсем не тот, которого можно было бы ожидать! Посмотрите на пример:

>>> import tensorflow as tf
>>> import numpy as np
>>> x = tf.constant(np.arange(5))
>>> y = tf.constant(np.arange(5).reshape(-1,1))
>>> z1 = x * y
>>> z2 = x + y
>>> z3 = x + z1
>>> x.shape
TensorShape([Dimension(5)])
>>> y.shape
TensorShape([Dimension(5), Dimension(1)])
>>> z1.shape
TensorShape([Dimension(5), Dimension(5)])
>>> z2.shape
TensorShape([Dimension(5), Dimension(5)])
>>> sess = tf.InteractiveSession()
>>> sess.run(z1)
array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16]])
>>> sess.run(z2)
array([[0, 1, 2, 3, 4],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [4, 5, 6, 7, 8]])
>>> sess.run(z3)
array([[ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16],
       [ 0,  5, 10, 15, 20]])

Сложение или умножение тензора формы [5] на тензор формы [5,1] даёт тензор формы [5,5]!

Если не сжать Q-функции, q_pi_targ получает форму [batch size, 1], и цель Беллмана — а с ней и вся функция потерь Q — полностью портится.

Ошибка бродкастинга 1: (1 - d_ph) * q_pi_targ становится тензором [batch size, batch size] с внешним произведением маски на значения целевой Q-сети.

Ошибка бродкастинга 2: затем r_ph трактуется как вектор-строка и прибавляется к каждой строке (1 - d_ph) * q_pi_targ по отдельности.

Ошибка бродкастинга 3: q_loss зависит от q - backup, где снова неправильный бродкастинг между q (форма [batch size, 1]) и backup (форма [batch size, batch size]).

Математически: пусть $q$, $q'$, $r$, $d$ — векторы Q-значений, целевых Q-значений, наград и флагов завершения для данного батча из $n$ элементов. Правильная цель Беллмана —

$$z_i = r_i + \gamma (1-d_i) q'_i,$$

а правильная функция потерь —

$$\frac{1}{n} \sum_{i=1}^n (q_i - z_i)^2.$$

Но из-за этих ошибок вычисляется матрица целей

$$z_{ij} = r_j + \gamma (1-d_j) q'_i,$$

и испорченная функция потерь

$$\frac{1}{n^2} \sum_{i=1}^n \sum_{j=1}^n (q_j - z_{ij})^2.$$

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

Как это всё портит: версия на PyTorch

Ровно те же фокусы с бродкастингом, что и в версии на TensorFlow. Посмотрите эту заметку в документации PyTorch.

Мини-демонстрация на PyTorch (добавлено переводчиком). Батч из трёх переходов, запускается за доли секунды:

import torch
r = torch.tensor([1., 2., 3.])          # награды, форма [3]
d = torch.tensor([0., 0., 1.])          # флаги «эпизод закончился», форма [3]
q_targ_ok = torch.tensor([10., 20., 30.])        # правильно: форма [3]
q_targ_bug = q_targ_ok.unsqueeze(-1)              # баг: форма [3, 1]
gamma = 0.9
print((r + gamma * (1 - d) * q_targ_ok).shape)   # torch.Size([3])
print((r + gamma * (1 - d) * q_targ_bug).shape)  # torch.Size([3, 3]) — тихий бродкастинг!
print(r + gamma * (1 - d) * q_targ_bug)

Вывод:

torch.Size([3])
torch.Size([3, 3])
tensor([[10., 11.,  3.],
        [19., 20.,  3.],
        [28., 29.,  3.]])

Строка $i$ — это $z_{ij} = r_j + \gamma(1-d_j)q'_i$: например, первая строка $[1 + 9,\ 2 + 9,\ 3 + 0]$. Ни одной ошибки, ни одного предупреждения — только неверные числа. Защита простая: assert q.shape == r.shape в функции потерь.

Что дальше

• Перед упражнениями 1.2 и 1.3 перечитайте VPG и PPO и TRPO, DDPG, TD3, SAC.
• Практика RL в нашем курсе на современном стеке — лабораторные PPO и DQN на CartPole в уроках 6.0.2 и 6.0.3.