狀態:穩定發布
在 TensorFlow 2 中實作 DreamerV2 代理程式。
如果您發現此程式碼有用,請在您的論文中引用:
@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。此儲存庫中的實作在訓練世界模型、訓練策略和收集經驗之間交替進行,並在單一 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
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
高效調試。您可以使用--configs atari debug
中debug
配置。這可以減少批次大小,增加評估頻率,並停用tf.function
圖形編譯,以便於逐行除錯。
無限梯度範數。這是正常現象,並在混合精度指南中的損失縮放下進行了描述。您可以透過將--precision 32
傳遞給訓練腳本來停用混合精確度。混合精度速度更快,但原則上會導致數值不穩定。
訪問記錄的指標。指標以 TensorBoard 和 JSON 行格式儲存。您可以使用pandas.read_json()
直接載入它們。繪圖腳本還將多次運行的分箱和聚合指標儲存到單一 JSON 檔案中,以便於手動繪圖。