Home Original page

GitHub - DemyCode/minecraft-copilot-ml

Minecraft Map Generation with Diffusion Models, VAE, and Flow Matching

This repository contains code for generating Minecraft maps using multiple generative approaches:

  1. Diffusion Models: A state-of-the-art generative model that gradually denoises random noise into coherent Minecraft structures.
  2. Variational Autoencoder (VAE): A neural network that encodes Minecraft maps into a latent space and decodes them back.
  3. Conditional Flow Matching (CFM): A generative model that learns to generate new Minecraft maps by modeling the flow in the latent space.

Overview

The project consists of the following components:

  1. Minecraft Dataset: A PyTorch dataset for loading Minecraft schematic files.
  2. Diffusion Model: A 3D UNet-based diffusion model for generating Minecraft structures.
  3. Variational Autoencoder (VAE): A neural network that encodes Minecraft maps into a latent space and decodes them back.
  4. Conditional Flow Matching (CFM): A generative model that learns to generate new Minecraft maps by modeling the flow in the latent space.

Requirements

  • Python 3.6+
  • PyTorch
  • NumPy
  • Matplotlib
  • tqdm
  • nbtlib (for loading schematic files)
  • scikit-learn (for t-SNE visualization)

Usage

1. Prepare the Dataset

Place your Minecraft schematic files in a directory (e.g., minecraft-schematics-raw). The dataset will automatically load and process these files.

2. Train the Diffusion Model

python main.py --batch_size 64 --learning_rate 1e-4 --diffusion_steps 1000 --noise_schedule linear

3. Sample from the Diffusion Model

python sample.py --model_path models/model000100.pt --num_samples 10

4. Train the VAE

python vae_demo.py --data_dir minecraft-schematics-raw --train --epochs 20 --save_model models/minecraft_vae.pth

5. Evaluate and Visualize the VAE

python vae_demo.py --data_dir minecraft-schematics-raw --load_model models/minecraft_vae.pth --evaluate --visualize_latent --generate --interpolate

6. Train the Flow Matching Network

python flow_matching.py --data_dir minecraft-schematics-raw --vae_model models/minecraft_vae.pth --train --epochs 20 --save_flow_model models/minecraft_flow.pth

7. Generate Samples with the Flow Matching Network

python flow_matching.py --data_dir minecraft-schematics-raw --vae_model models/minecraft_vae.pth --load_flow_model models/minecraft_flow.pth --generate --n_samples 10

Model Architecture

Diffusion Model

The diffusion model is a 3D UNet that:

  • Takes a 3D tensor of shape [batch_size, num_blocks, 16, 16, 16] as input
  • Gradually denoises random noise into coherent Minecraft structures
  • Uses a mask for conditioning on valid positions
  • Employs a standard diffusion process with a noise schedule

Variational Autoencoder (VAE)

The VAE consists of:

  • An encoder that maps Minecraft maps to a latent space
  • A decoder that reconstructs maps from the latent space
  • A reparameterization trick for sampling from the latent space

The VAE takes into account the mask to handle variable-sized inputs.

Conditional Flow Matching (CFM)

The CFM learns to model the vector field of a continuous normalizing flow in the latent space. It can be conditioned on additional context vectors for controlled generation.

Files

  • minecraft_dataset.py: PyTorch dataset for Minecraft schematic files
  • schematic_loader.py: Utility for loading schematic files
  • main.py: Main script for training the diffusion model
  • sample.py: Script for sampling from the trained diffusion model
  • minecraft_vae.py: Implementation of the Variational Autoencoder
  • vae_demo.py: Demo script for the VAE
  • flow_matching.py: Implementation of the Conditional Flow Matching network
  • test_dataset.py: Test script for the dataset

Examples

Diffusion Model Samples

The diffusion model can generate samples by gradually denoising random noise:

# Sample from the diffusion model
samples = diffusion.p_sample_loop(
    model,
    (batch_size, num_blocks, 16, 16, 16),
    clip_denoised=True,
    model_kwargs={"mask": mask},
    device="cuda",
)

VAE Samples

The VAE can generate samples by sampling from the latent space and decoding:

# Sample from the VAE
samples = vae.sample(num_samples=5, device="cuda")

Flow Matching Samples

The Flow Matching network can generate samples by solving the ODE:

# Generate samples with the Flow Matching network
samples = generate_samples_with_flow(vae, flow_model, device="cuda", num_samples=5)

Customization

You can customize the models by adjusting the following parameters:

Diffusion Model Parameters

  • model_channels: Number of channels in the UNet model
  • num_res_blocks: Number of residual blocks in each UNet layer
  • attention_resolutions: Resolutions at which to apply attention
  • dropout: Dropout rate
  • channel_mult: Channel multiplier for each UNet layer
  • diffusion_steps: Number of diffusion steps
  • noise_schedule: Type of noise schedule ("linear" or "cosine")

VAE Parameters

  • latent_dim: Dimension of the latent space
  • embedding_dim: Dimension of the block embeddings
  • hidden_dims: Dimensions of the hidden layers

Flow Matching Parameters

  • context_dim: Dimension of the context vector for conditional generation
  • hidden_dims: Dimensions of the hidden layers in the flow model

License

This project is licensed under the MIT License - see the LICENSE file for details.