Факторизованные бейзлайны в policy gradient: когда они реально помогают

Проблема, из-за которой всё затевалось

Оценка градиента политики методом REINFORCE несмещённая, но зашумлённая. Дисперсия оценки растёт с длиной эпизода и с размерностью пространства действий: в задаче вроде управления гуманоидом на 17–21 степень свободы один сэмпл градиента может отличаться от следующего на порядок. На практике это выливается в сотни миллионов шагов среды до сходимости и в чувствительность к random seed, из-за которой один и тот же алгоритм на одной и той же задаче даёт разброс итоговой награды в разы.

Классический приём борьбы с дисперсией — вычесть из отдачи бейзлайн. Пока бейзлайн не зависит от действия, оценка остаётся несмещённой: математическое ожидание вычитаемого члена равно нулю. Обычно берут state-value baseline b(s) — обученную оценку ценности состояния. Вопрос, вокруг которого выстроилась отдельная линия исследований: можно ли сделать бейзлайн зависящим ещё и от действия, не потеряв несмещённость, и получить дополнительное снижение дисперсии.

Откуда берётся action-dependent baseline

Ключевое наблюдение: если политика факторизуется по компонентам действия — то есть каждая координата a_i сэмплируется независимо при фиксированном состоянии, что верно для диагональной гауссовой политики, стандарта в непрерывном управлении, — то для градиента по i-й координате можно вычитать бейзлайн, зависящий от всех остальных координат, кроме i.

Бейзлайн, который «видит» все действия, кроме того, по которому берётся производная, остаётся несмещённым — потому что по этой конкретной координате он константа, а значит её вклад в ожидание производной логарифма плотности равен нулю.

Формально для факторизованной политики π(a|s) = ∏ᵢ πᵢ(aᵢ|s) градиент раскладывается покомпонентно, и для каждой компоненты допустим свой бейзлайн bᵢ(s, a₋ᵢ), где a₋ᵢ — все действия, кроме i-го. Отсюда и название — action-dependent factorized baseline.

Почему это в принципе должно снижать дисперсию

Интуиция такая. State-value baseline убирает вариацию, общую для всех действий в данном состоянии. Но остаётся вариация из-за того, как именно скоррелированы координаты действия и отдача. Если bᵢ учитывает остальные координаты, он поглощает часть этой остаточной вариации. В пределе оптимальный action-dependent baseline — это условное ожидание отдачи при фиксированных a₋ᵢ, и оно даёт дисперсию не выше, чем оптимальный state-only baseline.

Где обещание разошлось с практикой

Первые работы показывали заметное ускорение сходимости на задачах MuJoCo и на многоагентных сетапах. Затем появился разбор, который аккуратно отделил эффект самого бейзлайна от эффектов реализации. Вывод оказался отрезвляющим: значительная часть заявленного выигрыша объяснялась не свойством action-dependent baseline как такового, а сопутствующими факторами.

• Разная нормализация advantage. Экспериментальные ветки с новым бейзлайном и без часто нормализовали преимущество по-разному, и это само по себе двигало метрики.
• Смещение из-за обучаемого бейзлайна. Теоретическая несмещённость держится на том, что bᵢ — истинная функция от a₋ᵢ. Обучаемая нейросеть-аппроксиматор эту гарантию нарушает: появляется смещение, которое может как помогать, так и вредить, но которое уже не описывается красивой теоремой.
• Оценка величины эффекта. Реальное снижение дисперсии от учёта корреляций между координатами часто мало по сравнению с дисперсией, которую даёт сама траектория и оценка отдачи Монте-Карло.

Итог дискуссии: action-dependent factorized baseline — корректная идея, но её практический выигрыш в типичных непрерывных задачах управления скромнее, чем показывали ранние прогоны, и легко перебивается менее экзотическими приёмами.

Что действительно двигает дисперсию сильнее

Если цель — стабилизировать обучение, обычно больший эффект дают вещи попроще:

1. Generalized Advantage Estimation с настройкой λ — контроль bias-variance tradeoff на уровне самой оценки преимущества.
2. Нормализация advantage внутри батча к нулевому среднему и единичной дисперсии.
3. Достаточно большой батч и корректная нормализация наблюдений и наград.
4. Хорошо обученный state-value baseline с отдельной сетью и своим learning rate.

Сравнение подходов к бейзлайну

Как это выглядит в коде

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

# policy: диагональная гауссова, mean и log_std на состояние s
# actions: [batch, action_dim], возвращённые из среды
# returns: [batch] — отдача (или GAE-advantage) на траекторию

log_probs = policy.log_prob(actions)          # [batch, action_dim]
# state-value baseline: один скаляр на состояние
adv_state = returns - value_net(states)       # [batch]
pg_state = -(log_probs.sum(dim=1) * adv_state.detach()).mean()

# action-dependent factorized: свой бейзлайн на каждую координату
# b_i зависит от s и всех действий, кроме i-го
adv_i = returns.unsqueeze(1) - baseline_net(states, actions_masked_i)
# baseline_net не должен видеть a_i для i-й компоненты — иначе смещение
pg_factored = -(log_probs * adv_i.detach()).sum(dim=1).mean()

Главный практический нюанс скрыт в actions_masked_i: сеть-бейзлайн для координаты i обязана не получать на вход aᵢ. Нарушение этого условия — самый частый способ незаметно потерять несмещённость и получить «слишком хорошие» кривые обучения.

Когда всё-таки стоит пробовать

Action-dependent baseline имеет смысл рассматривать, когда действие по-настоящему многомерно и координаты сильно скоррелированы с отдачей — например, в многоагентных задачах, где действие агента естественно разложить по агентам, и вклад отдельного агента в общую награду хочется отделить от вклада остальных. Здесь конструкция ближе к counterfactual baseline и мотивирована структурой задачи, а не только желанием сбить дисперсию.

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

Короткий вывод: теорема верна, идея красивая, но проверяйте, откуда именно берётся ваше ускорение. В RL слишком легко принять артефакт реализации за свойство алгоритма.