dreamerv2

Estado: lanzamiento estable

Dominar Atari con modelos de mundos discretos

Implementación del agente DreamerV2 en TensorFlow 2. Se incluyen curvas de entrenamiento para los 55 juegos.

Si encuentra útil este código, haga referencia en su artículo:

 @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}
}

Método

DreamerV2 es el primer agente modelo mundial que logra un rendimiento a nivel humano en el punto de referencia de Atari. DreamerV2 también supera el rendimiento final de los mejores agentes sin modelos Rainbow e IQN utilizando la misma cantidad de experiencia y computación. La implementación en este repositorio alterna entre entrenar el modelo mundial, entrenar la política y recopilar experiencia y se ejecuta en una sola GPU.

DreamerV2 aprende un modelo del entorno directamente a partir de imágenes de entrada de alta dimensión. Para ello, predice el futuro utilizando estados aprendidos compactos. Los estados constan de una parte determinista y varias variables categóricas que se muestrean. El previo de estas categóricas se aprende mediante una pérdida de KL. El modelo mundial se aprende de un extremo a otro mediante gradientes directos, lo que significa que el gradiente de la densidad se establece en el gradiente de la muestra.

DreamerV2 aprende redes de actores y críticos a partir de trayectorias imaginadas de estados latentes. Las trayectorias comienzan en estados codificados de secuencias encontradas previamente. Luego, el modelo mundial predice el futuro utilizando las acciones seleccionadas y su estado aprendido antes. El crítico se entrena utilizando el aprendizaje de diferencias temporales y el actor se entrena para maximizar la función de valor mediante refuerzo y gradientes directos.

Para más información:

• Publicación del blog de IA de Google
• Sitio web del proyecto
• Trabajo de investigación

Usando el paquete

La forma más sencilla de ejecutar DreamerV2 en entornos nuevos es instalar el paquete mediante pip3 install dreamerv2 . El código detecta automáticamente si el entorno utiliza acciones discretas o continuas. A continuación se muestra un ejemplo de uso que entrena a DreamerV2 en el entorno 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 )

Instrucciones manuales

Para modificar el agente DreamerV2, clone el repositorio y siga las instrucciones a continuación. También hay un Dockerfile disponible, en caso de que no desees instalar las dependencias en tu sistema.

Obtener dependencias:

pip3 install tensorflow==2.6.0 tensorflow_probability ruamel.yaml ' gym[atari] ' dm_control

Tren en Atari:

python3 dreamerv2/train.py --logdir ~ /logdir/atari_pong/dreamerv2/1 
  --configs atari --task atari_pong

Entrenar en Control DM:

python3 dreamerv2/train.py --logdir ~ /logdir/dmc_walker_walk/dreamerv2/1 
  --configs dmc_vision --task dmc_walker_walk

Monitorear resultados:

tensorboard --logdir ~ /logdir

Generar parcelas:

python3 common/plot.py --indir ~ /logdir --outdir ~ /plots 
  --xaxis step --yaxis eval_return --bins 1e6

Instrucciones de la ventana acoplable

Dockerfile le permite ejecutar DreamerV2 sin instalar sus dependencias en su sistema. Esto requiere que tengas configurado Docker con acceso a GPU.

Verifique su configuración:

docker run -it --rm --gpus all tensorflow/tensorflow:2.4.2-gpu nvidia-smi

Tren en 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

Entrenar en Control 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

Consejos

• Depuración eficiente. Puede utilizar la configuración debug como en --configs atari debug . Esto reduce el tamaño del lote, aumenta la frecuencia de evaluación y deshabilita la compilación del gráfico tf.function para facilitar la depuración línea por línea.

• Normas de gradiente infinito. Esto es normal y se describe en escala de pérdidas en la guía de precisión mixta. Puede desactivar la precisión mixta pasando --precision 32 al script de entrenamiento. La precisión mixta es más rápida pero, en principio, puede provocar inestabilidades numéricas.

• Acceder a métricas registradas. Las métricas se almacenan en formato de líneas TensorBoard y JSON. Puedes cargarlos directamente usando pandas.read_json() . El script de trazado también almacena las métricas agrupadas y agregadas de múltiples ejecuciones en un único archivo JSON para facilitar el trazado manual.