BonanDing/WorldMem_Repro

WorldMem

Long-term consistent world simulation with memory.

Environment (conda)

conda create -n worldmem python=3.10
conda activate worldmem
pip install -r requirements.txt
conda install -c conda-forge ffmpeg=4.3.2

Data preparation (data folder)

1. Download the Minecraft dataset: https://huggingface.co/datasets/zeqixiao/worldmem_minecraft_dataset
2. Place it under data/ with this structure:

data/
└── minecraft/
    ├── training/
    ├── validation/
    └── test/

The training and evaluation scripts expect the dataset to live at data/minecraft by default.

Checkpoints

Pretrained checkpoints are hosted on Hugging Face: zeqixiao/worldmem_checkpoints.

Example download to checkpoints/:

huggingface-cli download zeqixiao/worldmem_checkpoints diffusion_only.ckpt --local-dir checkpoints
huggingface-cli download zeqixiao/worldmem_checkpoints vae_only.ckpt --local-dir checkpoints
huggingface-cli download zeqixiao/worldmem_checkpoints pose_prediction_model_only.ckpt --local-dir checkpoints

Then point your scripts or configs to these files, for example:

python -m main +name=train +diffusion_model_path=checkpoints/diffusion_only.ckpt +vae_path=checkpoints/vae_only.ckpt

Training

Run a single stage:

sh train_stage_1.sh
sh train_stage_2.sh
sh train_stage_3.sh

Run all stages:

sh train_3stages.sh

The stage scripts include dataset and checkpoint paths. Update those paths or override them on the CLI to match your local setup.

Training config (exp_video.yaml)

Defaults live in configurations/experiment/exp_video.yaml.

Common fields to edit:

• training.lr
• training.precision
• training.batch_size
• training.max_steps
• training.checkpointing.every_n_train_steps
• validation.val_every_n_step
• validation.batch_size
• test.batch_size

You can also override values from the CLI used in the scripts:

python -m main +name=train experiment.training.batch_size=8 experiment.training.max_steps=100000

W&B run IDs: configurations/training.yaml has resume and load fields. The run ID is the short token in the run URL (for example, ot7jqmgn).