dreamerv2

Status: versão estável

Dominando o Atari com modelos de mundo discretos

Implementação do agente DreamerV2 no TensorFlow 2. Curvas de treinamento para todos os 55 jogos estão incluídas.

Se você achar este código útil, faça referência em seu artigo:

 @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 é o primeiro agente modelo mundial que atinge desempenho de nível humano no benchmark Atari. O DreamerV2 também supera o desempenho final dos principais agentes livres de modelos, Rainbow e IQN, usando a mesma quantidade de experiência e computação. A implementação neste repositório alterna entre treinar o modelo mundial, treinar a política e coletar experiência e é executada em uma única GPU.

O DreamerV2 aprende um modelo do ambiente diretamente a partir de imagens de entrada de alta dimensão. Para isso, ele prevê antecipadamente usando estados aprendidos compactos. Os estados consistem em uma parte determinística e diversas variáveis ​​categóricas que são amostradas. O anterior para essas categóricas é aprendido por meio de uma perda KL. O modelo mundial é aprendido de ponta a ponta por meio de gradientes diretos, o que significa que o gradiente da densidade é definido como o gradiente da amostra.

DreamerV2 aprende redes de atores e críticos a partir de trajetórias imaginadas de estados latentes. As trajetórias começam em estados codificados de sequências encontradas anteriormente. O modelo mundial então prevê o futuro usando as ações selecionadas e seu estado aprendido anteriormente. O crítico é treinado usando aprendizagem de diferença temporal e o ator é treinado para maximizar a função de valor por meio de gradientes reforçados e diretos.

Para mais informações:

• Postagem do blog sobre IA do Google
• Site do projeto
• Artigo de pesquisa

Usando o pacote

A maneira mais fácil de executar o DreamerV2 em novos ambientes é instalar o pacote via pip3 install dreamerv2 . O código detecta automaticamente se o ambiente utiliza ações discretas ou contínuas. Aqui está um exemplo de uso que treina o DreamerV2 no ambiente 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 )

Instruções manuais

Para modificar o agente DreamerV2, clone o repositório e siga as instruções abaixo. Também existe um Dockerfile disponível, caso você não queira instalar as dependências em seu sistema.

Obtenha dependências:

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

Treine no Atari:

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

Treine no controle DM:

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

Monitore os resultados:

tensorboard --logdir ~ /logdir

Gerar gráficos:

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

Instruções do Docker

O Dockerfile permite executar o DreamerV2 sem instalar suas dependências em seu sistema. Isso requer que você tenha o Docker com acesso à GPU configurado.

Verifique sua configuração:

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

Treine no 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

Treine no controle 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

Pontas

• Depuração eficiente. Você pode usar a configuração debug como em --configs atari debug . Isso reduz o tamanho do lote, aumenta a frequência de avaliação e desativa a compilação do gráfico tf.function para facilitar a depuração linha por linha.

• Normas de gradiente infinito. Isso é normal e descrito na escala de perda no guia de precisão mista. Você pode desabilitar a precisão mista passando --precision 32 para o script de treinamento. A precisão mista é mais rápida, mas pode, em princípio, causar instabilidades numéricas.

• Acessando métricas registradas. As métricas são armazenadas no formato de linhas TensorBoard e JSON. Você pode carregá-los diretamente usando pandas.read_json() . O script de plotagem também armazena as métricas agrupadas e agrupadas de várias execuções em um único arquivo JSON para facilitar a plotagem manual.