Статус: Стабильная версия
Реализация агента DreamerV2 в TensorFlow 2. Включены кривые обучения для всех 55 игр.
Если вы найдете этот код полезным, укажите его в своей статье:
@article{hafner2020dreamerv2,
title={Mastering Atari with Discrete World Models},
author={Hafner, Danijar and Lillicrap, Timothy and Norouzi, Mohammad and Ba, Jimmy},
journal={arXiv preprint arXiv:2010.02193},
year={2020}
}
DreamerV2 — первый агент мировой модели, который достигает производительности человеческого уровня в тестах Atari. DreamerV2 также превосходит по конечной производительности лучших агентов без моделей Rainbow и IQN, используя тот же объем опыта и вычислений. Реализация в этом репозитории чередуется между обучением модели мира, обучением политики и сбором опыта и выполняется на одном графическом процессоре.
DreamerV2 изучает модель окружающей среды непосредственно из входных многомерных изображений. Для этого он прогнозирует будущее, используя компактные изученные состояния. Состояния состоят из детерминированной части и нескольких категориальных переменных, которые выбираются. Приоритет для этих категориалов изучается посредством потери KL. Модель мира изучается сквозным образом посредством сквозных градиентов, что означает, что градиент плотности устанавливается равным градиенту выборки.
DreamerV2 изучает сети актеров и критиков на основе воображаемых траекторий скрытых состояний. Траектории начинаются с закодированных состояний ранее встречавшихся последовательностей. Затем модель мира прогнозирует будущее, используя выбранные действия и предварительно изученное состояние. Критик обучается с использованием обучения временным разницам, а актер обучается максимизировать функцию ценности с помощью подкрепления и сквозных градиентов.
Для получения дополнительной информации:
Самый простой способ запустить DreamerV2 в новых средах — установить пакет с помощью pip3 install dreamerv2
. Код автоматически определяет, использует ли среда дискретные или непрерывные действия. Вот пример использования, который обучает DreamerV2 в среде MiniGrid:
import gym
import gym_minigrid
import dreamerv2 . api as dv2
config = dv2 . defaults . update ({
'logdir' : '~/logdir/minigrid' ,
'log_every' : 1e3 ,
'train_every' : 10 ,
'prefill' : 1e5 ,
'actor_ent' : 3e-3 ,
'loss_scales.kl' : 1.0 ,
'discount' : 0.99 ,
}). parse_flags ()
env = gym . make ( 'MiniGrid-DoorKey-6x6-v0' )
env = gym_minigrid . wrappers . RGBImgPartialObsWrapper ( env )
dv2 . train ( env , config )
Чтобы изменить агент DreamerV2, клонируйте репозиторий и следуйте инструкциям ниже. Также доступен Dockerfile на случай, если вы не хотите устанавливать зависимости в своей системе.
Получить зависимости:
pip3 install tensorflow==2.6.0 tensorflow_probability ruamel.yaml ' gym[atari] ' dm_control
Тренируйтесь на Atari:
python3 dreamerv2/train.py --logdir ~ /logdir/atari_pong/dreamerv2/1
--configs atari --task atari_pong
Тренируйтесь на DM Control:
python3 dreamerv2/train.py --logdir ~ /logdir/dmc_walker_walk/dreamerv2/1
--configs dmc_vision --task dmc_walker_walk
Мониторинг результатов:
tensorboard --logdir ~ /logdir
Генерировать графики:
python3 common/plot.py --indir ~ /logdir --outdir ~ /plots
--xaxis step --yaxis eval_return --bins 1e6
Dockerfile позволяет запускать DreamerV2 без установки его зависимостей в вашей системе. Для этого необходимо, чтобы у вас был настроен Docker с доступом к графическому процессору.
Проверьте вашу настройку:
docker run -it --rm --gpus all tensorflow/tensorflow:2.4.2-gpu nvidia-smi
Тренируйтесь на Atari:
docker build -t dreamerv2 .
docker run -it --rm --gpus all -v ~ /logdir:/logdir dreamerv2
python3 dreamerv2/train.py --logdir /logdir/atari_pong/dreamerv2/1
--configs atari --task atari_pong
Тренируйтесь на DM Control:
docker build -t dreamerv2 . --build-arg MUJOCO_KEY= " $( cat ~ /.mujoco/mjkey.txt ) "
docker run -it --rm --gpus all -v ~ /logdir:/logdir dreamerv2
python3 dreamerv2/train.py --logdir /logdir/dmc_walker_walk/dreamerv2/1
--configs dmc_vision --task dmc_walker_walk
Эффективная отладка. Вы можете использовать конфигурацию debug
, как в --configs atari debug
. Это уменьшает размер пакета, увеличивает частоту оценки и отключает компиляцию графа tf.function
для упрощения построчной отладки.
Бесконечные градиентные нормы. Это нормально и описано в разделе «Масштабирование потерь» в руководстве смешанной точности. Вы можете отключить смешанную точность, передав --precision 32
в сценарий обучения. Смешанная точность быстрее, но в принципе может вызвать числовую нестабильность.
Доступ к зарегистрированным метрикам. Метрики хранятся как в формате строк TensorBoard, так и в формате JSON. Вы можете загрузить их напрямую, используя pandas.read_json()
. Сценарий построения графиков также сохраняет объединенные и агрегированные показатели нескольких запусков в один файл JSON для упрощения построения графиков вручную.