Offline RL English

Глава 6: Ограничение политики и Actor-Critic (TD3+BC, AWAC)

«Держись того, что видел — но используй критика, чтобы опереться на лучшее из этого.»

Где заканчивается пессимизм по значениям

В главах 4 и 5 мы боролись с ошибкой экстраполяции, делая функцию ценности пессимистичной: CQL штрафует Q-значения за OOD-действия; IQL вообще не обращается к OOD, используя экспектильную регрессию и advantage-weighted регрессию.

Другой подход — оставить критика (Q или V) по сути без изменений и ограничить или регуляризовать актора так, чтобы обученная политика оставалась близка к поведенческой. Агент по-прежнему улучшается за счёт данных — критик указывает, какие действия в датасете были лучше — но политика не может уходить далеко в OOD-области.

Это семейство policy-constraint (или actor-regularized) offline RL. Методы являются actor-critic: обучаются и критик, и политика, но в целевую функцию политики явно входит член, притягивающий её к данным. Два распространённых метода: TD3+BC (минималистичный, детерминированный) и AWAC (advantage-weighted, in-sample обновления актора).

TD3+BC: минималистичный подход с регуляризацией политики

TD3+BC (Fujimoto & Gu, NeurIPS 2021) добавляет к loss актора один член: loss поведенческого клонирования, штрафующий отклонение от действий в датасете. Идея: актор должен максимизировать Q и оставаться близко к действиям из данных.

Идея

Целевая функция актора:

$$\pi^* = \arg\max_\pi \; \mathbb{E}_{(s,a) \sim \mathcal{D}} \left[ \lambda \, Q(s, \pi(s)) - \bigl(\pi(s) - a\bigr)^2 \right]$$

Первый член — использовать критика; второй — имитировать данные. Гиперпараметр $\lambda$ задаёт баланс. В статье Fujimoto & Gu Q нормализуют по среднему абсолютному значению по батчу: $q / \bigl( |B|^{-1} \sum |Q(s,a)| \bigr)$, чтобы оба члена были одного масштаба.

Формализация

Критик (Q): как в TD3 — два Q-сети, target-сети, TD loss. Актор: детерминированная политика $\pi_\phi(s)$; loss выше. Теоретической гарантии (в отличие от CQL) нет — эмпирический, простой в реализации метод.

AWAC: Advantage-Weighted Actor-Critic

AWAC (Nair et al., 2020) обновляет актора только по данным: политика подгоняется с весами по преимуществу (advantage), без сэмплирования из текущей политики.

$$\mathcal{L}_\pi(\phi) = -\mathbb{E}_{(s,a) \sim \mathcal{D}} \left[ \exp\!\left( \frac{1}{\beta} \bigl( Q(s,a) - V(s) \bigr) \right) \cdot \log \pi_\phi(a | s) \right]$$

$A(s,a) = Q(s,a) - V(s)$ — преимущество действия $a$ в состоянии $s$. Экспоненциальный вес усиливает хорошие действия; $\beta$ — температура. OOD-запросов к актору нет.

Реализация

📄 Полный код: td3bc.py

TD3+BC: сети и loss актора

TD3+BC использует ту же архитектуру, что и TD3: детерминированный актор, две Q-сети, target-сети.

class Actor(nn.Module):
    """Deterministic policy s -> a in [-1, 1]. Same as IQL DeterministicPolicy."""
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, action_dim), nn.Tanh(),
        )
    def forward(self, state):
        return self.net(state)
    def act(self, state):
        with torch.no_grad():
            return self.forward(state).cpu().numpy().squeeze()


def td3bc_actor_loss(actor, Q1, states, actions, lambda_=0.25):
    """
    TD3+BC actor loss (minimize):
      -lambda_ * Q_norm(s, pi(s)) + ||pi(s) - a||^2
    Эквивалентная максимизация:
      lambda_ * Q_norm(s, pi(s)) - ||pi(s) - a||^2
    Q нормализуется по среднему абсолютному значению по батчу (Fujimoto & Gu).
    """
    pi = actor(states)
    q = Q1(states, pi)
    q_norm = q / (q.abs().mean() + 1e-6)
    bc_loss = ((pi - actions) ** 2).mean()
    return -q_norm.mean() * lambda_ + bc_loss

Оба члена (Q и BC) вносят вклад в градиент; в td3bc.py используется нормализация по среднему абсолютному значению Q (как в статье), а не Z-нормализация, чтобы $\lambda$ соответствовал оригинальной настройке.

AWAC-style loss политики (advantage-weighted)

def awac_actor_loss(policy, Q, V, states, actions, beta=1.0):
    """
    Advantage-Weighted Regression: log pi(a|s) weighted by exp(A(s,a)/beta).
    A(s,a) = Q(s,a) - V(s). Requires stochastic policy that outputs log_prob.
    """
    with torch.no_grad():
        A = Q(states, actions) - V(states)
        weights = (A / beta).exp()
        weights = weights / (weights.mean() + 1e-6)  # stabilize
    log_prob = policy.log_prob(states, actions)
    return -(weights * log_prob).mean()

Для детерминированной политики (как в TD3) можно использовать гауссов с малой дисперсией вокруг $\pi(s)$ как суррогат log_prob или перейти на стохастическую голову.

Ограничения

Гарантии нижней границы нет. Чувствительность к $\lambda$ / $\beta$. Детерминированная политика (TD3+BC) не может быть мультимодальной.

Итог

Метод Где ограничение OOD актор? Теория
TD3+BC Loss актора (BC-штраф) Да Нет
AWAC Loss актора (веса по advantage) Нет Нет

Методы с ограничением политики дают простой способ улучшаться над поведенческой политикой, оставаясь близко к данным. TD3+BC — самый простой в реализации; для максимальной устойчивости по-прежнему предпочтительны CQL и IQL.

Глава 7 переходит к другой парадигме: Decision Transformers — offline RL как моделирование последовательностей без Bellman-бэкапа.

Литература