Policy Gradient и REINFORCE

1. Сеть политики

Пусть состояние — вектор из \(n\) вещественных чисел, действий — конечное множество из \(m\) штук. Политику \(\pi_\theta(a\,|\,s)\) — условную вероятность действия — аппроксимируем нейросетью с параметрами \(\theta\): \(n\) входов, \(m\) выходов-вероятностей. Знакомая по «Нейросетям» конструкция многоклассовой классификации, только «класс» — это действие, а разметки нет: правильных действий нам никто не скажет.

model = nn.Sequential(
    nn.Linear(nS, nH),        # nS - размерность состояния
    nn.ReLU(),
    nn.Linear(nH, nA),        # nA - число действий; выходы - логиты
)

def run_model(state):          # одно наблюдение или батч
    return model( torch.Tensor(state).view(-1, nS) )

Действие не выбирается жадно — оно разыгрывается по вероятностям: политика сама стохастическая, и разведка встроена в неё (ε-жадность больше не нужна). Конспект преподавателя приводит несколько равносильных способов розыгрыша; самый удобный — распределение Categorical, принимающее сырые логиты:

from torch.distributions.categorical import Categorical

def policy(state):
    return Categorical( logits=run_model(state) ).sample()

# эквивалентно: softmax + multinomial
def policy(state):
    probs = torch.softmax( run_model(state), 1)
    return torch.multinomial(probs, 1)

# а если действий меньше, чем в среде (MountainCar: action_space=[0,2]):
def policy(state):
    probs = torch.softmax( run_model(state), 1).numpy()
    return np.random.choice(action_space, p=probs)

2. Идея: кросс-энтропия, взвешенная возвратом

Чему учить такую сеть, если правильных действий нет? Есть кое-что взамен — итог игры. Сыграем текущей политикой \(M\) эпизодов, запишем каждую тройку «состояние, сделанное действие, суммарная награда эпизода» \((s^{(k)}, a^{(k)}, R^{(k)})\) и минимизируем

Сравните с кросс-энтропией многоклассовой классификации («Нейросети», конспект 3): там максимизировался логарифм вероятности правильного класса, здесь — логарифм вероятности сделанного действия, взятый с весом \(R^{(k)}\). Действия из удачных эпизодов усиливаются сильнее, из провальных — слабее (а при отрицательных наградах — подавляются). Это и есть семейство методов градиента политики (англ. policy gradient). Почему в (9.1) стоит именно логарифм — пока принято на веру; сейчас выведем это честно, заодно найдя у наивной формулы два изъяна.

3. Вывод: градиент, который не трогает среду

Зафиксируем эпизодическую задачу конечной длины \(T\). Траектория — \(\tau = (S_0, A_0, R_1, S_1, A_1, R_2, \ldots, S_{T-1}, A_{T-1}, R_T, S_T)\); её вероятность при политике \(\pi_\theta\) складывается из распределения старта, политики и динамики среды:

Максимизируем ожидаемый дисконтированный возврат:

3.1. Приём likelihood-ratio

Теорема 9.1 (градиент через траектории). Пусть \(\pi_\theta\) дифференцируема по \(\theta\) и её носитель не зависит от \(\theta\). Тогда

Доказательство. Горизонт конечен — дифференцируем под знаком суммы: \(\nabla_\theta J = \sum_\tau \nabla_\theta p_\theta(\tau) \cdot R(\tau)\). Применяем тождество \(\nabla_\theta p_\theta = p_\theta \nabla_\theta \log p_\theta\) (верное всюду, где \(p_\theta > 0\)) — сумма снова становится ожиданием по \(p_\theta\). Осталось взять логарифм (9.2):

где первый и третий члены от \(\theta\) не зависят и при дифференцировании исчезают. \(\blacksquare\)

Вот и обещанный фокус: динамика среды \(P\) выпала из градиента до того, как нам понадобилось её знать. Среда может оставаться чёрным ящиком — достаточно уметь в ней играть и дифференцировать собственную политику. Заодно легализован логарифм из (9.1): он пришёл из \(\nabla p = p\,\nabla\log p\), а не из аналогии с кросс-энтропией.

3.2. Причинность: прошлое не зависит от текущего действия

Формула (9.4) неэкономична: полный возврат эпизода умножается на градиенты всех действий — в том числе ранних, которые никак не могли повлиять на уже полученные награды. Интуиция подсказывает, что награды до момента \(t\) из веса действия \(A_t\) можно убрать. Докажем.

Лемма 9.1. Для любого \(t\):

Доказательство. Прошлая сумма \(C_t = \sum_{k<t} \gamma^k R_{k+1}\) определяется историей до момента \(t\) и от \(A_t\) не зависит. По формуле повторного ожидания достаточно показать, что \(\mathbb{E}[\nabla_\theta \log \pi_\theta(A_t|S_t) \,|\, S_t = s] = 0\):

\(\blacksquare\) Тождество (9.6) — «сумма вероятностей равна единице, градиент единицы — ноль» — ещё дважды сработает ниже; запомните его.

Теорема 9.2 (форма reward-to-go). Обозначив возврат с момента \(t\) как \(G_t = \sum_{k \ge t} \gamma^{k-t} R_{k+1}\) (конспект 1, формула (1.2)), имеем

Доказательство. В (9.4) разбиваем полный возврат при каждом \(t\) на прошлое \(C_t\) и будущее \(\gamma^t G_t\); слагаемые с \(C_t\) зануляет лемма 9.1. \(\blacksquare\)

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

3.3. Теорема о градиенте политики

Последний шаг — заменить в (9.7) случайный \(G_t\) его условным ожиданием: по определению (конспект 3, формула (3.5)) \(\mathbb{E}[G_t | S_t, A_t] = Q^{\pi_\theta}(S_t, A_t)\). Собирая слагаемые по состояниям, получаем классическую формулировку.

Теорема 9.3 (Policy Gradient Theorem). Пусть \(d^\gamma_{\pi_\theta}(s) = \sum_t \gamma^t \Pr(S_t = s)\) — дисконтированная мера посещения состояний. Тогда

Концептуальный итог: градиент качества политики выражен через локальные величины — «куда сдвинуть вероятности в состоянии \(s\)» и «сколько стоит действие», — и нигде не дифференцирует среду. Постоянно посещаемые состояния влияют на градиент сильнее редких: политика улучшается прежде всего там, где живёт.

3.4. Baseline: сдвиг, который ничего не смещает

Оценки (9.7) корректны, но дисперсия их велика. Классическое средство — вычесть из возврата опору (англ. baseline) \(b(s)\), зависящую только от состояния:

Математическое ожидание не изменилось: добавка \(\mathbb{E}[\nabla \log \pi_\theta \cdot b(s)]\) равна нулю — \(b\) выносится за ожидание по действию, и остаётся тождество (9.6). А вот дисперсия от \(b\) зависит, и удачный выбор (среднее \(R\) по батчу, а лучше всего — оценка ценности состояния \(V(S_t)\), тогда в скобках возникает знакомое по конспекту 8 преимущество \(G_t - V(S_t)\)) снижает её в разы. Из этой скобки в следующем конспекте вырастет актёр-критик.

4. REINFORCE

Алгоритм REINFORCE (Уильямс, 1992) — это формула (9.9), исполненная по Монте-Карло, как в конспекте 4: играем эпизод до конца, считаем реальные \(G_t\), делаем градиентный шаг.

Обратите внимание на слово on-policy: формулы раздела 3 верны для траекторий, разыгранных текущей политикой. После каждого шага обучения старые эпизоды становятся чужими — память переходов из DQN здесь неприменима, опыт потребляется один раз. Это честная цена за прямую оптимизацию политики.

5. Тренажёр: REINFORCE на двуруком бандите

Чтобы увидеть механику голыми глазами, возьмём наименьшую задачу, где она работает, — softmax-бандита из лекции: одно состояние, два действия, политика \(\pi_\theta(a{=}i) = e^{\theta_i} / (e^{\theta_1} + e^{\theta_2})\). Градиент логарифма считается вручную: \(\partial \log \pi_\theta(a{=}i) / \partial \theta_j = \mathbf{1}[i{=}j] - \pi_\theta(a{=}j)\). Средние награды рычагов задаются ползунками (агент их не знает и видит только зашумлённые выплаты); эпизод — один ход, поэтому \(G_t\) — это просто полученная награда.

Три опыта, которые стоит проделать. Первый: при наградах 1 против 2 вероятность лучшего рычага растёт к единице — политика «подталкивается» к удачным выборам, в точности по формуле шага 3. Второй: сделайте награды равными — вероятности блуждают вокруг 0,5, градиент в среднем нулевой, но шаг за шагом политика может надолго «залипать» у краёв: это и есть дисперсия Монте-Карло оценки. Третий: включите baseline при наградах 1 и 2 — кривая становится заметно глаже: вычитание среднего убирает из веса общий уровень наград, оставляя только их разницу.

6. Практика: CartPole решается напрямую

Теперь настоящая среда. В CartPole (конспект 1) REINFORCE особенно удачлив, и понятно почему: чем дольше эпизод, тем больше возврат, — уже сама длина эпизода отделяет хорошие розыгрыши от плохих; к тому же оптимальная политика в этой задаче почти линейна по состоянию, и маленькой сети достаточно. Ноутбук курса собирает REINFORCE в полсотни строк:

class Policy(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(nS, 128, bias=False),
            nn.Dropout(p=0.6),          # dropout заметно улучшает политику
            nn.ReLU(),
            nn.Linear(128, nA, bias=False),
            nn.Softmax(dim=-1)          # выход - вероятности действий
        )

Выбор действия — розыгрыш из Categorical с запоминанием логарифма вероятности (он понадобится в градиенте):

def select_action(state, log_probs):
    c = Categorical( policy(torch.from_numpy(state).float()) )
    action = c.sample()
    log_probs.append( c.log_prob(action) )   # log pi(a|s) - в историю эпизода
    return int(action)

Сердце метода — обновление по концу эпизода. Reward-to-go считается одним проходом с конца (тот же приём, что в MC-оценке конспекта 4), а нормировка возвратов по батчу — это baseline «среднее по эпизоду» плюс масштабирование к единичной дисперсии:

def update_policy(log_probs, rewards_episode):
    R, returns = 0.0, []
    for r in rewards_episode[::-1]:        # дисконтируем с конца эпизода
        R = r + gamma * R
        returns.insert(0, R)               # G_t для каждого шага t
    returns = torch.tensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)  # baseline
    loss = -(torch.stack(log_probs) * returns).sum()   # формула (9.9)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Мы прогнали этот код при подготовке конспекта: задача решается — средняя длина эпизода превышает порог 475 из 500 возможных. Слева кривая обучения этого запуска, справа — обученная политика в деле:

А теперь то же самое — вживую. Ниже настоящий CartPole: точная физика среды, а за «обученной» политикой стоят подлинные веса сети, обученной выше, — они экспортированы из нашего запуска прямо в страницу. Сеть каждый такт сообщает свои вероятности — можно подглядывать, что она «думает», даже управляя вручную:

Поиграйте за все три политики. Случайная падает за пару десятков шагов — вы, скорее всего, продержитесь дольше, но 500 шагов руками не выстоять: тележка уползает за границу быстрее, чем кажется. Обученная сеть держит стержень все 500 шагов мелкими симметричными подруливаниями — и обратите внимание на её вероятности: вдали от беды они близки к 50/50 (какая разница, куда шагнуть, когда всё стабильно), но стоит стержню накрениться — распределение резко перекашивается в сторону спасающего действия. Это выученная стохастическая политика из раздела 1, работающая у вас на глазах.

# reinforce_cartpole.py — REINFORCE (Monte-Carlo policy gradient) на CartPole-v1.
# По ноутбуку курса; API Gymnasium. Один файл: обучение, кривая, сохранение весов.
import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Categorical
import gymnasium as gym

learning_rate = 0.01
gamma = 0.99
env = gym.make("CartPole-v1")
torch.manual_seed(1); np.random.seed(1)

class Policy(nn.Module):
    def __init__(self):
        super().__init__()
        nS = env.observation_space.shape[0]      # 4 компоненты состояния
        nA = env.action_space.n                  # 2 действия: влево/вправо
        self.net = nn.Sequential(
            nn.Linear(nS, 128, bias=False),
            nn.Dropout(p=0.6),                   # dropout заметно улучшает политику
            nn.ReLU(),
            nn.Linear(128, nA, bias=False),
            nn.Softmax(dim=-1)                   # выход - вероятности действий
        )

    def forward(self, x):
        return self.net(x)

policy = Policy()
optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate)

def select_action(state, log_probs):
    """Выбор действия по вероятностям политики; логарифм вероятности - в историю."""
    state = torch.from_numpy(state).float()
    c = Categorical(policy(state))
    action = c.sample()
    log_probs.append(c.log_prob(action))
    return int(action)

def update_policy(log_probs, rewards_episode):
    """Шаг REINFORCE по формуле (9.9): reward-to-go + нормировка (baseline)."""
    R, returns = 0.0, []
    for r in rewards_episode[::-1]:              # дисконтируем с конца эпизода
        R = r + gamma * R
        returns.insert(0, R)                     # G_t для каждого шага
    returns = torch.tensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)  # baseline+масштаб
    loss = -(torch.stack(log_probs) * returns).sum()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

running, history = 10.0, []
for episode in range(1, 1001):
    state, _ = env.reset(seed=episode)
    log_probs, rewards_episode = [], []
    for t in range(1, 1001):
        action = select_action(state, log_probs)
        state, reward, terminated, truncated, _ = env.step(action)
        rewards_episode.append(reward)
        if terminated or truncated:
            break
    update_policy(log_probs, rewards_episode)
    history.append(t)
    running = running * 0.99 + t * 0.01          # скользящая средняя длина
    if episode % 50 == 0:
        print(f"эпизод {episode:4d}  длина {t:4d}  средняя {running:6.1f}", flush=True)
    if running > 475:                            # порог "решённости" CartPole-v1
        print(f"Решено на эпизоде {episode}: средняя длина {running:.1f}")
        break

torch.save(policy.state_dict(), r"C:\Users\Evgenie\AppData\Local\Temp\reinforce_policy.pt")
np.save(r"C:\Users\Evgenie\AppData\Local\Temp\reinforce_history.npy", np.array(history))

7. Где REINFORCE ломается

Успех на CartPole не должен обманывать. У метода два родовых недостатка, и оба видны уже в формулах.

Дисперсия. Возможных траекторий чудовищно много, а оцениваем мы ожидание по ним счётным числом эпизодов — оценка градиента шумит (тренажёр показывал это даже на бандите с двумя действиями). Шумный градиент вынуждает малые шаги, малые шаги — множество эпизодов; а поскольку метод on-policy, каждый эпизод используется один раз и выбрасывается. REINFORCE прожорлив к данным по построению.

Одинаковые возвраты. Веса при логарифмах различают действия, только если возвраты эпизодов различаются. Конспект преподавателя разбирает MountainCar: случайная политика в начале обучения никогда не доезжает до флага, все эпизоды длятся 200 шагов, и все возвраты равны −200 — алгоритму не за что зацепиться, «хорошие» действия неотличимы от «плохих». Переход к reward-to-go мало что меняет: награды на каждом шаге одинаковы (−1), и веса просто больше у шагов ближе к концу эпизода — независимо от того, вела машинка к цели или нет. Сравните с DQN из конспекта 7, который решает эту задачу: value-методам хватает редкого сигнала терминальности, растекающегося по бутстрапу.

Оба лекарства уже названы в этом конспекте. Дисперсию лечит baseline, а лучший baseline — ценность состояния \(V(S_t)\): в весе остаётся преимущество \(G_t - V(S_t)\), «насколько действие лучше обычного для этого состояния». Но \(V\) неоткуда взять — её надо учить параллельно с политикой, второй сетью. Политика-актёр и оценщик-критик, обучаемые вместе, — это и есть актёр-критик, герой следующего конспекта.

Контрольные вопросы

• Почему на непрерывных действиях value-методы буксуют, а policy-методы — нет?
  Value-методу для выбора действия нужен argmax Q(s,a) по a — на непрерывном пространстве это отдельная задача оптимизации на каждом шаге. Policy-метод параметризует само распределение π(a|s) и просто сэмплирует из него; для непрерывных действий сеть выдаёт параметры распределения (например, среднее и дисперсию гауссова).
• Откуда в функции ошибки policy gradient берётся логарифм вероятности?
  Из likelihood-ratio приёма: ∇p_θ = p_θ · ∇log p_θ. Он превращает градиент суммы по траекториям обратно в ожидание по ним (теорема 9.1), а логарифм произведения (9.2) рассыпается в сумму, из которой члены среды исчезают при дифференцировании. Аналогия с кросс-энтропией — следствие, а не причина.
• Почему градиент политики не требует дифференцировать модель среды?
  В log p_θ(τ) динамика среды входит слагаемыми log P(s′|s,a), не зависящими от θ, — при взятии ∇_θ они исчезают (теорема 9.1). Остаются только ∇log π_θ(a|s) — производные собственной политики. Поэтому метод работает со средой-«чёрным ящиком»: достаточно уметь играть.
• Докажите, что вычитание baseline b(s) не смещает градиент.
  E[∇log π_θ(A|s)·b(s)] = b(s)·Σ_a π_θ(a|s)∇log π_θ(a|s) = b(s)·Σ_a ∇π_θ(a|s) = b(s)·∇(Σ_a π_θ(a|s)) = b(s)·∇1 = 0 (тождество (9.6)). Ожидание градиента не меняется, а дисперсия зависит от b — правильный выбор (среднее по батчу, V(s)) её снижает.
• Что такое reward-to-go и почему полный возврат эпизода в весе действия — расточительство?
  G_t — возврат с момента t. В полном возврате сидят и награды ДО момента t, которые от действия A_t не зависят: их вклад в ожидание градиента равен нулю (лемма 9.1), но в дисперсию — нет. Замена полного возврата на G_t (теорема 9.2) даёт тот же градиент в среднем с меньшим шумом — принцип причинности.
• Почему REINFORCE нельзя снабдить памятью переходов, как DQN?
  Формулы градиента верны для траекторий, разыгранных текущей политикой (ожидание по τ ~ p_θ) — метод on-policy. После шага оптимизатора старые эпизоды принадлежат уже другой политике, и их вклад смещён. DQN off-policy: его цель построена на одношаговом переходе и от породившей политики не зависит.
• Почему REINFORCE решает CartPole, но не MountainCar?
  В CartPole длина эпизода сама ранжирует розыгрыши: дольше простоял — больше возврат. В MountainCar случайная политика никогда не доезжает, все возвраты равны −200, веса при логарифмах одинаковы — градиенту не за что зацепиться; reward-to-go лишь взвешивает шаги по близости к концу эпизода, а не к цели. Нужен baseline уровня V(s) и обучаемый критик — следующий конспект.
• Выпишите шаг REINFORCE для softmax-политики двурукого бандита.
  π_θ(a=i) = e^{θ_i}/(e^{θ_1}+e^{θ_2}); ∂log π(a=i)/∂θ_j = 1[i=j] − π(j). После розыгрыша действия i с наградой r: θ_j += α·(r − b)·(1[i=j] − π(j)). При положительном (r−b) вероятность сыгранного действия растёт, остальных — падает; при отрицательном — наоборот.

Источники

1. Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning / R. J. Williams // Machine Learning. — 1992. — Vol. 8. — P. 229–256. — DOI: 10.1007/BF00992696.
2. Sutton, R. S. Policy gradient methods for reinforcement learning with function approximation / R. S. Sutton, D. McAllester, S. Singh, Y. Mansour // Advances in Neural Information Processing Systems. — 2000. — Vol. 12. — P. 1057–1063.