Diffusion Models Are Real-Time Game Engines reproduced by Masao Taketani

The above GIFs (300 steps) show the results of the autoregressively generated environment with the trained models from this repo.

Warning

The quality of the model is assumed to be far from the original one, but I am releasing this repo because I am resigning from my current position at the company, HEROZ, Inc., which has let me use their GPU resource for my personal R&D. So, big thanks to them!

Note

If you would just like to check the behavior of the trained models, skip to Run autoregressive inference or Run playable environment. The necessary steps to run the inference codes are written there.

Note

As for the diffusion model traininig and the VAE decoder fine-tuning, I have used one or two NVIDIA A100 80GB GPU(s). So, please change the batch size or other settings written on the training scripts if you don't have the same GPU(s).

Warning

As for the data collection process, you are supposed to generate lots of data (about 5 TB by default if I remember correctly). If you don't have enough space for that, modify the total training time steps accordingly.

This is an unofficial repo of GameNGen. I have referred to arnaudstiegler/gameNgen-repro in order to create this repo.

Modifications

Here are the list of overall modifications from arnaudstiegler/gameNgen-repro. For more details, please look at each item listed in To Do.

• Fixed the environment setup with Anaconda and write down the precise commands in order to properly reproduce the code (For this development, I start from continuumio/anaconda3:latest Docker image). Thus, to precisely replicate the dev environment, set up the Docker environment with the following command.

cd GameNGen
docker run --log-opt max-size=10m --log-opt max-file=2 -it --rm --gpus '"device={input your GPU device ID(s) here}"' -v .:/work continuumio/anaconda3:latest

Although --log-opt max-file is optional, I strongly recommend using --log-opt max-size option since the RL training script for ViZDoom outputs lots of logs to stdout, which conssumes excessive disk space.

• Use the original paper's training setup as much as I can
• Incorporate much more efficient diffusion training method and data files for better performance
• Add a playable generative environment script

To Do

• Set MOVE_LEFT and MOVE_RIGHT as independent actions
• Add no-op action to the action space
• Enable distributed training
• Implement playable generative environments
• Try using the original settings
  • VizDoom
  • Add MOVE_BACKWARD action
  • Enable to collect VizDoom dataset with multiprocessing during RL policy training
  • Diffusion Models
  • [] Use 4 for the number of inference steps for training, evaluation, and autoregressive inference
    - [x] 2e-5 for the learning rate
    - [x] Adafactor optimizer
    (I've tried the above two, but got worse results compared to the settings used for arnaudstiegler/gameNgen-repro. So, I didn't use those at the end)
  • A context length of 64
  • Train iteratively with respect to steps, not epochs
    
  • Wrap unet and action_embedding as one model in order to synchronize those two models' parameters for multi-gpu training. For more details, refer to the issue
  • Pad every image frame to get 320x256 from 320x240
  • Create dataloader that makes sure it won't pick frames from two distinct episodes when creating each training sample
  • Zero out conditioning images in latent space instead of image space for CFG
  • Make conditioning noise level adjustable during inference
• Efficient Diffusion Training
  • Convert image dataset into latent embedding one and train only using the embeddings from the beginning, not handling image dataset at all during the training so that you can much more effectively utilize GPU memory and use much larger batch size
  • Convert dataset file format from .parquet (for images & actions) to .pt (for latent embeddings & actions) to directly handle Torch tensors to train

Trained Models

All the trained models except the policy model for VizDoom, which is located in ViZDoomPPO/logs/models/deathmatch_simple/best_model.zip, are available on Hugging Face Hub:

• Finetuned Decoder
• Diffusion Model

Datasets

The datasets I provide here are meant just for inference because the original training datasets are too large to upload to HF Hub. As for training datasets, you should generate them locally with the script written at Generate training data for diffusion models while training an RL agent and Convert image dataset to latent embedding one, but with the datasets provided here, you can also check how the training scripts work before seriously conducting the heavy lifting.

• Pixel image dataset used during inference
• Latent image dataset used during inference

Scripts

Generate training data for diffusion models while training an RL agent

First, follow the commands below in order to create an environment to train Vizdoom agent.

conda create -n vizdoom python=3.11 -y
conda activate vizdoom
apt update && apt install libgl1 swig g++ -y
pip install setuptools==65.5.0 pip==21 wheel==0.38.0
cd ViZDoomPPO/
pip install -r vizdoom_requirements.txt

Then, run the following command to train an agent and build the dataset on vizdoom at the same time.

python train_ppo_and_collect_data_parallel.py --out_base_dir [dataset directory path]

Once the agent is trained, generate episodes and upload them as a HF dataset using:

python load_model_generate_dataset.py --episodes {number of episodes} --output parquet --upload --hf_repo {name of the repo}

Note: you can also generate a gif file to QA the behavior of the agent by running:

python load_model_generate_dataset.py --episodes 1 --output gif

Train the diffusion model

Second step is that follow the commands below in order to create an environment to train diffusion model.

conda deactivate
conda create -n diffusion python=3.11 -y
conda activate diffusion
pip install -r diffusion_requirements.txt

Convert image dataset to latent embedding one

You are going to convert the collected image dataset (.parquet files) into latent embedding one (.pt files), so that you can save lots of memory during diffusion training. For that, you are recommended to use at least one GPU. Use the following command to execute with single GPU.

python encode_images.py --dataset_basepath [directory path under which parquet files are placed] --save_dir_path [pt file save directory path] --dataloader_num_workers [number of workers for dataloader] --batch_size [batch size to process for one step] --dtype [data type used for inference]

You can also use multiple GPUs to process the conversion in parallel (each GPU processes handles a different episode). In order to do that, specify number of chunks for parallel processing. Also, specify which chunk ID each GPU process needs to handle. Lastly, assign GPU ID to run the command. For example, if you decide to use 3 GPUs, the number of chunks should be 3, chunk ID and GPU ID should be a number between 0 and 2. Use the following command to execute with multiple GPUs.

python encode_images.py --dataset_basepath [directory path under which parquet files are placed] --save_dir_path [pt file save directory path] --dataloader_num_workers [number of workers for dataloader] --batch_size [batch size to process for one step] --dtype [data type used for inference] --num_chunks [number of chunks to split your dataset] --chunk_id [chunk ID] --gpu_id [GPU ID]

Single GPU

If you only have single gpu, after modifying some arguments (such as dataset path) of train_diffusion_scripts/single_gpu_latents.sh, follow the instruction below to train the diffusion model.

sh train_diffusion_scripts/single_gpu_latents.sh

Multiple GPUs

Warning

As for the diffusion model training, even if you have mutiple GPUs available, the bottleneck for training speed with the current program is the data loading part. So, you may not efficiently be able to train the model with multiple GPUs. Thus, I totally recommend training it with single GPU at first.

If you have more than single GPU, after modifying some arguments (such as dataset path) of train_diffusion_scripts/multi_gpus_latents.sh, follow the instruction below to train the diffusion model.

sh train_diffusion_scripts/multi_gpus_latents.sh

Fine-tune the latent decoder

Single GPU

After modifying some arguments (such as dataset path) of finetune_vae_scripts/single_gpus.sh, follow the instruction below to finetune the decoder of VAE.

sh finetune_vae_scripts/single_gpus.sh

Multiple GPUs

If you have more than single GPU, after modifying some arguments (such as dataset path) of finetune_vae_scripts/multi_gpus.sh, follow the instruction below to finetune the decoder of VAE.

sh finetune_vae_scripts/multi_gpus.sh

Run autoregressive inference

Tip

Since the memory requirement gets bigger if you start with pixels instead of latents, I recommend not raising the flag --start_from_pixels if you don't have good amount of VRAM.

If you would just like to see how the trained model behave, download either pixel or latent inference dataset to run the autoregressive inference script.

Download the pixel dataset with the following command.

hf download Masao-Taketani/vizdoom-inference-pixel-dataset --local-dir [specify your local dir to download] --repo-type dataset

Or, download the latent dataset with the following command.

hf download Masao-Taketani/vizdoom-inference-latent-dataset --local-dir [specify your local dir to download] --repo-type dataset

The following command will generate rollouts, where each new frame is generated by the model conditioned on the previous frames and actions. We initially fill the buffer using 64 frames of training data and their corresponding actions, and only use actions from the dataset after that (i.e it matches what the agent did in the episode). Here are the explanations of arguments that are used when executing.

• start_from_pixels (store_true): Start autoregressive inference using original pixel images. If not raised, the program starts autoregressive inference using latent images
• dataset_basepath (str): Specify your parquet directory path if --start_from_pixels flag is raised. Otherwise, specify your pt directory path
• num_episodes (int): Total number of episodes to generate
• episode_length (int): Total steps for each episode to generate
• unet_model_folder (str): Specify your trained U-net folder. If not specified, the code will download the trained diffusion model
• vae_ft_model_folder (str): Specify your finetuned VAE folder. If not specified, the code will download the finetuned VAE model
• num_inference_steps (int): Number of inference steps to generate each frame
• discretized_noise_level (int): Discretized noise level used for noise augmentation. It takes values from 0 to 9 where 0 means the model is conditioned on the info that the smallest noise is added and 9 means the largest noise is added

python run_autoregressive.py --dataset_basepath [your parquet or pt path] --num_episodes [mumber of episodes] --episode_length [number of steps] --unet_model_folder [your unet folder] --vae_ft_model_folder [your vae folder] --num_inference_steps [number of inference steps] --discretized_noise_level [noise level used for noise augmentation] ([optional boolean flag] --start_from_pixels)

The generated images are saved as a GIF for each episode in the folder named rollouts by default, but you can modify the output directory with an argument named gif_outdir.

Run playable environment

You can also run the playable environment. In oder to do that, you need additional installation of the following modules.

apt -y update
apt -y install libopencv-dev
pip install keyboard

After the installation, use the following command to start running the playable environment.

Warning

1. If you are running the command remotely without a monitor, always use --conduct_headless_test argument to just test the script. In this case, you can also use --action_key_for_headless argument to specify an action key to use it repeatly without a monitor.
2. I haven't tested this playable environment script with UI since I don't have a capable GPU to run the code locally. Thus, there might be some bugs. So, please be aware of that.

python run_playable_env.py --dataset_basepath [your parquet or pt path] --num_inference_steps [number of inference steps] --num_episode_steps [number of steps]

Optionally, you can save the rollout as video (mp4) or GIF and also the action log you used for the rollout. Please refer to the following argument explanation.

• --cv2_rec: A mp4 file is saved
• --gif_rec: A GIF file is saved
• --rec_path_wo_ext: A base file name to be used to save the recording (without file extension)
• --action_log_dir: A directory path to save the action log. If None, it is not saved (default)

References

Paper

• Diffusion Models Are Real-Time Game Engines

GitHub Repos

• arnaudstiegler/gameNgen-repro
• lkiel/rl-doom