FB Sequence Planner

Планирование последовательностей на основе замороженных Forward-Backward (FB) представлений для задач долгосрочного zero-shot обучения с подкреплением.

Этот репозиторий расширяет кодовую базу работы Switching Successor Measures for Hierarchical Zero-shot Reinforcement Learning высокоуровневым планировщиком, который явно оценивает последовательности латентных намерений, вместо выбора только одного намерения за один шаг.

Основной эксперимент проводится в среде ogbench-antmaze-medium-navigate-v0 с использованием предоставленного предобученного FB-чекпоинта. FB-представление и низкоуровневая политика остаются замороженными; предложенный планировщик не вводит дополнительных обучаемых моделей и использует только офлайн-датасет OGBench для построения кандидатных последовательностей путевых точек.

Основной результат

Средний success rate по 5 seed, 5 задачам и 5 эпизодам на каждую задачу:

Метод	Средний success rate
Single-intention baseline	0.776
FB Sequence Planner, H=4	0.784
FB Sequence Planner, H=8	0.816

Лучшая конфигурация повышает средний success rate на 4.0 процентных пункта по сравнению с предоставленным single-intention baseline.

Метод

Для текущего состояния s и латентного представления задачи z_r планировщик:

1. извлекает близкие состояния из офлайн-датасета, сохраняя кандидаты из разных траекторий;
2. строит последовательности путевых точек фиксированной длины из реальных фрагментов офлайн-траекторий;
3. отображает каждую путевую точку w_i в намерение z_i = B(w_i);
4. оценивает FB-reachability между последовательными намерениями;
5. рекурсивно оценивает полную последовательность с помощью multi-switch FB-разложения значения;
6. использует минимальную оценку по двум головам FB-ансамбля как консервативную оценку последовательности;
7. исполняет первое намерение лучшей валидной последовательности с положительным advantage;
8. выполняет перепланирование через фиксированный высокоуровневый интервал.

Если ни одна валидная кандидатная последовательность не имеет положительного robust advantage, контроллер возвращается к предоставленной single-intention high-level политике.

Реализация планировщика находится в:

planning/fb_sequence_planner.py

Логика исполнения интегрирована в:

utils/evaluation.py

Установка

Исходная кодовая база использует Python 3.11. Следующая настройка соответствует оригинальному репозиторию.

conda create -n ssm python=3.11
conda activate ssm
pip install -r requirements.txt

Датасеты OGBench загружаются автоматически при первом запуске и обычно сохраняются в:

~/.ogbench/data

Для установки JAX с поддержкой GPU может потребоваться настройка CUDA, зависящая от конкретной машины.

Предобученный чекпоинт

В экспериментах используется предоставленный предобученный чекпоинт для среды AntMaze Medium.

Разместите чекпоинт так, чтобы был доступен путь:

checkpoints/medium

Файлы чекпоинта не обязательно добавлять в этот репозиторий, если они распространяются отдельно вместе с заданием или оригинальной реализацией.

Запуск sequence planner

Финальная конфигурация с H=8 задаётся в main.py следующим образом:

planner = FBSequencePlanner(
    agent=agent,
    dataset=train_dataset,
    planning_horizon=8,
    waypoint_spacing=1.0,
    num_retrieved_trajectories=10,
    retrieval_pool_size=5_000,
)

Для запуска одного evaluation seed используйте:

python main.py \
  --env_name=ogbench-antmaze-medium-navigate-v0 \
  --agent=agents/fbpiswitch.py \
  --restore_path=checkpoints/medium \
  --train_steps=1 \
  --agent.lr=0.0 \
  --eval_episodes=5 \
  --video_episodes=0 \
  --enable_wandb=0 \
  --seed=0

Повторите запуск для:

--seed=0
--seed=1
--seed=2
--seed=3
--seed=4

Для абляции с H=4 измените только:

planning_horizon=4

Остальные параметры планировщика оставьте без изменений.

Запуск single-intention baseline

Baseline — это исходный высокоуровневый контроллер FB pi-Switch без sequence planner.

В вызове evaluate(...) в main.py укажите:

planner=None

вместо:

planner=planner

После этого запустите ту же команду:

python main.py \
  --env_name=ogbench-antmaze-medium-navigate-v0 \
  --agent=agents/fbpiswitch.py \
  --restore_path=checkpoints/medium \
  --train_steps=1 \
  --agent.lr=0.0 \
  --eval_episodes=5 \
  --video_episodes=0 \
  --enable_wandb=0 \
  --seed=0

И аналогично повторите для seed 0--4.

Протокол оценки

Для каждого метода:

• среда: ogbench-antmaze-medium-navigate-v0
• seed: 0, 1, 2, 3, 4
• задачи: все 5 evaluation-задач AntMaze-medium
• количество эпизодов на задачу: 5
• общее количество эпизодов на метод: 125
• метрика: success rate

Предобученное FB-представление и низкоуровневая политика остаются замороженными во время оценки. Дополнительное взаимодействие со средой для обучения sequence planner не используется.

Параметры планировщика

Параметр	Значение
Расстояние между путевыми точками	1.0
Интервал перепланирования	10 environment steps
Максимальное значение reachability	1.1
Максимальное расхождение между головами ансамбля	0.25
Число голов FB-ансамбля	2
Проверенные горизонты планирования	4, 8

Происхождение репозитория

Этот проект основан на официальной реализации работы:

Stefan Stojanovic and Alexandre Proutiere, Switching Successor Measures for Hierarchical Zero-shot Reinforcement Learning, 2026.

Оригинальный репозиторий:

https://github.com/stestoKTH/switching-successor-measures

Страница проекта:

https://stestokth.github.io/switching-successors/

Вклад данного репозитория заключается в sequence-aware высокоуровневом механизме планирования и его оценке поверх замороженных предобученных FB-компонентов.

Благодарности

Проект основан на оригинальной реализации Switching Successor Measures и OGBench.

• Switching Successor Measures: https://github.com/stestoKTH/switching-successor-measures
• OGBench: https://github.com/seohongpark/ogbench

Лицензия

Оригинальный репозиторий распространяется под лицензией MIT. При распространении модифицированных версий необходимо сохранять исходную лицензию и указание авторства.