ML: Методы основанные на политике

Введение

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

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

Policy Gradients

Предположим сначала, что состояние среды описывается вектором $n$ вещественных чисел, а действия дискретны (один из элементов множества $m$ целых чисел).

Будем аппроксимировать функцию политики $\pi_\theta(a|s)$ (условная вероятность) нейронной сетью с параметрами $\theta$, с $n$ входами, $m$ выходами на которых находятся вероятности каждого из $m$ действий.

Для обучения сети необходимы обучающие данные $\{s^{(k)},\,\pi^{(k)}\},~k=1...N$, где $s^{(k)}$ - $n$-мерный вектор для $k$-того примера, а $\pi^{(k)}=(p_1,...,p_m)^{(k)}$ нормированный на единицу неотрицательный $m$-мерный вектор (вероятности).

Возьмём начальную функцию $\pi_\theta(a|s)$ со случайно инициализированными параметрами. С текущей функцией политики проведём $M$ эпизодов (игровых сессий), вычисляя в каждом получаемое суммарное вознаграждение $R$. На основе этой информации сформируем обучающие данные в виде тройки $(s_t,\,a_t,\,R)$, где в данном эпизоде во всех тройках одинаковый $R$.

Каждую тройку $(s,a,R)$ из нескольких эпизодов будем нумеровать верхним индексом $k$. При обучении минимизируется сумма: $$ L = -\sum_i R^{(k)} \log p^{(k)}, ~~~~~~~~~p^{(k)} = \pi_\theta(a^{(k)}\,|\,s^{(k)}), $$ где $p^{(k)}$ вероятность действия $a^{(k)}$. Таким образом, при прямом распространении на вход сети подаётся пример $s^{(k)}$, на выходе получается распределение вероятностей $(p_1,...,p_m)^{(k)}$. Из этих вероятностей выбирается та, которая соответствует сделанному в примере действию $a^{(k)}$.

Так как $\log p < 0$, минимизация $L$ максимизиирует вероятности сделанных действий. При этом больший вес отдаётся примерам в которых в данном эпизоде получилось большее вознаграждение $R$. То что вместо максимизации вероятностей следует максимизировать их логарифмы следует их несложных математических выкладок приведенных ниже. Впрочем, если вспомнить ошибку кросс-энтропии, используемую при мультиклассовой классификации, такой выбор достаточно естественен. Данный метод лежит в основе группы методов которые называются градиентом политики (Policy Gradients).

Нейронная сеть политики

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

Будем при помощи сигмоида ограничивать выход нейронной сети (nS=n - размерность состояния, nA=m - число возможных действий, nH - нейронов в скрытом слое): 

model = nn.Sequential(
          nn.Linear(nS, nH),   
          nn.ReLU(),    
          nn.Linear(nH, nA), 
          nn.Sigmoid() )
В общем случае на вход может поступать вектор компонент одиночного наблюдения (состояния) как numpy массив или сразу несколько наблюдений (при пакетном вычислении). Поэтому следует состояние преобразовать: 
def run_model(state):
    return model( torch.Tensor(state).view(-1, nS) )
Тогда выбор действия можно осуществлять при помощи функции multinomial: 
def policy(state):
    return torch.multinomial( run_model(state), 1)

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

Сигмоид на выходе можно (лучше) не ставить. Тогда выходы следует пропустить через softmax-функцию. 

def policy(state):
    probs = torch.softmax( run_model(state), 1)
    return  torch.multinomial(probs, 1)
Вместо комбинации softmax и multinomial можно использовать распределение Categorical принимающее ненормированные вероятности (произвольные вещественные числа): 
from torch.distributions.categorical import Categorical

def policy(state):     
    return  Categorical( logits = run_model(state) ).sample()
Наконец, если используются не все действия и они перечислены в списке action_space (например в MountainCar-v0 это action_space=[0,2]), можно воспользоваться библиотекой numpy: 
def policy(state):  
    probs = torch.softmax( run_model(state), 1).numpy()      
    return np.random.choice(action_space, p=probs)

Математическая мотивировка $^*$

Рассмотрим математическую мотивировку функции ошибки с логарифмами вероятности политики. Будем считать, что награда на каждом шаге зависит от текущего состояния среды $r_t=r(s_t)$. Нас интересует максимальная суммарная награда в течении эпизода ($T$ - время попадания в терминальное состояние): $$ R(\tau) = r_1+r_2+...+r_T. $$ Эта награда реализуется на конкретной траектории $\tau=\{s_0,a_0,\,s_1,a_1,\,...,s_{T-1},a_{T-1}\,s_T\}$ и зависит от политики агента $\pi_\theta(s\,|\,a)$ с параметрами $\theta$. Вероятность траектории определяется распределением начального состояния $p(s_0)$, условными вероятностями политики $\pi_\theta(s|a)$ и моделью среды $P(s'|s,a)$: $$ P_\theta(\tau) = p(s_0)\cdot\prod^{T-1}_{t=0}\pi_\theta(a_t|s_t)\,P(s_{t+1}|s_t,a_t). $$ Усреднение по всем траекториям имеет вид: $$ \langle R \rangle = \sum_{\tau} P_\theta(\tau)\cdot R(\tau), $$ где сумма по $\tau$ - это суммы по всем $\{s_0,a_0,\,s_1,a_1,\,...,\,s_T\}$. Чтобы максимизировать это выражение при помощи градиентного метода, необходимо вычислить градиент $\nabla_\theta\langle R \rangle$ по параметрам $\theta$ модели политики $\pi_\theta(s|a)$. Запишем логарифм вероятности траектории и возьмём его градиент: $$ \log P_\theta(\tau) = \sum^{T-1}_{t=0}\log\pi_\theta(a_t|s_t) + [\mathrm{не~зависят~от}~ \theta]~~~~\Rightarrow~~~~ \nabla_\theta P_\theta(\tau) = P_\theta(\tau)\, \sum^{T-1}_{t=0}\nabla\log\pi_\theta(a_t|s_t). $$ Таким образом: $$ \nabla_\theta \langle R \rangle = \sum_{\tau} P_\theta(\tau)\cdot R(\tau)\cdot \sum^{T-1}_{t=0}\nabla\log\pi_\theta(a_t|s_t). $$

Если модель среды $P(s_1|s_0,a_0)$ неизвестна, градиент оценивается по методу Монте-Карло: $$ \nabla_\theta \langle R \rangle = \frac{1}{N}\sum^{N}_{k=1} \sum^{T_k-1}_{t=0}R^{(k)}\cdot\nabla\log\pi_\theta(a^{(k)}_t|s^{(k)}_t), $$ где $R^{(k)}$ - награда, полученная в $k$-том эпизоде, а $\{s^{(k)}_t, a^{(k)}_t\}$ - значения состояния и действия, полученные в $k$-том эпизоде на временном шаге $t$. Проблемы такой оценки связаны с двумя неприятностями:

Заметим, что, если в исходной формуле для $\langle R\rangle$ из $R$ вычесть константу (baseline): $R\mapsto R-b$, то значение градиента не изменится: $$ \nabla\langle R \rangle = \nabla\sum_{\tau} P(\tau)\cdot (R-b) =\nabla\sum_{\tau} P(\tau)\cdot R - b\,\nabla\sum_{\tau} P(\tau) =\nabla\sum_{\tau} P(\tau)\cdot R, $$ так как сумма от $P(\tau)$ равна единице, градиент от которой нулевой. В тоже время диспресия градиента от $b$ зависит, что при правильном выборе $b$ позволяет её снизить.

Метод REINFORCE

Метод REINFORCE (Williams, 1992) отличается от описанного выше тем, что в тройках вместо суммарной награды за эпизод стоит будущее дисконтированное суммарное вознаграждение с момента $t$ в данном эпизоде (causality principle): $$ R_{t+1} = r_{t+1} + \gamma\,r_{t+2} + \gamma^2\,r_{t+3} ... $$ где $\gamma$ - дисконтирующий множитель (гиперпараметр). Таким образом имеем тройки вида $(s_t,\,a_t,\,R_{t+1})$. Функция ошибки по-прежнему имеет вид: $$ L = - \sum^b_{k=1} R^{(k)}\,\log \pi(a^{(k)}|s^{(k)}) $$

Если из будущего вознаграждения вычитают константу, то такой метод называется REINFORCE with baseline $$ L = - \sum^b_{k=1} (R^{(k)}-b)\,\log \pi(a^{(k)}|s^{(k)}) $$ В качестве b можно использовать награду за данный эпизод, среднее значение $R$ по батчу и т.п.

Подобные методы неплохо справляются с простыми задачами. Например, в CartPole вначале обучения стержень часто падает и мы имеем набор эпизодов с различными $R$ и следовательно различные веса для усиления "правильных" действий и подавления "неправильных". Заметим, что в этой задаче функция политики имеет тривиальную линейную разделяющую поверхность (плоскость): $a = s_0+s_1+3s_2+s_4 > 0$, поэтому она удобна для Gradient Policy.

В сложных задачах, типа MountainCar-v0, в начале обучения случайные действия не приводят к достижению цели и все суммарные вознаграждения (в GradientPolicy) будут одинаковыми. Следовательно алгоритм не может отличить "хорошие" действия от "плохих". Использование будущего вознаграждения (REINFORCE) ситуацию не сильно меняет. Так как вознаграждения на каждом шаге одинаковые (-1), веса будут больше для тех состояний-действий которые находятся ближе к конце эпизода (хотя они богут быть далеки от целевого состояния).

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

Метод Актёра-Критика

Правильный выбор весов при логарифмах вероятностей играет очень важную роль. Предположим что нам известна ценность состояний $V(s)$ и состояний-действий $Q(s,a)$. Ценность состояния $V(s)$, как и политика, аппроксимируется нейронной сетью $V_\omega(s)$ с параметрами $\omega$. Для её оптимизации минимизируются квадраты отклонений от уравнения Беллмана: $$ L_\omega = \sum \Bigr( r'+\gamma\,V_\omega(s') - V_\omega(s) \Bigr)^2. $$ Так как она вычисляется по сделанным действиям (после которых получена награда $r'$) - это ценность состояния при текущей политике.

Функцию ошибки для политики запишем в следующим виде: $$ L_\theta = -\sum A(s,a) \,\log\pi_\theta(a|s),~~~~~~~~~A(s,a) = Q(s,a)-V(s), $$ где веса $A(s,a)$ характеризуют преимущество выполненного действия $a$ (advantage function). Действительно, $Q(s,a)$ - это будущая награда в состоянии $s$ при выполнении действии $a$. Из неё вычитается ценность $V(s)$ состояния $s$ при данной политике. Если действие не оптимально преимущество будет отрицательным $Q(s,a) \le V(s)$. Минимизации ошибки, как и раньше, проводится по параметрам $\theta$.

Заметим, что, если действие оптимально, то $A(s,a^*)=0$, что не хорошо (?), так как подавляет максимизацию вероятности такого действия. Возможно, лучше от этого выражения взять экспоненту $A(s,a)\mapsto \exp\bigr[\mu A(s,a)\bigr]$, где $\mu$ - гиперпараметр. Тогда оптимальное действие будет иметь единичный вес, а неоптимальные - меньше единицы.

В качестве $Q(s_t,a_t)$ можно взять реальное значение будущей суммарной награды $R_{t+1}=r_{t+1}+\gamma\,r_{t+2}+...$ в данном эпизоде или её оценку $r_{t+1}+\gamma\,V(s_{t+1})$ по $V-функции$. Наконец, можно использовать "n-шаговую" оценку: $$ Q(s_t,a_t)=r_{t+1}+\gamma \,r_{t+2} +...+\gamma^{n-1}\, r_{t+n} + \gamma^n\,V(s_{t+n+1}). $$

Этот алгоритм называется актёром-критиком с преимуществом (Advantage Actor-Critic, A2C). Актёром считается функция политики, а критиком функция преимущества, которая "критикует"

Непрерывные действия

До сих пор мы максимизировали вероятности $m$ дискретных действий. Нейронная сеть $\pi_\theta(a|s)$ имела $m$ выходов, на которых находилась вероятность каждого действия.

Пусть теперь есть $m$-компонентный вектор действия с вещественным элементами. Будем считать, что вероятность $m$-той компоненты описываются нормальным распределением: $$ \pi(a|s) =\prod^{m}_{k=1}\frac{1}{\sqrt{2\pi D_k}}\, e^{-(a_k-\mu_k)^2/2D_k}. $$ Будем учить сеть получать средние значения $\mu_k$ и дисперсии $D_k=\sigma^2_k$. Соответственно логарифм функции политики будет равен: $$ \log \pi(a|s) = - \sum^{m}_{k=1}\Bigr[ \frac{(a_k-\mu_k)^2}{2 D_k} + \log\sqrt{2\pi D_k}\,\Bigr]. $$

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

class ModelA2C(nn.Module):
    def __init__(self, n, m, h):
        super(ModelA2C, self).__init__()

        self.base = nn.Sequential(
            nn.Linear(n, h),
            nn.ReLU() )
        self.mu = nn.Sequential(
            nn.Linear(h, m),
            nn.Tanh() )
        self.var = nn.Sequential(
            nn.Linear(h, m),
            nn.Softplus() )
        self.value = nn.Linear(h, 1)

    def forward(self, x):
        y = self.base(x)
        return self.mu(y),self.var(y),self.value(y)




  • n - компонент состояния
  • m - компонент действия
  • h - нейронов в скрытом слое

    • Дисперсии var должны быть положительными, поэтому на выходе для них используется активация Softplus ("гладкий" ReLU). Так как изменение компонент действий находится в симметричном пределе, на выходе mu стоит Tanh. Наконец ценность состояния может быть любой, поэтому для неё активационной функции нет.

    PPO

    PPO расшифровывается как Proximal Policy Optimization.

    Градиенты детерминированной политики

    Deterministic Policy Gradient (DPG) работает с детерминированной политикой $a=\pi(s)$.