Statut : version stable
Implémentation de l'agent DreamerV2 dans TensorFlow 2. Les courbes d'entraînement pour les 55 jeux sont incluses.
Si vous trouvez ce code utile, veuillez le référencer dans votre article :
@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 est le premier agent modèle mondial à atteindre des performances de niveau humain sur le benchmark Atari. DreamerV2 surpasse également les performances finales des meilleurs agents sans modèle Rainbow et IQN en utilisant la même quantité d'expérience et de calcul. L'implémentation dans ce référentiel alterne entre la formation du modèle mondial, la formation de la politique et la collecte d'expériences et s'exécute sur un seul GPU.
DreamerV2 apprend un modèle de l'environnement directement à partir d'images d'entrée de haute dimension. Pour cela, il prédit à l’avance en utilisant des états appris compacts. Les états sont constitués d'une partie déterministe et de plusieurs variables catégorielles échantillonnées. Le prior de ces catégoriques est appris grâce à une perte de KL. Le modèle mondial est appris de bout en bout via des gradients directs, ce qui signifie que le gradient de la densité est défini sur le gradient de l'échantillon.
DreamerV2 apprend les réseaux d'acteurs et de critiques à partir de trajectoires imaginées d'états latents. Les trajectoires commencent aux états codés de séquences rencontrées précédemment. Le modèle mondial prédit ensuite à l'avance en utilisant les actions sélectionnées et son état appris auparavant. Le critique est formé à l'aide de l'apprentissage des différences temporelles et l'acteur est formé à maximiser la fonction de valeur via des gradients renforcés et directs.
Pour plus d'informations :
Le moyen le plus simple d'exécuter DreamerV2 sur de nouveaux environnements est d'installer le package via pip3 install dreamerv2
. Le code détecte automatiquement si l'environnement utilise des actions discrètes ou continues. Voici un exemple d'utilisation qui entraîne DreamerV2 sur l'environnement 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 )
Pour modifier l'agent DreamerV2, clonez le référentiel et suivez les instructions ci-dessous. Un Dockerfile est également disponible, au cas où vous ne souhaiteriez pas installer les dépendances sur votre système.
Obtenez les dépendances :
pip3 install tensorflow==2.6.0 tensorflow_probability ruamel.yaml ' gym[atari] ' dm_control
Entraînez-vous sur Atari :
python3 dreamerv2/train.py --logdir ~ /logdir/atari_pong/dreamerv2/1
--configs atari --task atari_pong
Entraînez-vous sur le contrôle DM :
python3 dreamerv2/train.py --logdir ~ /logdir/dmc_walker_walk/dreamerv2/1
--configs dmc_vision --task dmc_walker_walk
Surveiller les résultats :
tensorboard --logdir ~ /logdir
Générez des tracés :
python3 common/plot.py --indir ~ /logdir --outdir ~ /plots
--xaxis step --yaxis eval_return --bins 1e6
Le Dockerfile vous permet d'exécuter DreamerV2 sans installer ses dépendances sur votre système. Cela nécessite que vous ayez configuré Docker avec accès GPU.
Vérifiez votre configuration :
docker run -it --rm --gpus all tensorflow/tensorflow:2.4.2-gpu nvidia-smi
Entraînez-vous sur 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
Entraînez-vous sur le contrôle 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
Débogage efficace. Vous pouvez utiliser la configuration debug
comme dans --configs atari debug
. Cela réduit la taille du lot, augmente la fréquence d'évaluation et désactive la compilation de graphiques tf.function
pour un débogage facile ligne par ligne.
Normes de gradient infinies. Ceci est normal et décrit sous la rubrique mise à l’échelle des pertes dans le guide de précision mixte. Vous pouvez désactiver la précision mixte en passant --precision 32
au script de formation. La précision mixte est plus rapide mais peut en principe provoquer des instabilités numériques.
Accès aux métriques enregistrées. Les métriques sont stockées au format de lignes TensorBoard et JSON. Vous pouvez les charger directement en utilisant pandas.read_json()
. Le script de traçage stocke également les métriques regroupées et agrégées de plusieurs exécutions dans un seul fichier JSON pour faciliter le traçage manuel.