利用离散世界模型掌握 Atari
在 TensorFlow 2 中实现DreamerV2代理。包含全部 55 款游戏的训练曲线。
如果您发现此代码有用，请在您的论文中引用：
<span style="background-color:var(--bgColor-muted, var(--color-canvas-subtle))"><span style="color:#1f2328"><span style="color:var(--fgColor-default, var(--color-fg-default))"><span style="background-color:var(--bgColor-muted, var(--color-canvas-subtle))"><code>@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}
}
</code></span></span></span></span>方法
DreamerV2 是第一个在 Atari 基准上达到人类水平表现的世界模型代理。在使用相同数量的经验和计算的情况下，DreamerV2 的最终表现也优于顶级无模型代理 Rainbow 和 IQN。此存储库中的实现在训练世界模型、训练策略和收集经验之间交替进行，并在单个 GPU 上运行。
DreamerV2 直接从高维输入图像中学习环境模型。为此，它使用紧凑的学习状态进行预测。这些状态由确定性部分和几个采样的分类变量组成。这些分类的先验是通过 KL 损失学习的。世界模型是通过直通梯度端到端学习的，这意味着密度的梯度被设置为样本的梯度。
DreamerV2 从潜在状态的想象轨迹中学习参与者和评论家网络。轨迹从之前遇到的序列的编码状态开始。然后，世界模型使用选定的动作及其学习到的状态进行预测。评论家使用时间差异学习进行训练，参与者通过强化和直通梯度进行训练以最大化价值函数。
更多信息：
使用包
在新的环境中运行 DreamerV2 的最简单方法是通过 安装包pip3 install dreamerv2。代码会自动检测环境是使用离散动作还是连续动作。以下是在 MiniGrid 环境中训练 DreamerV2 的使用示例：
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 控制上进行训练：
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
Docker 说明
Dockerfile可让您运行 DreamerV2 ， 而无需在系统中安装其依赖项。这要求您设置具有 GPU 访问权限的 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 控制上进行训练：
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 文件中，以便于手动绘图。