Package: dreamerv2
Status:Stable releaseMastering Atari with Discrete World ModelsImplementation of theDreamerV2agent in TensorFlow 2. Training
curves for all 55 games are included.If you find this code useful, please reference in your paper:@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}
}MethodDreamerV2 is the first world model agent that achieves human-level performance
on the Atari benchmark. DreamerV2 also outperforms the final performance of the
top model-free agents Rainbow and IQN using the same amount of experience and
computation. The implementation in this repository alternates between training
the world model, training the policy, and collecting experience and runs on a
single GPU.DreamerV2 learns a model of the environment directly from high-dimensional
input images. For this, it predicts ahead using compact learned states. The
states consist of a deterministic part and several categorical variables that
are sampled. The prior for these categoricals is learned through a KL loss. The
world model is learned end-to-end via straight-through gradients, meaning that
the gradient of the density is set to the gradient of the sample.DreamerV2 learns actor and critic networks from imagined trajectories of latent
states. The trajectories start at encoded states of previously encountered
sequences. The world model then predicts ahead using the selected actions and
its learned state prior. The critic is trained using temporal difference
learning and the actor is trained to maximize the value function via reinforce
and straight-through gradients.For more information:Google AI Blog postProject websiteResearch paperUsing the PackageThe easiest way to run DreamerV2 on new environments is to install the package
viapip3 install dreamerv2. The code automatically detects whether the
environment uses discrete or continuous actions. Here is a usage example that
trains DreamerV2 on the MiniGrid environment:importgymimportgym_minigridimportdreamerv2.apiasdv2config=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)Manual InstructionsTo modify the DreamerV2 agent, clone the repository and follow the instructions
below. There is also a Dockerfile available, in case you do not want to install
the dependencies on your system.Get dependencies:pip3installtensorflow==2.6.0tensorflow_probabilityruamel.yaml'gym[atari]'dm_controlTrain on Atari:python3dreamerv2/train.py--logdir~/logdir/atari_pong/dreamerv2/1\--configsatari--taskatari_pongTrain on DM Control:python3dreamerv2/train.py--logdir~/logdir/dmc_walker_walk/dreamerv2/1\--configsdmc--taskdmc_walker_walkMonitor results:tensorboard--logdir~/logdirGenerate plots:python3common/plot.py--indir~/logdir--outdir~/plots\--xaxisstep--yaxiseval_return--bins1e6Docker InstructionsTheDockerfilelets you run DreamerV2 without installing its dependencies in your system. This
requires you to have Docker with GPU access set up.Check your setup:dockerrun-it--rm--gpusalltensorflow/tensorflow:2.4.2-gpunvidia-smiTrain on Atari:dockerbuild-tdreamerv2.
dockerrun-it--rm--gpusall-v~/logdir:/logdirdreamerv2\python3dreamerv2/train.py--logdir/logdir/atari_pong/dreamerv2/1\--configsatari--taskatari_pongTrain on DM Control:dockerbuild-tdreamerv2.--build-argMUJOCO_KEY="$(cat~/.mujoco/mjkey.txt)"dockerrun-it--rm--gpusall-v~/logdir:/logdirdreamerv2\python3dreamerv2/train.py--logdir/logdir/dmc_walker_walk/dreamerv2/1\--configsdmc--taskdmc_walker_walkTipsEfficient debugging.You can use thedebugconfig as in--configs atari debug. This reduces the batch size, increases the evaluation
frequency, and disablestf.functiongraph compilation for easy line-by-line
debugging.Infinite gradient norms.This is normal and described under loss scaling in
themixed precisionguide. You can disable mixed precision by passing--precision 32to the training script. Mixed precision is faster but can in
principle cause numerical instabilities.Accessing logged metrics.The metrics are stored in both TensorBoard and
JSON lines format. You can directly load them usingpandas.read_json(). The
plotting script also stores the binned and aggregated metrics of multiple runs
into a single JSON file for easy manual plotting.