Temporal segment models: как предсказывать не шаги, а куски

Проблема, которую решают сегментные модели

Классическая модель динамики в обучении с подкреплением предсказывает один шаг: по состоянию s_t и действию a_t выдаёт s_{t+1}. Чтобы спланировать поведение на горизонте в 50 шагов, вы прогоняете эту модель 50 раз подряд, подставляя собственный прогноз на вход. Каждый прогон добавляет ошибку, и к середине горизонта предсказанная траектория уходит в область, которую модель никогда не видела в данных. Это называют накоплением ошибки (compounding error), и именно оно ограничивает длину планирования в model-based RL.

Идея temporal segment models — учить модель предсказывать не один переход, а целый отрезок траектории фиксированной длины H сразу, как единый объект. Вместо цепочки из H однокадровых прогнозов вы получаете один прогноз на H шагов вперёд. Ошибка при этом не накручивается автогрессивно: модель не подставляет свой же выход себе на вход внутри сегмента.

Откуда термин

Работа Nikhil Mishra, Pieter Abbeel и Igor Mordatch «Prediction and Control with Temporal Segment Models» (ICML 2017) предложила рассматривать участок траектории как латентную переменную и обучать порождающую модель на уровне сегментов, а не переходов. Дальнейшие работы по латентному планированию (PlaNet, Dreamer) развивают близкую линию: свернуть длинную динамику в компактное латентное представление и планировать в нём.

Как устроен сегментный прогноз

Ключевая конструкция — латентная переменная z, которая кодирует весь сегмент (s_t, a_t, ..., s_{t+H}). Обучение идёт по схеме вариационного автоэнкодера:

• Энкодер q(z | segment) сжимает наблюдаемый отрезок траектории в распределение над z.
• Декодер p(segment | z, actions, s_t) восстанавливает состояния сегмента по латентному коду, начальному состоянию и последовательности действий.
• Приор p(z | s_t) задаёт, какие сегменты правдоподобны из данного начала, чтобы во время планирования можно было сэмплировать z без доступа к будущему.

Оптимизируется нижняя оценка правдоподобия (ELBO*): реконструкция сегмента плюс KL-расхождение между энкодером и приором. После обучения для прогноза вам достаточно взять начальное состояние, план действий, сэмплировать z из приора и декодировать сразу весь отрезок.

Одношаговая модель отвечает на вопрос «что будет через шаг». Сегментная — на вопрос «как выглядит правдоподобный кусок будущего целиком». Второе ближе к тому, что нужно планировщику.

Почему это уменьшает накопление ошибки

Автогрессивный прогноз ошибается, потому что обучается на реальных данных, а исполняется на собственных выходах — возникает разрыв распределений (distribution shift). Сегментная модель декодирует H шагов из одного латентного кода, не пропуская промежуточные состояния через свой же вход. Ошибка на 40-м шаге не рождается из ошибки на 39-м — она рождается из неточности единого кода z. Это не отменяет ошибку, но меняет её характер: она не растёт экспоненциально по длине горизонта.

Планирование поверх сегментной модели

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

1. Сэмплируете кандидатов. Генерируете множество последовательностей действий длиной H — случайно или из текущей политики.
2. Прогнозируете сегменты. Для каждого кандидата сэмплируете z из приора и декодируете весь отрезок за один проход.
3. Оцениваете. Считаете суммарную награду по предсказанному сегменту.
4. Отбираете и уточняете. В духе Cross-Entropy Method** берёте топ кандидатов, пересчитываете распределение действий, повторяете несколько итераций.
5. Исполняете первый шаг. Применяете первое действие лучшего плана, наблюдаете реальное состояние и перепланируете (model predictive control).

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

С чем сравнивать по стоимости

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

Где это работает, а где нет

Сегментные модели показывают себя на задачах непрерывного управления с плавной динамикой — робототехнические симуляторы, локомоция, манипуляции. Там фиксированный горизонт H и допущение о гладкости латентного кода оправданы.

Слабые места стоит держать в голове:

• Фиксированная длина сегмента. H — гиперпараметр. Слишком короткий сегмент возвращает вас к накоплению ошибки при склейке; слишком длинный тяжело выучить.
• Разрывная динамика. Контакты, столкновения, дискретные переключения режимов ломают предположение о гладком отрезке — единый латентный код плохо кодирует скачок.
• Стоимость обучения. Вариационная модель сегментов обучается сложнее одношаговой: нужно балансировать реконструкцию и KL, следить за коллапсом приора.
• Не заменяет исследование. Модель хороша ровно настолько, насколько разнообразны данные. На незнакомых участках среды сегментный прогноз так же ненадёжен, как любой другой.

* ELBO (Evidence Lower Bound) — нижняя оценка логарифма правдоподобия данных, которую максимизируют при обучении вариационных моделей вместо самого правдоподобия, вычислить которое напрямую нельзя.

** Cross-Entropy Method (CEM) — итеративный метод оптимизации: сэмплируем кандидатов из распределения, отбираем лучшие, пересчитываем распределение по ним и повторяем, постепенно стягивая его к хорошим решениям.