Spinning Up, часть 3: введение в оптимизацию политики

Где в курсе то же самое проще

Тот же вывод «на пальцах», с числовым примером на двух действиях, — в уроке 6.0.3, разделы «Градиент политики на пальцах», «REINFORCE и reward-to-go» и «Baseline и преимущество». Для языковых моделей (действие — токен, траектория — ответ) — в уроке 6.3.1 «Основы RL для языковых моделей». Здесь те же три результата доказаны строго и связаны со 128 строками кода.

Как читать перевод

Обозначения оригинала сохранены: $\underset{\tau \sim \pi_\theta}{\mathrm{E}}[\,\cdot\,]$ — матожидание по траекториям, которые порождает политика $\pi_\theta$ (в курсе — $\mathbb{E}_{\tau \sim \pi_\theta}[\,\cdot\,]$). Абзацы «Словами» добавлены переводчиком. Комментарии в коде переведены; номера строк в тексте — как в файле оригинала.

В этом разделе мы обсудим математические основы алгоритмов оптимизации политики и свяжем материал с примерами кода. Мы разберём три ключевых результата теории градиентов политикиградиент политикиEN policy gradient∇J = E[r · ∇ log π(y|x)] — основа REINFORCE, PPO и GRPO.:

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

В конце мы соберём эти результаты вместе и опишем выражение градиента политики через преимущество — ту версию, которую мы используем в нашей реализации Vanilla Policy Gradient.

Вывод простейшего градиента политики

Рассмотрим случай стохастической параметризованной политики $\pi_{\theta}$. Мы хотим максимизировать ожидаемую отдачу $J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[R(\tau)\right]$. Для этого вывода возьмём в качестве $R(\tau)$ конечногоризонтную недисконтированную отдачу, но вывод для бесконечногоризонтной дисконтированной отдачи почти такой же.

Мы хотели бы оптимизировать политику градиентным подъёмом, например

$$\theta_{k+1} = \theta_k + \alpha \left. \nabla_{\theta} J(\pi_{\theta}) \right|_{\theta_k}.$$

Словами: новые параметры = старые параметры + шаг $\alpha$ в направлении, где ожидаемая отдача растёт быстрее всего; градиент вычислен в точке $\theta_k$ (вертикальная черта с индексом).

Градиент качества политики $\nabla_{\theta} J(\pi_{\theta})$ называется градиентом политики (policy gradient), а алгоритмы, оптимизирующие политику таким способом, — алгоритмами градиента политики. (Примеры — Vanilla Policy Gradient и TRPO. PPO часто тоже называют алгоритмом градиента политики, хотя это не совсем точно.)

Чтобы на деле использовать этот алгоритм, нужно выражение для градиента политики, которое можно посчитать численно. Это делается в два шага: 1) вывести аналитический градиент качества политики, который, как окажется, имеет вид матожидания, а затем 2) построить выборочную оценку этого матожидания, которую можно посчитать по данным из конечного числа шагов взаимодействия агента со средой.

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

Начнём с нескольких фактов, полезных для вывода аналитического градиента.

1. Вероятность траектории. Вероятность траектории $\tau = (s_0, a_0, ..., s_{T+1})$ при условии, что действия выбирает $\pi_{\theta}$, равна

$$P(\tau|\theta) = \rho_0 (s_0) \prod_{t=0}^{T} P(s_{t+1}|s_t, a_t) \pi_{\theta}(a_t |s_t).$$

2. Трюк с логарифмической производной (log-derivative trick). Он основан на простом правиле из матанализа: производная $\log x$ по $x$ равна $1/x$. Если его переписать и объединить с правилом дифференцирования сложной функции, получим:

$$\nabla_{\theta} P(\tau | \theta) = P(\tau | \theta) \nabla_{\theta} \log P(\tau | \theta).$$

Словами: $\nabla \log P = \nabla P / P$; умножив обе части на $P$, получаем формулу. Она позволяет превратить градиент вероятности обратно в матожидание.

3. Логарифм вероятности траектории. Логарифм вероятности траектории равен просто

$$\log P(\tau|\theta) = \log \rho_0 (s_0) + \sum_{t=0}^{T} \bigg( \log P(s_{t+1}|s_t, a_t) + \log \pi_{\theta}(a_t |s_t)\bigg).$$

4. Градиенты функций среды. Среда не зависит от $\theta$, поэтому градиенты $\rho_0(s_0)$, $P(s_{t+1}|s_t, a_t)$ и $R(\tau)$ равны нулю.

5. Градиент логарифма вероятности траектории. Поэтому градиент логарифма вероятности траектории равен

$$\begin{aligned} \nabla_{\theta} \log P(\tau | \theta) &= \cancel{\nabla_{\theta} \log \rho_0 (s_0)} + \sum_{t=0}^{T} \bigg( \cancel{\nabla_{\theta} \log P(s_{t+1}|s_t, a_t)} + \nabla_{\theta} \log \pi_{\theta}(a_t |s_t)\bigg) \\ &= \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t). \end{aligned}$$

Собрав всё вместе, получаем:

Вывод базового градиента политики

$$\begin{aligned} \nabla_{\theta} J(\pi_{\theta}) &= \nabla_{\theta} \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[R(\tau)\right] & \\ &= \nabla_{\theta} \int_{\tau} P(\tau|\theta) R(\tau) & \text{раскрываем матожидание} \\ &= \int_{\tau} \nabla_{\theta} P(\tau|\theta) R(\tau) & \text{вносим градиент под интеграл} \\ &= \int_{\tau} P(\tau|\theta) \nabla_{\theta} \log P(\tau|\theta) R(\tau) & \text{трюк с логарифмической производной} \\ &= \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log P(\tau|\theta) R(\tau)\right] & \text{возвращаемся к форме матожидания} \\ \therefore \nabla_{\theta} J(\pi_{\theta}) &= \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau)\right] & \text{выражение для градиента логарифма вероятности} \end{aligned}$$

Это матожидание, а значит, его можно оценить выборочным средним. Если собрать набор траекторий $\mathcal{D} = \{\tau_i\}_{i=1,...,N}$, где каждая траектория получена, когда агент действовал в среде по политике $\pi_{\theta}$, градиент политики можно оценить так:

$$\hat{g} = \frac{1}{|\mathcal{D}|} \sum_{\tau \in \mathcal{D}} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau),$$

где $|\mathcal{D}|$ — число траекторий в $\mathcal{D}$ (здесь $N$).

Словами: сыграйте $N$ эпизодов текущей политикой; для каждого шага каждого эпизода возьмите градиент логарифма вероятности выбранного действия и умножьте на отдачу всего эпизода; всё сложите и поделите на $N$. Действия из эпизодов с большой отдачей станут вероятнее. Модель среды при этом не нужна.

Это последнее выражение — простейшая версия вычислимого выражения, которое мы искали. Если политика представлена так, что можно посчитать $\nabla_{\theta} \log \pi_{\theta}(a|s)$, и если мы можем запускать политику в среде, чтобы собрать набор траекторий, то мы можем вычислить градиент политики и сделать шаг обновления.

Реализация простейшего градиента политики

Короткая реализация этой простой версии алгоритма градиента политики на PyTorch лежит в spinup/examples/pytorch/pg_math/1_simple_pg.py (её можно посмотреть и на GitHub). В ней всего 128 строк, поэтому очень советуем прочитать её внимательно. Мы не будем разбирать весь код здесь, но выделим и объясним несколько важных фрагментов.

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

Раньше этот раздел был написан с примером на Tensorflow. Старый вариант можно найти здесь.

1. Создание сети политики.

# строки 30–40 файла 1_simple_pg.py

    # создаём ядро сети политики
    logits_net = mlp(sizes=[obs_dim]+hidden_sizes+[n_acts])

    # функция, вычисляющая распределение действий
    def get_policy(obs):
        logits = logits_net(obs)
        return Categorical(logits=logits)

    # функция выбора действия (выдаёт целые действия, сэмплированные из политики)
    def get_action(obs):
        return get_policy(obs).sample().item()

Этот блок строит модули и функции для категориальной политики на полносвязной нейросети. (Чтобы освежить материал, см. раздел «Стохастические политики» в части 1.) Выход модуля logits_net можно использовать, чтобы построить логарифмы вероятностей и вероятности действий, а функция get_action сэмплирует действия по вероятностям, вычисленным из логитов. (Замечание: эта конкретная функция get_action предполагает, что на вход подаётся только одно наблюдение obs и на выходе получается одно целое действие. Поэтому в ней используется .item() — метод, который достаёт содержимое тензора из одного элемента.)

Большую часть работы в этом примере делает объект Categorical в строке 36. Это объект Distribution из PyTorch, который заключает в себе математические функции, связанные с распределениями вероятностей. В частности, у него есть метод сэмплирования из распределения (мы используем его в строке 40) и метод вычисления логарифмов вероятностей заданных сэмплов (его мы используем позже). Распределения PyTorch очень полезны для RL, поэтому загляните в их документацию, чтобы понять, как они работают.

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

Дружеское напоминание! Когда мы говорим, что у категориального распределения есть «логиты», мы имеем в виду, что вероятности исходов получаются функцией softmax от логитов. То есть вероятность действия $j$ при категориальном распределении с логитами $x_j$ равна:

$$p_j = \frac{\exp(x_j)}{\sum_{i} \exp(x_i)}$$

2. Создание функции потерь.

# строки 42–45 файла 1_simple_pg.py

    # функция потерь, градиент которой на правильных данных равен градиенту политики
    def compute_loss(obs, act, weights):
        logp = get_policy(obs).log_prob(act)
        return -(logp * weights).mean()

В этом блоке мы строим функцию «потерь» для алгоритма градиента политики. Если подставить в неё правильные данные, её градиент будет равен градиенту политики. Правильные данные — это набор кортежей (состояние, действие, вес), собранных во время действий по текущей политике, где вес для пары «состояние—действие» — отдача эпизода, которому эта пара принадлежит. (Хотя, как мы покажем в следующих подразделах, в качестве веса можно подставлять и другие величины, и всё будет работать правильно.)

Словами: минус нужен потому, что оптимизаторы PyTorch минимизируют, а нам нужно максимизировать $\sum \log\pi \cdot R$. Градиентный спуск по $-L$ — это градиентный подъём по $L$.

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

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

1. Распределение данных зависит от параметров. Функцию потерь обычно определяют на фиксированном распределении данных, которое не зависит от оптимизируемых параметров. Здесь это не так: данные нужно сэмплировать самой свежей политикой.

2. Она не измеряет качество. Функция потерь обычно оценивает ту метрику качества, которая нам важна. Здесь нам важна ожидаемая отдача $J(\pi_{\theta})$, но наша «функция потерь» совсем её не приближает, даже в среднем. Эта «функция потерь» полезна нам только потому, что, если её вычислить при текущих параметрах на данных, порождённых текущими параметрами, её градиент равен градиенту качества со знаком минус.

Но после первого шага градиентного спуска связи с качеством больше нет. Это значит, что минимизация такой «функции потерь» на данном батче данных никак не гарантирует улучшения ожидаемой отдачи. Можно загнать эту функцию потерь в $-\infty$, а качество политики при этом может рухнуть; на самом деле обычно так и происходит. Иногда исследователь глубокого RL опишет такой исход как «переобучение» политики на батч данных. Описание наглядное, но понимать его буквально не стоит: речь не об ошибке обобщения.

Мы обращаем на это внимание, потому что специалисты по машинному обучению часто считают функцию потерь полезным сигналом во время обучения: «если loss падает, всё хорошо». Для градиентов политики эта интуиция неверна, и следить нужно только за средней отдачей. Функция потерь не значит ничего.

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

Способ, которым здесь получается тензор logp, — вызов метода log_prob у объекта Categorical из PyTorch, — может потребовать доработки для других видов распределений.

Например, если вы используете нормальное распределение (для диагональной гауссовой политики), вызов policy.log_prob(act) вернёт тензор с отдельными логарифмами вероятностей для каждой компоненты каждого векторного действия. То есть на вход подаётся тензор формы (batch, act_dim), и на выходе тоже тензор формы (batch, act_dim), а для функции потерь RL нужен тензор формы (batch,). В этом случае логарифмы вероятностей компонент нужно сложить, чтобы получить логарифмы вероятностей действий. То есть вычислить:

logp = get_policy(obs).log_prob(act).sum(axis=-1)  # сумма по координатам действия

3. Запуск одной эпохи обучения.

# строки 50–111 файла 1_simple_pg.py

    # обучение политики
    def train_one_epoch():
        # пустые списки для логирования
        batch_obs = []          # наблюдения
        batch_acts = []         # действия
        batch_weights = []      # веса R(tau) в градиенте политики
        batch_rets = []         # для измерения отдач эпизодов
        batch_lens = []         # для измерения длин эпизодов

        # сбрасываем переменные, относящиеся к эпизоду
        obs = env.reset()       # первое наблюдение — из начального распределения
        done = False            # сигнал среды о том, что эпизод закончился
        ep_rews = []            # список наград, накопленных за эпизод

        # отрисовываем первый эпизод каждой эпохи
        finished_rendering_this_epoch = False

        # собираем опыт: действуем в среде текущей политикой
        while True:

            # отрисовка
            if (not finished_rendering_this_epoch) and render:
                env.render()

            # сохраняем наблюдение
            batch_obs.append(obs.copy())

            # действуем в среде
            act = get_action(torch.as_tensor(obs, dtype=torch.float32))
            obs, rew, done, _ = env.step(act)

            # сохраняем действие и награду
            batch_acts.append(act)
            ep_rews.append(rew)

            if done:
                # если эпизод закончился, записываем сведения о нём
                ep_ret, ep_len = sum(ep_rews), len(ep_rews)
                batch_rets.append(ep_ret)
                batch_lens.append(ep_len)

                # вес каждого logprob(a|s) равен R(tau)
                batch_weights += [ep_ret] * ep_len

                # сбрасываем переменные эпизода
                obs, done, ep_rews = env.reset(), False, []

                # в этой эпохе больше не отрисовываем
                finished_rendering_this_epoch = True

                # выходим из цикла сбора опыта, если его набралось достаточно
                if len(batch_obs) > batch_size:
                    break

        # делаем один шаг обновления по градиенту политики
        optimizer.zero_grad()
        batch_loss = compute_loss(obs=torch.as_tensor(batch_obs, dtype=torch.float32),
                                  act=torch.as_tensor(batch_acts, dtype=torch.int32),
                                  weights=torch.as_tensor(batch_weights, dtype=torch.float32)
                                  )
        batch_loss.backward()
        optimizer.step()
        return batch_loss, batch_rets, batch_lens

Функция train_one_epoch() выполняет одну «эпоху» градиента политики, которую мы определяем как

1. шаг сбора опыта (строки 67–102), где агент действует в среде несколько эпизодов по самой свежей политике, и следующий за ним
2. один шаг обновления по градиенту политики (строки 104–111).

Главный цикл алгоритма просто раз за разом вызывает train_one_epoch().

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

Если вы ещё не знакомы с оптимизацией в PyTorch, обратите внимание на шаблон одного шага градиентного спуска в строках 104–111. Сначала очищаются буферы градиентов. Затем вычисляется функция потерь. Затем по ней делается обратный проход — он накапливает свежие градиенты в буферах. Наконец, оптимизатор делает шаг.

Лемма о матожидании градиента логарифма вероятности

В этом подразделе мы выведем промежуточный результат, который широко используется во всей теории градиентов политики. Мы будем называть его леммой о матожидании градиента логарифма вероятности (Expected Grad-Log-Prob, EGLP).¹

Лемма EGLP. Пусть $P_{\theta}$ — параметризованное распределение вероятностей случайной величины $x$. Тогда:

$$\underset{x \sim P_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log P_{\theta}(x)\right] = 0.$$

Доказательство

Вспомним, что все распределения вероятностей нормированы:

$$\int_x P_{\theta}(x) = 1.$$

Возьмём градиент от обеих частей условия нормировки:

$$\nabla_{\theta} \int_x P_{\theta}(x) = \nabla_{\theta} 1 = 0.$$

С помощью трюка с логарифмической производной получаем:

$$\begin{aligned} 0 &= \nabla_{\theta} \int_x P_{\theta}(x) \\ &= \int_x \nabla_{\theta} P_{\theta}(x) \\ &= \int_x P_{\theta}(x) \nabla_{\theta} \log P_{\theta}(x) \\ \therefore 0 &= \underset{x \sim P_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log P_{\theta}(x)\right]. \end{aligned}$$

Словами: если усреднить направление «сделать сэмпл вероятнее» по всем сэмплам самого распределения, получится ноль: сумма вероятностей всегда равна 1, и поднять вероятность одних исходов можно только за счёт других. Пример: монетка с вероятностью орла $p = 0.3$. Производная $\log p$ по $p$ равна $1/p$, производная $\log(1-p)$ равна $-1/(1-p)$. Среднее: $0.3 \cdot \tfrac{1}{0.3} + 0.7 \cdot \left(-\tfrac{1}{0.7}\right) = 1 - 1 = 0$.

Не позволяйте прошлому отвлекать вас

Посмотрим на наше последнее выражение для градиента политики:

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau)\right].$$

Шаг по этому градиенту поднимает логарифмы вероятностей каждого действия пропорционально $R(\tau)$ — сумме всех когда-либо полученных наград. Но в этом мало смысла.

Агент должен подкреплять действия только на основе их последствий. Награды, полученные до действия, никак не говорят о том, насколько хорошим оно было: важны только награды, полученные после.

Оказывается, эта интуиция проявляется и в математике: можно показать, что градиент политики можно записать и так:

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1})\right].$$

В этой форме действия подкрепляются только наградами, полученными после того, как они совершены.

Мы будем называть эту форму «градиентом политики с reward-to-go», потому что сумма наград после некоторой точки траектории,

$$\hat{R}_t \doteq \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1}),$$

называется reward-to-go («оставшаяся награда») от этой точки, а это выражение градиента политики зависит от reward-to-go пар «состояние—действие».

Словами: $\hat{R}_t$ — сумма наград от шага $t$ до конца эпизода (знак $\doteq$ значит «по определению равно»). Пример: награды эпизода $[1, 0, 2]$ дают reward-to-go $[3, 2, 2]$, а отдача всего эпизода $R(\tau) = 3$ одинакова для всех шагов. В курсе эту величину обозначают $G_t$ и называют отдачей с шага $t$.

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

Но чем это лучше? Ключевая проблема градиентов политики — сколько нужно выборочных траекторий, чтобы получить их оценку с малой дисперсией. В формуле, с которой мы начали, были слагаемые, подкрепляющие действия пропорционально прошлым наградам. Все они имели нулевое среднее, но ненулевую дисперсию: в итоге они просто добавляли шум в выборочные оценки градиента политики. Убрав их, мы уменьшаем число нужных выборочных траекторий.

(Необязательное) доказательство этого утверждения приведено в приложении А; в конечном счёте оно опирается на лемму EGLP.

Реализация градиента политики с reward-to-go

Короткая реализация градиента политики с reward-to-go на PyTorch лежит в spinup/examples/pytorch/pg_math/2_rtg_pg.py (её можно посмотреть и на GitHub).

По сравнению с 1_simple_pg.py изменилось только одно: теперь в функции потерь другие веса. Изменение в коде совсем небольшое: мы добавляем новую функцию и меняем две другие строки. Новая функция:

# строки 17–22 файла 2_rtg_pg.py

def reward_to_go(rews):
    n = len(rews)
    rtgs = np.zeros_like(rews)
    for i in reversed(range(n)):                             # идём с конца эпизода
        rtgs[i] = rews[i] + (rtgs[i+1] if i+1 < n else 0)    # R̂_i = r_i + R̂_{i+1}
    return rtgs

А старые строки 91–92:

                    # вес каждого logprob(a|s) равен R(tau)
                    batch_weights += [ep_ret] * ep_len

мы меняем на:

                    # вес каждого logprob(a_t|s_t) равен reward-to-go от шага t
                    batch_weights += list(reward_to_go(ep_rews))

(в файле 2_rtg_pg.py это строки 98–99).

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

Раньше этот раздел был написан с примером на Tensorflow. Старый вариант можно найти здесь.

Базовые линии в градиентах политики

Непосредственное следствие леммы EGLP: для любой функции $b$, которая зависит только от состояния,

$$\underset{a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t|s_t) b(s_t)\right] = 0.$$

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

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \left(\sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1}) - b(s_t)\right)\right].$$

Любую функцию $b$, используемую таким образом, называют базовой линиейбазовая линия (baseline)EN baselineЧисло, которое вычитают из награды; не смещает градиент политики, но снижает его дисперсию. (baseline).

Самый распространённый выбор базовой линии — функция ценности on-policy $V^{\pi}(s_t)$. Напомним: это средняя отдача, которую получает агент, если начинает в состоянии $s_t$ и затем до конца жизни действует по политике $\pi$.

На практике выбор $b(s_t) = V^{\pi}(s_t)$ даёт желаемый эффект — уменьшает дисперсию выборочной оценки градиента политики. В результате политика учится быстрее и стабильнее. Этот выбор привлекателен и с концептуальной точки зрения: он выражает интуицию, что если агент получил то, чего ожидал, он должен «чувствовать» себя нейтрально.

Словами: вес действия теперь — «насколько исход оказался лучше ожидаемого». Если из состояния обычно набирают $V = 10$, а в этот раз reward-to-go 12, вес $+2$ (действие поощряется); если 8 — вес $-2$ (действие становится реже). Без базовой линии оба веса были бы положительными, и разницу пришлось бы «вычитать» шумом многих эпизодов.

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

На практике $V^{\pi}(s_t)$ нельзя вычислить точно, поэтому её приходится аппроксимировать. Обычно это делают нейросетью $V_{\phi}(s_t)$, которую обновляют одновременно с политикой (так, чтобы сеть ценности всегда аппроксимировала функцию ценности самой свежей политики).

Самый простой способ учить $V_{\phi}$, используемый в большинстве реализаций алгоритмов оптимизации политики (включая VPG, TRPO, PPO и A2C), — минимизировать среднеквадратичную ошибку:

$$\phi_k = \arg \min_{\phi} \underset{s_t, \hat{R}_t \sim \pi_k}{\mathrm{E}}\left[\left( V_{\phi}(s_t) - \hat{R}_t \right)^2\right],$$

где $\pi_k$ — политика на эпохе $k$. Это делается одним или несколькими шагами градиентного спуска, начиная с предыдущих параметров ценности $\phi_{k-1}$.

Другие формы градиента политики

Мы видели, что градиент политики имеет общий вид

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Phi_t\right],$$

где $\Phi_t$ может быть любым из

$$\Phi_t = R(\tau),$$

или

$$\Phi_t = \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1}),$$

или

$$\Phi_t = \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1}) - b(s_t).$$

Все эти варианты дают одно и то же матожидание градиента политики, хотя дисперсии у них разные. Оказывается, есть ещё два допустимых выбора весов $\Phi_t$, которые важно знать.

1. Функция ценности действия on-policy. Выбор

$$\Phi_t = Q^{\pi_{\theta}}(s_t, a_t)$$

тоже допустим. (Необязательное) доказательство этого утверждения — в приложении Б.

2. Функция преимущества. Вспомним, что преимущество действия, определённое как $A^{\pi}(s_t,a_t) = Q^{\pi}(s_t,a_t) - V^{\pi}(s_t)$, описывает, насколько оно в среднем лучше или хуже других действий (относительно текущей политики). Выбор

$$\Phi_t = A^{\pi_{\theta}}(s_t, a_t)$$

тоже допустим. Доказательство: это то же самое, что взять $\Phi_t = Q^{\pi_{\theta}}(s_t, a_t)$ и затем вычесть базовую линию — функцию ценности, что мы всегда вправе сделать.

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

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

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

Более подробное изложение этой темы — в статье о Generalized Advantage Estimation (GAEобобщённая оценка преимущества (GAE)EN generalized advantage estimation (GAE)Schulman и др., 2015. Смешивает оценки преимущества разной длины с коэффициентом λ.), где в разделах с предпосылками подробно разобраны разные варианты выбора $\Phi_t$.

Дальше статья описывает GAE — метод аппроксимации функции преимущества в алгоритмах оптимизации политики, который очень широко используется. Например, его используют реализации VPG, TRPO и PPO в Spinning Up. Поэтому настоятельно советуем её изучить.

Итоги

В этой главе мы описали базовую теорию методов градиента политики и связали некоторые ранние результаты с примерами кода. Заинтересованному студенту стоит продолжить, изучив, как более поздние результаты (базовая линия — функция ценности — и запись градиентов политики через преимущество) воплощаются в реализации Vanilla Policy Gradient в Spinning Up.

---

Приложение А. Доказательство для раздела «Не позволяйте прошлому отвлекать вас»

Перевод страницы «Extra Material: Proof for Don’t Let the Past Distract You».

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

Раскроем $R(\tau)$ в выражении для простейшего градиента политики:

$$\begin{aligned} \nabla_{\theta} J(\pi_{\theta}) &= \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau)\right] \\ &= \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \sum_{t'=0}^T R(s_{t'}, a_{t'}, s_{t'+1})\right] \\ &= \sum_{t=0}^{T} \sum_{t'=0}^T \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(s_{t'}, a_{t'}, s_{t'+1})\right], \end{aligned}$$

и рассмотрим слагаемое

$$\underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(s_{t'}, a_{t'}, s_{t'+1})\right].$$

Мы покажем, что в случае $t' < t$ (награда получена раньше подкрепляемого действия) это слагаемое равно нулю. Это полное доказательство исходного утверждения: после удаления слагаемых с $t' < t$ из выражения остаётся форма градиента политики с reward-to-go, что и требовалось:

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1})\right]$$

1. Используем маргинальное распределение. Чтобы двигаться дальше, нужно разложить матожидание в $\underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right]$. Это матожидание по траекториям, но выражение под ним касается лишь нескольких состояний и действий: $s_t$, $a_t$, $s_{t'}$, $a_{t'}$ и $s_{t'+1}$. Поэтому при вычислении матожидания нам важно только маргинальное распределение этих случайных величин.

Выводим:

$$\begin{aligned} \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] &= \int_{\tau} P(\tau|\pi_{\theta}) f(t,t') \\ &= \int_{s_t, a_t, s_{t'}, a_{t'}, s_{t'+1}} P(s_t, a_t, s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta}) f(t,t') \\ &= \underset{s_t, a_t, s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right]. \end{aligned}$$

2. Цепное правило вероятностей. Совместные распределения можно выразить через условные и маргинальные вероятности по цепному правилу: $P(A,B) = P(B|A) P(A)$. Здесь мы используем это правило, чтобы вычислить

$$P(s_t, a_t, s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta}) = P(s_t, a_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) P(s_{t'}, a_{t'}, s_{t'+1} | \pi_{\theta})$$

3. Разделяем матожидания по нескольким случайным величинам. Если есть матожидание по двум случайным величинам $A$ и $B$, его можно разбить на внутреннее и внешнее, где внутреннее считает величину из внешнего матожидания константой. Возможность такого разбиения опирается на цепное правило вероятностей. Математически:

$$\begin{aligned} \underset{A,B}{\mathrm{E}}\left[f(A,B)\right] &= \int_{A,B} P(A,B) f(A,B) \\ &= \int_{A} \int_B P(B|A) P(A) f(A,B) \\ &= \int_A P(A) \int_B P(B|A) f(A,B) \\ &= \int_A P(A) \underset{B}{\mathrm{E}}\left[f(A,B) \Big| A\right] \\ &= \underset{A}{\mathrm{E}}\left[\underset{B}{\mathrm{E}}\left[f(A,B) \Big| A\right] \right] \end{aligned}$$

Значит, матожидание по $s_t, a_t, s_{t'}, a_{t'}, s_{t'+1}$ можно записать как

$$\begin{aligned} \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] &= \underset{s_t, a_t, s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] \\ &= \underset{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[\underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t') \Big| s_{t'}, a_{t'}, s_{t'+1}\right]\right] \end{aligned}$$

4. Константы можно выносить за знак матожидания. Если слагаемое под матожиданием постоянно относительно величины, по которой берётся матожидание, его можно вынести. Пример: снова матожидание по двум случайным величинам $A$ и $B$, но теперь $f(A,B) = h(A) g(B)$. Тогда, используя прежний результат:

$$\begin{aligned} \underset{A,B}{\mathrm{E}}\left[f(A,B)\right] &= \underset{A}{\mathrm{E}}\left[\underset{B}{\mathrm{E}}\left[f(A,B) \Big| A\right]\right] \\ &= \underset{A}{\mathrm{E}}\left[\underset{B}{\mathrm{E}}\left[h(A) g(B) \Big| A\right]\right]\\ &= \underset{A}{\mathrm{E}}\left[h(A) \underset{B}{\mathrm{E}}\left[g(B) \Big| A\right]\right]. \end{aligned}$$

Функция под нашим матожиданием раскладывается именно так, и мы можем записать:

$$\begin{aligned} \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] &= \underset{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[\underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t') \Big| s_{t'}, a_{t'}, s_{t'+1}\right]\right] \\ &= \underset{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[\underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(s_{t'}, a_{t'}, s_{t'+1}) \Big| s_{t'}, a_{t'}, s_{t'+1}\right]\right] \\ &= \underset{s_{t'}, a_{t'}, s_{t'+1} \sim \pi_{\theta}}{\mathrm{E}}\left[R(s_{t'}, a_{t'}, s_{t'+1}) \underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big| s_{t'}, a_{t'}, s_{t'+1}\right]\right]. \end{aligned}$$

5. Применяем лемму EGLP. Последний шаг доказательства опирается на лемму EGLP. Сейчас нас будет интересовать только самое внутреннее матожидание:

$$\underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big| s_{t'}, a_{t'}, s_{t'+1}\right] = \int_{s_t, a_t} P(s_t, a_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) \nabla_{\theta} \log \pi_{\theta}(a_t |s_t).$$

Теперь нужно различить два случая: $t' < t$ — награда получена до действия, и $t' \geq t$ — не до.

Случай первый: награда до действия. Если $t' < t$, то условные вероятности действий $a_t$ задаёт политика:

$$P(s_t, a_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) = \pi_{\theta}(a_t | s_t) P(s_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}),$$

и самое внутреннее матожидание можно разложить дальше:

$$\begin{aligned} \underset{s_t, a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big| s_{t'}, a_{t'}, s_{t'+1}\right] &= \int_{s_t, a_t} P(s_t, a_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \\ &= \int_{s_t} P(s_t | \pi_{\theta}, s_{t'}, a_{t'}, s_{t'+1}) \int_{a_t} \pi_{\theta}(a_t | s_t) \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \\ &= \underset{s_t \sim \pi_{\theta}}{\mathrm{E}}\left[ \underset{a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big| s_t \right] \Big| s_{t'}, a_{t'}, s_{t'+1}\right]. \end{aligned}$$

Лемма EGLP говорит, что

$$\underset{a_t \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big| s_t \right] = 0,$$

откуда заключаем, что при $t' < t$ выполняется $\underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[f(t,t')\right] = 0$.

Случай второй: награда после действия. А что со случаем $t' \geq t$? Почему та же логика к нему не применима? В этом случае условные вероятности $a_t$ нельзя разложить так же, потому что мы ставим условие на будущее. Представьте себе так: каждый день утром вы выбираете между пробежкой и тем, чтобы пораньше пойти на работу, и шансы 50 на 50. Если поставить условие на будущее, в котором вы пришли на работу рано, каковы шансы, что вы бегали? Ясно, что не бегали. А если ставить условие на прошлое — до того, как вы приняли решение, — каковы шансы, что вы потом побежите? Снова 50 на 50.

Поэтому в случае $t' \geq t$ условное распределение действий $a_t$ — не $\pi(a_t|s_t)$, и лемма EGLP неприменима.

Приложение Б. Доказательство для Q-функции в формуле градиента политики

Перевод страницы «Extra Material: Proof for Using Q-Function in Policy Gradient Formula».

В этом разделе мы покажем, что

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \Big( \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \Big) Q^{\pi_{\theta}}(s_t, a_t)\right],$$

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

Доказательство опирается на закон повторного матожидания (law of iterated expectations). Сначала перепишем выражение для градиента политики, начиная с формы с reward-to-go (для краткости используем обозначение $\hat{R}_t = \sum_{t'=t}^T R(s_{t'}, a_{t'}, s_{t'+1})$):

$$\begin{aligned} \nabla_{\theta} J(\pi_{\theta}) &= \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \hat{R}_t\right] \\ &= \sum_{t=0}^{T} \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[\nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \hat{R}_t\right] \end{aligned}$$

Определим $\tau_{:t} = (s_0, a_0, ..., s_t, a_t)$ как траекторию до момента $t$, а $\tau_{t:}$ — как остаток траектории после него. По закону повторного матожидания предыдущее выражение можно разбить так:

$$\nabla_{\theta} J(\pi_{\theta}) = \sum_{t=0}^{T} \underset{\tau_{:t} \sim \pi_{\theta}}{\mathrm{E}}\left[ \underset{\tau_{t:} \sim \pi_{\theta}}{\mathrm{E}}\left[ \left. \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \hat{R}_t \right| \tau_{:t}\right]\right]$$

Градиент логарифма вероятности постоянен относительно внутреннего матожидания (потому что он зависит от $s_t$ и $a_t$, которые внутреннее матожидание считает фиксированными в $\tau_{:t}$), поэтому его можно вынести, и остаётся:

$$\nabla_{\theta} J(\pi_{\theta}) = \sum_{t=0}^{T} \underset{\tau_{:t} \sim \pi_{\theta}}{\mathrm{E}}\left[ \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) \underset{\tau_{t:} \sim \pi_{\theta}}{\mathrm{E}}\left[ \left. \hat{R}_t \right| \tau_{:t}\right]\right]$$

В марковских процессах принятия решений будущее зависит только от последнего состояния и действия. Поэтому внутреннее матожидание — матожидание по будущему при условии всего прошлого (всего до момента $t$) — равно тому же матожиданию при условии только последнего шага (только $(s_t,a_t)$):

$$\underset{\tau_{t:} \sim \pi_{\theta}}{\mathrm{E}}\left[ \left. \hat{R}_t \right| \tau_{:t}\right] = \underset{\tau_{t:} \sim \pi_{\theta}}{\mathrm{E}}\left[ \left. \hat{R}_t \right| s_t, a_t\right],$$

а это и есть определение $Q^{\pi_{\theta}}(s_t, a_t)$: ожидаемая отдача при старте из состояния $s_t$ и действия $a_t$, если до конца траектории действовать on-policy.

Результат следует немедленно.

---

Приложение В (от переводчика). Полный пример на современном Gymnasium

Файл 1_simple_pg.py написан под старый gym 0.15: env.reset() возвращает только наблюдение, а env.step() — четыре значения. В современном Gymnasium reset() возвращает пару (наблюдение, info), а step() — пять значений, где конец эпизода разделён на terminated (агент проиграл или выиграл) и truncated (кончилось время). Ниже — оба примера (простейший и reward-to-go) в одном файле с этими правками. Изменённые места помечены комментарием «было:». Нужны torch, numpy и gymnasium (pip install gymnasium), GPU не нужен.

# simple_pg_gymnasium.py — 1_simple_pg.py и 2_rtg_pg.py из Spinning Up на Gymnasium
import numpy as np
import torch
import torch.nn as nn
from torch.distributions.categorical import Categorical
from torch.optim import Adam
import gymnasium as gym                      # было: import gym
from gymnasium.spaces import Discrete, Box


def mlp(sizes, activation=nn.Tanh, output_activation=nn.Identity):
    # Собираем полносвязную сеть: Linear → активация → … → Linear
    layers = []
    for j in range(len(sizes) - 1):
        act = activation if j < len(sizes) - 2 else output_activation
        layers += [nn.Linear(sizes[j], sizes[j + 1]), act()]
    return nn.Sequential(*layers)


def reward_to_go(rews):
    # Награда «до конца» для каждого шага: сумма наград с этого шага и дальше
    n = len(rews)
    rtgs = np.zeros_like(rews)
    for i in reversed(range(n)):
        rtgs[i] = rews[i] + (rtgs[i + 1] if i + 1 < n else 0)
    return rtgs


def train(env_name='CartPole-v1', hidden_sizes=[32], lr=1e-2,
          epochs=50, batch_size=5000, use_rtg=False, seed=0):
    torch.manual_seed(seed)                  # фиксируем случайность сети
    env = gym.make(env_name)
    assert isinstance(env.observation_space, Box), "нужны непрерывные наблюдения"
    assert isinstance(env.action_space, Discrete), "нужны дискретные действия"
    obs_dim = env.observation_space.shape[0]
    n_acts = env.action_space.n

    # Сеть политики: наблюдение → логиты действий
    logits_net = mlp(sizes=[obs_dim] + hidden_sizes + [n_acts])

    def get_policy(obs):
        return Categorical(logits=logits_net(obs))

    def get_action(obs):
        return get_policy(obs).sample().item()

    # «Функция потерь», градиент которой на данных текущей политики равен градиенту политики со знаком минус
    def compute_loss(obs, act, weights):
        logp = get_policy(obs).log_prob(act)
        return -(logp * weights).mean()

    optimizer = Adam(logits_net.parameters(), lr=lr)

    def train_one_epoch(epoch):
        batch_obs, batch_acts, batch_weights = [], [], []
        batch_rets, batch_lens = [], []
        obs, _ = env.reset(seed=seed + epoch)     # было: obs = env.reset()
        ep_rews = []
        while True:
            batch_obs.append(obs.copy())
            act = get_action(torch.as_tensor(obs, dtype=torch.float32))
            # было: obs, rew, done, _ = env.step(act)
            obs, rew, terminated, truncated, _ = env.step(act)
            done = terminated or truncated
            batch_acts.append(act)
            ep_rews.append(rew)
            if done:
                ep_ret, ep_len = sum(ep_rews), len(ep_rews)
                batch_rets.append(ep_ret)
                batch_lens.append(ep_len)
                if use_rtg:
                    batch_weights += list(reward_to_go(ep_rews))   # вес — reward-to-go с шага t
                else:
                    batch_weights += [ep_ret] * ep_len             # вес — отдача всего эпизода R(tau)
                obs, _ = env.reset()              # было: obs, done, ep_rews = env.reset(), False, []
                ep_rews = []
                if len(batch_obs) > batch_size:
                    break
        # Один шаг градиентного подъёма по качеству политики
        optimizer.zero_grad()
        batch_loss = compute_loss(
            obs=torch.as_tensor(np.array(batch_obs), dtype=torch.float32),
            act=torch.as_tensor(batch_acts, dtype=torch.int32),
            weights=torch.as_tensor(batch_weights, dtype=torch.float32))
        batch_loss.backward()
        optimizer.step()
        return batch_loss, batch_rets, batch_lens

    for i in range(epochs):
        batch_loss, batch_rets, batch_lens = train_one_epoch(i)
        print('epoch: %3d \t loss: %.3f \t return: %.3f \t ep_len: %.3f' %
              (i, batch_loss.item(), np.mean(batch_rets), np.mean(batch_lens)))


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--env_name', '--env', type=str, default='CartPole-v1')
    parser.add_argument('--lr', type=float, default=1e-2)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--rtg', action='store_true')   # включить веса reward-to-go
    args = parser.parse_args()
    train(env_name=args.env_name, lr=args.lr, epochs=args.epochs, use_rtg=args.rtg)

Запуск: python simple_pg_gymnasium.py (простейший вариант) и python simple_pg_gymnasium.py --rtg (reward-to-go). В оригинале среда по умолчанию — CartPole-v0 (максимум 200 шагов), в Gymnasium мы взяли CartPole-v1 (максимум 500). В нашем прогоне на CPU (seed 0, около полутора минут на вариант) получилось так, каждая пятая эпоха:

простейший вариант (веса R(tau))          reward-to-go (флаг --rtg)
epoch:   7   return:  41.7                epoch:   9   return:  57.0
epoch:  17   return:  75.5                epoch:  19   return: 139.4
epoch:  27   return:  93.9                epoch:  29   return: 434.4
epoch:  37   return: 117.0                epoch:  39   return: 468.2
epoch:  47   return: 190.3                epoch:  49   return: 479.9

Ваши числа будут немного другими (другие версии библиотек дают другую случайность), но картина та же: с reward-to-go отдача растёт заметно быстрее — меньше шума в оценке градиента. А колонка loss, которую печатает скрипт, растёт вместе с отдачей — наглядная иллюстрация предупреждения выше: эта «функция потерь» не измеряет качество.

Словарик

English	Русский
policy gradient	градиент политики
log-derivative trick	трюк с логарифмической производной
grad-log-prob	градиент логарифма вероятности
sample estimate, sample mean	выборочная оценка, выборочное среднее
reward-to-go	reward-to-go («оставшаяся награда», в курсе — отдача $G_t$)
baseline	базовая линия (baseline)
variance	дисперсия (разброс)
advantage function	функция преимущества
marginal distribution	маргинальное распределение
law of iterated expectations	закон повторного матожидания

Назад: Часть 2. Дальше: VPG и PPO или как стать исследователем глубокого RL.

Footnotes

1. Автору статьи неизвестно, чтобы у этой леммы где-то в литературе было стандартное название. Но раз она встречается так часто, кажется вполне разумным дать ей какое-нибудь имя, чтобы на неё было удобно ссылаться. ↩