update trainning pipeline and datasets

README.md

@@ -39,6 +39,7 @@ https://github.com/user-attachments/assets/fb8a32e2-9470-4819-a93d-c38caf76d72c
 conda create python=3.10 -n worldmem
 conda activate worldmem
 pip install -r requirements.txt
+conda install -c conda-forge ffmpeg=4.3.2
 ```
 
 
@@ -48,11 +49,81 @@ pip install -r requirements.txt
 python app.py
 ```
 
+## Training and Inference
+
+To enable cloud logging with [Weights & Biases (wandb)](https://wandb.ai/site), follow these steps:
+
+1. Sign up for a wandb account.
+2. Run the following command to log in:
+
+    ```bash
+    wandb login
+    ```
+
+3. Open `configurations/training.yaml` and set the `entity` and `project` field to your wandb username.
+
+---
+
+### Training
+
+We observe that gradually increasing task difficulty improves performance. Thus, we adopt a multi-stage training strategy:
+
+```bash
+sh train_stage_1.sh   # Small range, no vertical turning
+sh train_stage_2.sh   # Large range, no vertical turning
+sh train_stage_3.sh   # Large range, with vertical turning
+```
+
+To resume training from a previous checkpoint, configure the `resume` and `output_dir` variables in the corresponding `.sh` script.
+
+---
+
+### Inference
+
+To run inference:
+
+```bash
+sh infer.sh
+```
+
+You can either **load the diffusion model and VAE separately**:
+
+```bash
++diffusion_model_path=yslan/worldmem_checkpoints/diffusion_only.ckpt \
++vae_path=yslan/worldmem_checkpoints/vae_only.ckpt \
++customized_load=true \
++seperate_load=true \
+```
+
+Or **load a combined checkpoint**:
+
+```bash
++load=your_model_path \
++customized_load=true \
++seperate_load=false \
+```
+
+---
+
+## Dataset
+
+Download the Minecraft dataset from [Hugging Face](https://huggingface.co/datasets/zeqixiao/worldmem_minecraft_dataset)
+
+Place the dataset in the following directory structure:
+
+```
+data/
+└── minecraft/
+    ├── training/
+    └── validation/
+```
+
+
 ## TODO
 
 - [x] Release inference models and weights;
-- [ ] Release training pipeline on Minecraft;
-- [ ] Release training data on Minecraft;
+- [x] Release training pipeline on Minecraft;
+- [x] Release training data on Minecraft;
 
 
 

app.py

@@ -23,11 +23,23 @@ from huggingface_hub import hf_hub_download
 import tempfile
 import os
 import requests
+from huggingface_hub import model_info
+
+
+def is_huggingface_model(path: str) -> bool:
+    hf_ckpt = str(path).split('/')
+    repo_id = '/'.join(hf_ckpt[:2])
+    try:
+        model_info(repo_id)
+        return True
+    except:
+        return False
+    
 
 torch.set_float32_matmul_precision("high")
 
 def load_custom_checkpoint(algo, checkpoint_path):
-    try:
+    if is_huggingface_model(str(checkpoint_path)):
         hf_ckpt = str(checkpoint_path).split('/')
         repo_id = '/'.join(hf_ckpt[:2])
         file_name = '/'.join(hf_ckpt[2:])
@@ -35,9 +47,17 @@ def load_custom_checkpoint(algo, checkpoint_path):
                             filename=file_name)
         ckpt = torch.load(model_path, map_location=torch.device('cpu'))
 
-        algo.load_state_dict(ckpt['state_dict'], strict=True)
+        filtered_state_dict = {}
+        for k, v in ckpt['state_dict'].items():
+            if "frame_timestep_embedder" in k:
+                new_k = k.replace("frame_timestep_embedder", "timestamp_embedding")
+                filtered_state_dict[new_k] = v
+            else:
+                filtered_state_dict[k] = v
+
+        algo.load_state_dict(filtered_state_dict, strict=True)
         print("Load: ", model_path)
-    except:
+    else:
         ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
 
         filtered_state_dict = {}
@@ -192,13 +212,21 @@ SUNFLOWERS_RAIN_IMAGE = "assets/rain_sunflower_plains.png"
 device = torch.device('cuda')
 
 def save_video(frames, path="output.mp4", fps=10):
+    temp_path = path[:-4] + "_temp.mp4"
     h, w, _ = frames[0].shape
-    out = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+
+    out = cv2.VideoWriter(temp_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
     for frame in frames:
         out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
     out.release()
 
-    return path
+    ffmpeg_cmd = [
+        "ffmpeg", "-y", "-i", temp_path,
+        "-c:v", "libx264", "-crf", "23", "-preset", "medium",
+        path
+    ]
+    subprocess.run(ffmpeg_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    os.remove(temp_path)
 
 cfg = OmegaConf.load("configurations/huggingface.yaml")
 worldmem = WorldMemMinecraft(cfg)

configurations/algorithm/base_algo.yaml

@@ -0,0 +1,3 @@
+# This will be passed as the cfg to Algo.__init__(cfg) of your algorithm class
+
+debug: ${debug} # inherited from configurations/config.yaml

configurations/algorithm/base_pytorch_algo.yaml

@@ -0,0 +1,4 @@
+defaults:
+  - base_algo # inherits from configurations/algorithm/base_algo.yaml
+
+lr: ${experiment.training.lr}

configurations/algorithm/df_base.yaml

@@ -0,0 +1,42 @@
+defaults:
+  - base_pytorch_algo
+
+# dataset-dependent configurations
+x_shape: ${dataset.observation_shape}
+frame_stack: 1
+frame_skip: 1
+data_mean: ${dataset.data_mean}
+data_std: ${dataset.data_std}
+external_cond_dim: 0 #${dataset.action_dim}
+context_frames: ${dataset.context_length}
+# training hyperparameters
+weight_decay: 1e-4
+warmup_steps: 10000
+optimizer_beta: [0.9, 0.999]
+# diffusion-related
+uncertainty_scale: 1
+guidance_scale: 0.0
+chunk_size: 1 # -1 for full trajectory diffusion, number to specify diffusion chunk size
+scheduling_matrix: autoregressive
+noise_level: random_all
+causal: True
+
+diffusion:
+  # training
+  objective: pred_x0
+  beta_schedule: cosine
+  schedule_fn_kwargs: {}
+  clip_noise: 20.0
+  use_snr: False
+  use_cum_snr: False
+  use_fused_snr: False
+  snr_clip: 5.0
+  cum_snr_decay: 0.98
+  timesteps: 1000
+  # sampling
+  sampling_timesteps: 50 # fixme, numer of diffusion steps, should be increased
+  ddim_sampling_eta: 1.0
+  stabilization_level: 10
+  # architecture
+  architecture:
+    network_size: 64

configurations/algorithm/df_video_worldmemminecraft.yaml

@@ -0,0 +1,43 @@
+defaults:
+  - df_base
+
+n_frames: ${dataset.n_frames}
+frame_skip: ${dataset.frame_skip}
+metadata: ${dataset.metadata}
+
+# training hyperparameters
+weight_decay: 2e-3
+warmup_steps: 1000
+optimizer_beta: [0.9, 0.99]
+action_cond_dim: 25
+use_plucker: true
+
+diffusion:
+  # training
+  beta_schedule: sigmoid
+  objective: pred_v
+  use_fused_snr: True
+  cum_snr_decay: 0.96
+  clip_noise: 20.
+  # sampling
+  sampling_timesteps: 20
+  ddim_sampling_eta: 0.0
+  stabilization_level: 15
+  # architecture
+  architecture:
+    network_size: 64
+    attn_heads: 4
+    attn_dim_head: 64
+    dim_mults: [1, 2, 4, 8]
+    resolution: ${dataset.resolution}
+    attn_resolutions: [16, 32, 64, 128]
+    use_init_temporal_attn: True
+    use_linear_attn: True
+    time_emb_type: rotary
+
+metrics:
+  # - fvd
+  # - fid
+  # - lpips
+
+_name: df_video_worldmemminecraft

configurations/dataset/base_dataset.yaml

@@ -0,0 +1,3 @@
+# This will be passed as the cfg to Dataset.__init__(cfg) of your dataset class
+
+debug: ${debug} # inherited from configurations/config.yaml

configurations/dataset/base_video.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - base_dataset
+
+metadata: "data/${dataset.name}/metadata.json"
+data_mean: "data/${dataset.name}/data_mean.npy"
+data_std: "data/${dataset.name}/data_std.npy"
+save_dir: ???
+n_frames: 32
+context_length: 4
+resolution: 128
+observation_shape: [3, "${dataset.resolution}", "${dataset.resolution}"]
+external_cond_dim: 0
+validation_multiplier: 1
+frame_skip: 1

configurations/dataset/video_minecraft.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - base_video
+
+save_dir: data/minecraft_simple_backforward
+n_frames: 16 # TODO: increase later
+resolution: 128
+data_mean: 0.5
+data_std: 0.5
+action_cond_dim: 25
+context_length: 1
+frame_skip: 1
+validation_multiplier: 1
+
+_name: video_minecraft_oasis

configurations/experiment/base_experiment.yaml

@@ -0,0 +1,2 @@
+debug: ${debug} # inherited from configurations/config.yaml
+tasks: [main] # tasks to run sequantially, such as [training, test], useful when your project has multiple stages and you want to run only a subset of them.

configurations/experiment/base_pytorch.yaml

@@ -0,0 +1,50 @@
+# inherites from base_experiment.yaml
+# most of the options have docs at https://lightning.ai/docs/pytorch/stable/common/trainer.html
+
+defaults:
+  - base_experiment
+
+tasks: [training] # tasks to run sequantially, change when your project has multiple stages and you want to run only a subset of them.
+num_nodes: 1 # number of gpu servers used in large scale distributed training
+
+training:
+  precision: 16-mixed # set float precision, 16-mixed is faster while 32 is more stable
+  compile: False # whether to compile the model with torch.compile
+  lr: 0.001 # learning rate
+  batch_size: 16 # training batch size; effective batch size is this number * gpu * nodes iff using distributed training
+  max_epochs: 1000 # set to -1 to train forever
+  max_steps: -1 # set to -1 to train forever, will override max_epochs
+  max_time: null # set to something like "00:12:00:00" to enable
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing.
+    shuffle: True # whether training data will be shuffled
+  optim:
+    accumulate_grad_batches: 1 # accumulate gradients for n batches before backprop
+    gradient_clip_val: 0 # clip gradients with norm above this value, set to 0 to disable
+  checkpointing:
+    # these are arguments to pytorch lightning's callback, `ModelCheckpoint` class
+    every_n_train_steps: 5000 # save a checkpoint every n train steps
+    every_n_epochs: null # mutually exclusive with ``every_n_train_steps`` and ``train_time_interval``
+    train_time_interval: null # in format of "00:12:00:00", mutually exclusive with ``every_n_train_steps`` and ``every_n_epochs``.
+    enable_version_counter: False # If this is ``False``, later checkpoint will be overwrite previous ones.
+
+validation:
+  precision: 16-mixed
+  compile: False # whether to compile the model with torch.compile
+  batch_size: 16 # validation batch size per GPU; effective batch size is this number * gpu * nodes iff using distributed training
+  val_every_n_step: 2000 # controls how frequent do we run validation, can be float (fraction of epoches) or int (steps) or null (if val_every_n_epoch is set)
+  val_every_n_epoch: null # if you want to do validation every n epoches, requires val_every_n_step to be null.
+  limit_batch: null # if null, run through validation set. Otherwise limit the number of batches to use for validation.
+  inference_mode: True # whether to run validation in inference mode (enable_grad won't work!)
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing, for validation.
+    shuffle: False # whether validation data will be shuffled
+
+test:
+  precision: 16-mixed
+  compile: False # whether to compile the model with torch.compile
+  batch_size: 4 # test batch size per GPU; effective batch size is this number * gpu * nodes iff using distributed training
+  limit_batch: null # if null, run through test set. Otherwise limit the number of batches to use for test.
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing, for test.
+    shuffle: False # whether test data will be shuffled

configurations/experiment/exp_video.yaml

@@ -0,0 +1,31 @@
+defaults:
+  - base_pytorch
+
+tasks: [training]
+
+training:
+  lr: 2e-5
+  precision: 16-mixed
+  batch_size: 4
+  max_epochs: -1
+  max_steps: 2000005
+  checkpointing:
+    every_n_train_steps: 2500
+  optim:
+    gradient_clip_val: 1.0
+
+validation:
+  val_every_n_step: 2500
+  val_every_n_epoch: null
+  batch_size: 4
+  limit_batch: 1
+
+test:
+  limit_batch: 1
+  batch_size: 1
+
+logging:
+  metrics:
+    # - fvd
+    # - fid
+    # - lpips

configurations/huggingface.yaml

@@ -54,7 +54,7 @@ noise_level: random_all
 causal: True
 x_shape: [3, 360, 640]
 context_frames: 1
-diffusion_path: /mnt/xiaozeqi/worldmem_github/diffusion_only_2025-04-06-16-24-11_epoch1step720000.ckpt
+diffusion_path: yslan/worldmem_checkpoints/diffusion_only.ckpt
 vae_path: yslan/worldmem_checkpoints/vae_only.ckpt
 pose_predictor_path: yslan/worldmem_checkpoints/pose_prediction_model_only.ckpt
 next_frame_length: 1

configurations/training.yaml

@@ -0,0 +1,16 @@
+# configuration parsing starts here
+defaults:
+  - experiment: exp_video # experiment yaml file name in configurations/experiments folder [fixme]
+  - dataset: video_minecraft # dataset yaml file name in configurations/dataset folder [fixme]
+  - algorithm: df_video_worldmemminecraft # algorithm yaml file name in configurations/algorithm folder [fixme]
+  - cluster: null # optional, cluster yaml file name in configurations/cluster folder. Leave null for local compute
+
+debug: false # global debug flag will be passed into configuration of experiment, dataset and algorithm
+
+wandb:
+  entity: xizaoqu # wandb account name / organization name [fixme]
+  project: worldmem # wandb project name; if not provided, defaults to root folder name [fixme]
+  mode: online # set wandb logging to online, offline or dryrun
+
+resume: null # wandb run id to resume logging and loading checkpoint from
+load: null # wandb run id containing checkpoint or a path to a checkpoint file

datasets/README.md

@@ -0,0 +1,11 @@
+The `datasets` folder is used to contain dataset code or environment code.
+Don't store actual data like images here! For those, please use the `data` folder instead of `datasets`.
+
+Create a folder to create your own pytorch dataset definition. Then, update the `__init__.py`
+at every level to register all datasets.
+
+Each dataset class takes in a DictConfig file `cfg` in its `__init__`, which allows you to pass in arguments via configuration file in `configurations/dataset` or [command line override](https://hydra.cc/docs/tutorials/basic/your_first_app/simple_cli/).
+
+---
+
+This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research template [repo](https://github.com/buoyancy99/research-template). By its MIT license, you must keep the above sentence in `README.md` and the `LICENSE` file to credit the author.

datasets/__init__.py

@@ -0,0 +1 @@
+from .video import MinecraftVideoDataset

datasets/video/__init__.py

@@ -0,0 +1 @@
+from .minecraft_video_dataset import MinecraftVideoDataset

datasets/video/base_video_dataset.py

@@ -0,0 +1,153 @@
+from typing import Sequence
+import torch
+import random
+import os
+import numpy as np
+import cv2
+from omegaconf import DictConfig
+from torchvision import transforms
+from pathlib import Path
+from abc import abstractmethod, ABC
+import json
+
+
+class BaseVideoDataset(torch.utils.data.Dataset, ABC):
+    """
+    Base class for video datasets. Videos may be of variable length.
+
+    Folder structure of each dataset:
+    - [save_dir] (specified in config, e.g., data/phys101)
+        - /[split] (one per split)
+            - /data_folder_name (e.g., videos)
+            metadata.json
+    """
+
+    def __init__(self, cfg: DictConfig, split: str = "training"):
+        super().__init__()
+        self.cfg = cfg
+        self.split = split
+        self.resolution = cfg.resolution
+        self.external_cond_dim = cfg.external_cond_dim
+        self.n_frames = (
+            cfg.n_frames * cfg.frame_skip
+            if split == "training"
+            else cfg.n_frames * cfg.frame_skip * cfg.validation_multiplier
+        )
+        self.frame_skip = cfg.frame_skip
+        self.save_dir = Path(cfg.save_dir)
+        self.save_dir.mkdir(exist_ok=True, parents=True)
+        self.split_dir = self.save_dir / f"{split}"
+        self.data_paths = self.get_data_paths(self.split)
+
+        if self.split == 'training':
+            self.metadata = [1300] * len(self.data_paths) # total 1500 f, randomly sampled ~1300 clips
+        else:
+            self.metadata = [1] * len(self.data_paths) # total 1500 f
+
+        self.clips_per_video = np.clip(np.array(self.metadata) - self.n_frames + 1, a_min=1, a_max=None).astype(
+            np.int32
+        )
+        self.cum_clips_per_video = np.cumsum(self.clips_per_video)
+        self.transform = transforms.Resize((self.resolution, self.resolution), antialias=True)
+
+        # shuffle but keep the same order for each epoch, so validation sample is diverse yet deterministic
+        random.seed(0)
+        self.idx_remap = list(range(self.__len__()))
+        random.shuffle(self.idx_remap)
+        
+    @abstractmethod
+    def download_dataset(self) -> Sequence[int]:
+        """
+        Download dataset from the internet and build it in save_dir
+
+        Returns a list of video lengths
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_data_paths(self, split):
+        """Return a list of data paths (e.g. xxx.mp4) for a given split"""
+        raise NotImplementedError
+
+    def get_data_lengths(self, split):
+        """Return a list of num_frames for each data path (e.g. xxx.mp4) for a given split"""
+        lengths = []
+        for path in self.get_data_paths(split):
+            length = cv2.VideoCapture(str(path)).get(cv2.CAP_PROP_FRAME_COUNT)
+            lengths.append(length)
+        return lengths
+
+    def split_idx(self, idx):
+        video_idx = np.argmax(self.cum_clips_per_video > idx)
+        frame_idx = idx - np.pad(self.cum_clips_per_video, (1, 0))[video_idx]
+        return video_idx, frame_idx
+
+    @staticmethod
+    def load_video(path: Path):
+        """
+        Load video from a path
+        :param filename: path to the video
+        :return: video as a numpy array
+        """
+
+        cap = cv2.VideoCapture(str(path))
+
+        frames = []
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if ret:
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                frames.append(frame)
+            else:
+                break
+
+        cap.release()
+        frames = np.stack(frames, dtype=np.uint8)
+        return np.transpose(frames, (0, 3, 1, 2))  # (T, C, H, W)
+
+    @staticmethod
+    def load_image(filename: Path):
+        """
+        Load image from a path
+        :param filename: path to the image
+        :return: image as a numpy array
+        """
+        image = cv2.imread(str(filename))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        return np.transpose(image, (2, 0, 1))
+
+    def __len__(self):
+        return self.clips_per_video.sum()
+
+    def __getitem__(self, idx):
+        idx = self.idx_remap[idx]
+        video_idx, frame_idx = self.split_idx(idx)
+        video_path = self.data_paths[video_idx]
+        video = self.load_video(video_path)[frame_idx : frame_idx + self.n_frames]
+
+        pad_len = self.n_frames - len(video)
+
+        nonterminal = np.ones(self.n_frames)
+        if len(video) < self.n_frames:
+            video = np.pad(video, ((0, pad_len), (0, 0), (0, 0), (0, 0)))
+            nonterminal[-pad_len:] = 0
+
+        video = torch.from_numpy(video / 256.0).float()
+        video = self.transform(video)
+
+        if self.external_cond_dim:
+            external_cond = np.load(
+                # pylint: disable=no-member
+                self.condition_dir
+                / f"{video_path.name.replace('.mp4', '.npy')}"
+            )
+            if len(external_cond) < self.n_frames:
+                external_cond = np.pad(external_cond, ((0, pad_len),))
+            external_cond = torch.from_numpy(external_cond).float()
+            return (
+                video[:: self.frame_skip],
+                external_cond[:: self.frame_skip],
+                nonterminal[:: self.frame_skip],
+            )
+        else:
+            return video[:: self.frame_skip], nonterminal[:: self.frame_skip]

datasets/video/minecraft_video_dataset.py

@@ -0,0 +1,250 @@
+import os
+import io
+import tarfile
+import numpy as np
+import torch
+from typing import Sequence, Mapping
+from omegaconf import DictConfig
+from pytorchvideo.data.encoded_video import EncodedVideo
+import random
+
+from .base_video_dataset import BaseVideoDataset
+
+
+
+
+ACTION_KEYS = [
+    "inventory",
+    "ESC",
+    "hotbar.1",
+    "hotbar.2",
+    "hotbar.3",
+    "hotbar.4",
+    "hotbar.5",
+    "hotbar.6",
+    "hotbar.7",
+    "hotbar.8",
+    "hotbar.9",
+    "forward",
+    "back",
+    "left",
+    "right",
+    "cameraY",
+    "cameraX",
+    "jump",
+    "sneak",
+    "sprint",
+    "swapHands",
+    "attack",
+    "use",
+    "pickItem",
+    "drop",
+]
+
+def convert_action_space(actions):
+    vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
+    vec_25[actions[:,0]==1, 11] = 1
+    vec_25[actions[:,0]==2, 12] = 1
+    vec_25[actions[:,4]==11, 16] = -1
+    vec_25[actions[:,4]==13, 16] = 1
+    vec_25[actions[:,3]==11, 15] = -1
+    vec_25[actions[:,3]==13, 15] = 1
+    vec_25[actions[:,5]==6, 24] = 1
+    vec_25[actions[:,5]==1, 24] = 1
+    vec_25[actions[:,1]==1, 13] = 1
+    vec_25[actions[:,1]==2, 14] = 1
+    vec_25[actions[:,7]==1, 2] = 1
+    return vec_25
+
+# Dataset class
+class MinecraftVideoDataset(BaseVideoDataset):
+    """
+    Minecraft video dataset for training and validation.
+
+    Args:
+        cfg (DictConfig): Configuration object.
+        split (str): Dataset split ("training" or "validation").
+    """
+    def __init__(self, cfg: DictConfig, split: str = "training"):
+        if split == "test":
+            split = "validation"
+        self.wo_updown = getattr(cfg, "wo_updown", False)
+        super().__init__(cfg, split)
+        self.n_frames = cfg.n_frames_valid if split == "validation" and hasattr(cfg, "n_frames_valid") else cfg.n_frames
+        self.memory_condition_length = cfg.memory_condition_length
+        self.customized_validation = cfg.customized_validation
+        if split == "training":
+            self.angle_range = cfg.angle_range
+            self.pos_range = cfg.pos_range
+        self.add_timestamp_embedding = cfg.add_timestamp_embedding
+        self.training_dropout = 0.1
+        self.memory_condition_length = getattr(cfg, "memory_condition_length", False)
+        self.sample_more_event = getattr(cfg, "sample_more_event", False)
+        self.causal_frame = getattr(cfg, "causal_frame", False)
+
+    def get_data_paths(self, split: str):
+        """
+        Retrieve all video file paths for the given split.
+
+        Args:
+            split (str): Dataset split ("training" or "validation").
+
+        Returns:
+            List[Path]: List of video file paths.
+        """
+        data_dir = self.save_dir / split
+        paths = sorted(list(data_dir.glob("**/*.mp4")), key=lambda x: x.name)
+
+        if self.wo_updown:
+            # Filter out paths containing "w_updown"
+            paths = [p for p in paths if "w_updown" not in str(p)]
+        
+        if split == "validation" and self.wo_updown:
+            paths = [p for p in paths if "w_updown" not in str(p)]
+        elif split == "validation":
+            paths = [p for p in paths if "w_updown" in str(p)]
+
+        if not paths:
+            sub_dirs = os.listdir(data_dir)
+            for sub_dir in sub_dirs:
+                sub_path = data_dir / sub_dir
+                paths += sorted(list(sub_path.glob("**/*.mp4")), key=lambda x: x.name)
+        return paths
+
+    def download_dataset(self):
+        pass
+    
+    def __getitem__(self, idx: int):
+        """
+        Retrieve a single data sample by index.
+
+        Args:
+            idx (int): Index of the data sample.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, np.ndarray, np.ndarray]: Video, actions, poses, and timestamps.
+        """
+        max_retries = 1000
+        for _ in range(max_retries):
+            try:
+                return self.load_data(idx)
+            except Exception as e:
+                print(f"Retrying due to error: {e}")
+                idx = (idx + 1) % len(self)
+
+    def load_data(self, idx):
+        # === 1. Remap index and skip first few frames ===
+        idx = self.idx_remap[idx]
+        file_idx, frame_idx = self.split_idx(idx)
+        frame_idx += 100 # initial few frames are low quality
+
+        # === 2. Load paths and data arrays ===
+        video_path = self.data_paths[file_idx]
+        action_path = video_path.with_suffix(".npz")
+        data = np.load(action_path)
+        actions_pool = convert_action_space(data["actions"])
+        poses_pool = data["poses"]
+
+        # Fix corrupted height (maybe) in the first frame
+        poses_pool[0, 1] = poses_pool[1, 1]
+        assert poses_pool[:, 1].ptp() < 2, f"Height variation too large: {poses_pool[:, 1].ptp()} - {video_path}"
+
+        # Pad poses if shorter than actions
+        if len(poses_pool) < len(actions_pool):
+            poses_pool = np.pad(poses_pool, ((1, 0), (0, 0)))
+
+        # === 3. Load video clip ===
+        video_raw = EncodedVideo.from_path(video_path, decode_audio=False)
+        fps = 10
+        clip = video_raw.get_clip(
+            start_sec=frame_idx / fps,
+            end_sec=(frame_idx + self.n_frames) / fps
+        )["video"]
+        video = clip.permute(1, 2, 3, 0).numpy()
+
+        actions = np.copy(actions_pool[frame_idx : frame_idx + self.n_frames])
+        poses = np.copy(poses_pool[frame_idx : frame_idx + self.n_frames])
+
+        # === 4. Normalize poses relative to current segment ===
+        def normalize_pose(pose, ref_pose):
+            pose[:, :3] -= ref_pose[:1, :3]
+            pose[:, -1] = -pose[:, -1]
+            pose[:, 3:] %= 360
+            return pose
+
+        poses_pool = normalize_pose(poses_pool, poses)
+        poses = normalize_pose(poses, poses)
+
+        assert len(video) >= self.n_frames, f"{video_path}"
+
+        # === 5. Sample memory frames for training ===
+        if self.split == "training" and self.memory_condition_length > 0:
+            use_memory = random.uniform(0, 1) > self.training_dropout
+
+            if use_memory:
+                # Compute pose distance between current and candidate frames
+                dis = np.abs(poses[:, None] - poses_pool[None, :])
+                dis[..., 3:][dis[..., 3:] > 180] = 360 - dis[..., 3:][dis[..., 3:] > 180]
+
+                spatial_match = (dis[..., :3] <= self.pos_range).sum(-1) >= 3
+                angular_match = (dis[..., 3:] <= self.angle_range).sum(-1) >= 2
+                not_exact_match = ((dis[..., :3] > 0).sum(-1) >= 1) | ((dis[..., 3:] > 0).sum(-1) >= 1)
+
+                valid_index = (spatial_match & angular_match & not_exact_match).sum(0)
+                valid_index[:100] = 0  # skip unstable early frames
+
+                # Exclude future if causality and timestamp are enabled
+                if self.add_timestamp_embedding and self.causal_frame and (actions_pool[:frame_idx, 24] == 1).sum() > 0:
+                    valid_index[frame_idx:] = 0
+
+                # Select indices satisfying condition
+                mask = valid_index >= 1
+                mask[0] = False
+                candidate_indices = np.argwhere(mask)
+
+                # Backup candidates with weaker conditions
+                mask2 = valid_index >= 0
+                mask2[0] = False
+
+                count = min(self.memory_condition_length, candidate_indices.shape[0])
+                selected = candidate_indices[np.random.choice(candidate_indices.shape[0], count, replace=True)][:, 0]
+
+                if count < self.memory_condition_length:
+                    extra = np.argwhere(mask2)
+                    extra = extra[np.random.choice(extra.shape[0], self.memory_condition_length - count, replace=True)][:, 0]
+                    selected = np.concatenate([selected, extra])
+
+                # Prioritize event-trigger frames if applicable
+                if self.sample_more_event and random.uniform(0, 1) > 0.3:
+                    event_idx = torch.nonzero(actions_pool[:frame_idx, 24] == 1)[:, 0]
+                    if len(event_idx) > self.memory_condition_length // 2:
+                        event_idx = event_idx[-self.memory_condition_length // 2:]
+                    if len(event_idx) > 0:
+                        selected[-len(event_idx):] = event_idx + 4
+
+            else:
+                selected = np.full(self.memory_condition_length, random.randint(0, frame_idx))
+
+            # === 6. Retrieve video frames for selected memory indices ===
+            video_pool = []
+            for si in selected:
+                frame = video_raw.get_clip(start_sec=si / fps, end_sec=(si + 1) / fps)["video"][:, 0].permute(1, 2, 0)
+                video_pool.append(frame)
+
+            video = np.concatenate([video, np.stack(video_pool)], axis=0)
+            actions = np.concatenate([actions, actions_pool[selected]], axis=0)
+            poses = np.concatenate([poses, poses_pool[selected]], axis=0)
+            timestamp = np.concatenate([np.arange(frame_idx, frame_idx + self.n_frames), selected])
+        else:
+            timestamp = np.arange(self.n_frames)
+
+        # === 7. Convert video to torch format ===
+        video = torch.from_numpy(video / 255.0).float().permute(0, 3, 1, 2).contiguous()
+
+        # === 9. Return all items ===
+        return (
+            video[:: self.frame_skip],
+            actions[:: self.frame_skip],
+            poses[:: self.frame_skip],
+            timestamp
+        )

infer.sh

@@ -0,0 +1,22 @@
+wandb enabled
+python -m main +name=infer \
+    experiment.tasks=[validation] \
+    dataset.validation_multiplier=1 \
+    +diffusion_model_path=yslan/worldmem_checkpoints/diffusion_only.ckpt \
+    +vae_path=yslan/worldmem_checkpoints/vae_only.ckpt \
+    +load=/mnt/xiaozeqi/checkpoints/2025-03-31-22-24-06-epoch1step240000_2d_largerange.ckpt \
+    +customized_load=true \
+    +seperate_load=false \
+    dataset.n_frames=8 \
+    dataset.save_dir=data/minecraft \
+    +dataset.n_frames_valid=700 \
+    +dataset.memory_condition_length=8 \
+    +dataset.customized_validation=true \
+    +dataset.add_timestamp_embedding=true \
+    +algorithm.n_tokens=8 \
+    +algorithm.memory_condition_length=8 \
+    algorithm.context_frames=600 \
+    +algorithm.relative_embedding=true \
+    +algorithm.log_video=true \
+    +algorithm.add_timestamp_embedding=true \
+    algorithm.metrics=[lpips,psnr] \

train_stage_1.sh

@@ -0,0 +1,29 @@
+wandb enabled
+
+# set -e
+python -m main +name=train \
+    +diffusion_model_path=your_diffusion_model_path \
+    +vae_path=your_vae_path \
+    +customized_load=true \
+    +seperate_load=true \
+    +zero_init_gate=true \
+    dataset.n_frames=8 \
+    dataset.save_dir=data/minecraft \
+    +dataset.n_frames_valid=700 \
+    +dataset.angle_range=110 \
+    +dataset.pos_range=2 \
+    +dataset.memory_condition_length=8 \
+    +dataset.customized_validation=true \
+    +dataset.add_timestamp_embedding=true \
+    +dataset.wo_updown=true \
+    +algorithm.n_tokens=8 \
+    +algorithm.memory_condition_length=8 \
+    algorithm.context_frames=600 \
+    +algorithm.relative_embedding=true \
+    +algorithm.log_video=true \
+    +algorithm.add_timestamp_embedding=true \
+    algorithm.metrics=[lpips,psnr] \
+    experiment.training.checkpointing.every_n_train_steps=2500 \
+    experiment.training.max_steps=120000
+
+

train_stage_2.sh

@@ -0,0 +1,25 @@
+wandb enabled
+
+# set -e
+python -m main +name=train \
+    dataset.n_frames=8 \
+    dataset.save_dir=data/minecraft \
+    +dataset.n_frames_valid=700 \
+    +dataset.angle_range=110 \
+    +dataset.pos_range=8 \
+    +dataset.memory_condition_length=8 \
+    +dataset.customized_validation=true \
+    +dataset.add_timestamp_embedding=true \
+    +dataset.wo_updown=true \
+    +algorithm.n_tokens=8 \
+    +algorithm.memory_condition_length=8 \
+    algorithm.context_frames=600 \
+    +algorithm.relative_embedding=true \
+    +algorithm.log_video=true \
+    +algorithm.add_timestamp_embedding=true \
+    algorithm.metrics=[lpips,psnr] \
+    experiment.training.checkpointing.every_n_train_steps=2500 \
+    resume=your_wandb_job_id e.g.yhht29bz \
+    +output_dir=your_saving_path e.g. outputs/2025-05-18/15-16-32 \
+    experiment.training.max_steps=240000
+

train_stage_3.sh

@@ -0,0 +1,24 @@
+wandb enabled
+
+# set -e
+python -m main +name=train \
+    dataset.n_frames=8 \
+    dataset.save_dir=data/minecraft \
+    +dataset.n_frames_valid=700 \
+    +dataset.angle_range=110 \
+    +dataset.pos_range=8 \
+    +dataset.memory_condition_length=8 \
+    +dataset.customized_validation=true \
+    +dataset.add_timestamp_embedding=true \
+    +dataset.wo_updown=false \
+    +algorithm.n_tokens=8 \
+    +algorithm.memory_condition_length=8 \
+    algorithm.context_frames=600 \
+    +algorithm.relative_embedding=true \
+    +algorithm.log_video=true \
+    +algorithm.add_timestamp_embedding=true \
+    algorithm.metrics=[lpips,psnr] \
+    experiment.training.checkpointing.every_n_train_steps=2500 \
+    resume=your_wandb_job_id e.g.yhht29bz \
+    +output_dir=your_saving_path e.g. outputs/2025-05-18/15-16-32 \
+    experiment.training.max_steps=700000