update evaluation and fix psnr

README.md

@@ -49,7 +49,7 @@ conda install -c conda-forge ffmpeg=4.3.2
 python app.py
 ```
 
-## Training and Inference
+## Run
 
 To enable cloud logging with [Weights & Biases (wandb)](https://wandb.ai/site), follow these steps:
 
@@ -92,8 +92,8 @@ 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 \
++diffusion_model_path=zeqixiao/worldmem_checkpoints/diffusion_only.ckpt \
++vae_path=zeqixiao/worldmem_checkpoints/vae_only.ckpt \
 +customized_load=true \
 +seperate_load=true \
 ```
@@ -106,6 +106,41 @@ Or **load a combined checkpoint**:
 +seperate_load=false \
 ```
 
+### Evaluation
+
+To run evaluation:
+
+```bash
+sh evaluate.sh
+```
+
+This script reproduces the results in Table 1 (beyond context window). Evaluating 1 case on 1 A100 GPU takes approximately 6 minutes. You can adjust `experiment.test.limit_batch` to specify the number of cases to evaluate.
+
+Visual results will be saved by default to a timestamped directory (e.g., `outputs/2025-11-30/00-02-42`).
+
+To calculate the FID score, run:
+
+```bash
+python calculate_fid.py --videos_dir <path_to_videos>
+```
+
+For example:
+
+```bash
+python calculate_fid.py --videos_dir outputs/2025-11-30/00-02-42/videos/test_vis
+```
+
+**Expected Results:**
+
+| Metric | Value  |
+|--------|--------|
+| PSNR   | 19.34  |
+| LPIPS  | 0.1667 |
+| FID    | 15.13  |
+
+*Note: FID is computed over 5000 frames.*
+*Note: Previous versions incorrectly used `data_range=2.0` for PSNR calculation, but the decoded video data is in the range [0, 1], so `data_range=1.0` should be used. This bug inflated PSNR values by approximately 6 dB. We have now corrected this in the latest version by clipping predictions to [0, 1] before metric computation. The relative performance comparisons and conclusions remain unchanged.*
+
 ---
 
 ## Dataset
@@ -119,14 +154,29 @@ data/
 └── minecraft/
     ├── training/
     └── validation/
+    └── test/
 ```
 
+## Data Generation
+
+After setting up the environment as described in [MineDojo's GitHub repository](https://github.com/MineDojo/MineDojo), you can generate data using the following command:
+
+```bash
+xvfb-run -a python data_generator.py -o data/test -z 4 --env_type plains
+```
+
+**Parameters:**
+- `-o`: Output directory for generated data
+- `-z`: Number of parallel workers
+- `--env_type`: Environment type (e.g., `plains`, `forest`, `desert`)
+
 
 ## TODO
 
 - [x] Release inference models and weights;
 - [x] Release training pipeline on Minecraft;
 - [x] Release training data on Minecraft;
+- [x] Release evaluation scripts and data generator.
 
 
 

algorithms/worldmem/df_base.py

@@ -194,7 +194,7 @@ class DiffusionForcingBase(BasePytorchAlgo):
     def test_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
         return self.validation_step(*args, **kwargs, namespace="test")
 
-    def test_epoch_end(self) -> None:
+    def on_test_epoch_end(self) -> None:
         self.on_validation_epoch_end(namespace="test")
 
     def _generate_noise_levels(self, xs: torch.Tensor, masks: Optional[torch.Tensor] = None) -> torch.Tensor:

algorithms/worldmem/df_video.py

@@ -341,16 +341,18 @@ class WorldMemMinecraft(DiffusionForcingBase):
         self.memory_condition_length = cfg.memory_condition_length
         self.pose_cond_dim = getattr(cfg, "pose_cond_dim", 5)
 
-        self.use_plucker = cfg.use_plucker
-        self.relative_embedding = cfg.relative_embedding
+        self.use_plucker = getattr(cfg, "use_plucker", True)
+        self.relative_embedding = getattr(cfg, "relative_embedding", True)
         self.state_embed_only_on_qk = getattr(cfg, "state_embed_only_on_qk", True)
         self.use_memory_attention = getattr(cfg, "use_memory_attention", True)
-        self.add_timestamp_embedding = cfg.add_timestamp_embedding
+        self.add_timestamp_embedding = getattr(cfg, "add_timestamp_embedding", True)
         self.ref_mode = getattr(cfg, "ref_mode", 'sequential')
         self.log_curve = getattr(cfg, "log_curve", False)
         self.focal_length =  getattr(cfg, "focal_length", 0.35)
         self.log_video = cfg.log_video
-        self.self_consistency_eval = getattr(cfg, "self_consistency_eval", False)
+        self.save_local = getattr(cfg, "save_local", True)
+        self.local_save_dir = getattr(cfg, "local_save_dir", None)
+        self.lpips_batch_size = getattr(cfg, "lpips_batch_size", 16)
         self.next_frame_length = getattr(cfg, "next_frame_length", 1)
         self.require_pose_prediction = getattr(cfg, "require_pose_prediction", False)
 
@@ -497,12 +499,20 @@ class WorldMemMinecraft(DiffusionForcingBase):
                 namespace=namespace + "_vis",
                 context_frames=self.context_frames,
                 logger=self.logger.experiment,
+                save_local=self.save_local,
+                local_save_dir=self.local_save_dir,
             )
 
         if xs is not None:
+            # Move data to the same device as LPIPS model for metric calculation
+            device = next(self.validation_lpips_model.parameters()).device
+            xs_pred_device = xs_pred.to(device)
+            xs_device = xs.to(device)
+            
             metric_dict = get_validation_metrics_for_videos(
-                xs_pred, xs, 
-                lpips_model=self.validation_lpips_model)
+                xs_pred_device, xs_device, 
+                lpips_model=self.validation_lpips_model,
+                lpips_batch_size=self.lpips_batch_size)
             
             self.log_dict(
                 {"mse": metric_dict['mse'],
@@ -524,19 +534,6 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
                 self.logger.experiment.log({"frame_wise_psnr_plot": line_plot})
       
-        elif self.self_consistency_eval:
-            metric_dict = get_validation_metrics_for_videos(
-                xs_pred[:1],
-                xs_pred[-1:],
-                lpips_model=self.validation_lpips_model,
-            )            
-            self.log_dict(
-                {"lpips": metric_dict['lpips'],
-                "mse": metric_dict['mse'],
-                "psnr": metric_dict['psnr']},
-                sync_dist=True
-            )
-
         self.validation_step_outputs.clear()
 
     def _preprocess_batch(self, batch):
@@ -786,8 +783,8 @@ class WorldMemMinecraft(DiffusionForcingBase):
         xs_pred = self.decode(xs_pred[n_context_frames:].to(conditions.device))
         xs_decode = self.decode(xs[n_context_frames:].to(conditions.device))
 
-        # Store results for evaluation
-        self.validation_step_outputs.append((xs_pred, xs_decode))
+        # Store results for evaluation (move to CPU to save GPU memory)
+        self.validation_step_outputs.append((xs_pred.detach().cpu(), xs_decode.detach().cpu()))
         return
 
     @torch.no_grad()

app.py

@@ -25,54 +25,10 @@ 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
-    
+from experiments.exp_base import load_custom_checkpoint
 
 torch.set_float32_matmul_precision("high")
 
-def load_custom_checkpoint(algo, checkpoint_path):
-    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:])
-        model_path = hf_hub_download(repo_id=repo_id, 
-                            filename=file_name)
-        ckpt = torch.load(model_path, map_location=torch.device('cpu'))
-
-        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)
-    else:
-        ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
-
-        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)
-
-        # algo.load_state_dict(ckpt['state_dict'], strict=True)      
-        print("Load: ", checkpoint_path)  
-
 def download_assets_if_needed():
     ASSETS_URL_BASE = "https://huggingface.co/spaces/yslan/worldmem/resolve/main/assets/examples"
     ASSETS_DIR = "assets/examples"

compute_video_psnr.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""
+Compute PSNR for each video pair in gt and pred directories.
+"""
+import os
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+import pandas as pd
+import imageio
+
+
+def load_video(video_path):
+    """
+    Load a video file and return frames as a numpy array.
+    
+    Args:
+        video_path: Path to the video file
+        
+    Returns:
+        np.ndarray: Video frames of shape (T, C, H, W) normalized to [0, 1]
+    """
+    reader = imageio.get_reader(str(video_path))
+    frames = []
+    for frame in reader:
+        # frame is already RGB
+        frame = frame.astype(np.float32) / 255.0
+        frames.append(frame)
+    reader.close()
+    
+    if len(frames) == 0:
+        raise ValueError(f"No frames loaded from {video_path}")
+    
+    # Convert to numpy array: (T, H, W, C) -> (T, C, H, W)
+    frames = np.stack(frames, axis=0)
+    frames = np.transpose(frames, (0, 3, 1, 2))
+    
+    return frames
+
+
+def calculate_psnr(pred, gt, data_range=1.0):
+    """
+    Calculate PSNR between two images/videos.
+    
+    Args:
+        pred: Predicted frames
+        gt: Ground truth frames
+        data_range: Data range (default: 1.0)
+        
+    Returns:
+        float: PSNR value in dB
+    """
+    mse = np.mean((pred - gt) ** 2)
+    if mse == 0:
+        return float('inf')
+    return 10 * np.log10((data_range ** 2) / mse)
+
+
+def compute_psnr(pred_frames, gt_frames, data_range=1.0):
+    """
+    Compute PSNR between predicted and ground truth frames.
+    
+    Args:
+        pred_frames: Predicted frames array (T, C, H, W)
+        gt_frames: Ground truth frames array (T, C, H, W)
+        data_range: Data range of the frames (default: 1.0 for [0, 1] range)
+        
+    Returns:
+        dict: Dictionary containing overall PSNR and frame-wise PSNR
+    """
+    # Ensure same shape
+    assert pred_frames.shape == gt_frames.shape, \
+        f"Shape mismatch: pred {pred_frames.shape} vs gt {gt_frames.shape}"
+    
+    T, C, H, W = pred_frames.shape
+    
+    # Compute frame-wise PSNR
+    frame_wise_psnr = []
+    for t in range(T):
+        psnr = calculate_psnr(pred_frames[t], gt_frames[t], data_range=data_range)
+        frame_wise_psnr.append(psnr)
+    
+    # Compute overall PSNR
+    overall_psnr = calculate_psnr(pred_frames, gt_frames, data_range=data_range)
+    
+    return {
+        'overall_psnr': overall_psnr,
+        'frame_wise_psnr': frame_wise_psnr,
+        'mean_frame_psnr': np.mean(frame_wise_psnr),
+        'std_frame_psnr': np.std(frame_wise_psnr),
+        'num_frames': T
+    }
+
+
+def main():
+    # Directories
+    gt_dir = Path("/mnt/WorldMem/outputs/2025-12-01/23-39-44/videos/test_vis/gt")
+    pred_dir = Path("/mnt/WorldMem/outputs/2025-12-01/23-39-44/videos/test_vis/pred")
+    
+    # Get all video files in gt directory
+    gt_videos = sorted(gt_dir.glob("*.mp4"))
+    
+    print(f"Found {len(gt_videos)} videos to process")
+    print("=" * 80)
+    
+    results = []
+    
+    for gt_path in tqdm(gt_videos, desc="Computing PSNR"):
+        video_name = gt_path.name
+        pred_path = pred_dir / video_name
+        
+        if not pred_path.exists():
+            print(f"Warning: Prediction video not found for {video_name}")
+            continue
+        
+        try:
+            # Load videos
+            gt_frames = load_video(gt_path)
+            pred_frames = load_video(pred_path)
+            
+            # Compute PSNR
+            psnr_results = compute_psnr(pred_frames, gt_frames, data_range=1.0)
+            
+            # Store results
+            result = {
+                'video_name': video_name,
+                'overall_psnr': psnr_results['overall_psnr'],
+                'mean_frame_psnr': psnr_results['mean_frame_psnr'],
+                'std_frame_psnr': psnr_results['std_frame_psnr'],
+                'num_frames': psnr_results['num_frames']
+            }
+            results.append(result)
+            
+        except Exception as e:
+            print(f"Error processing {video_name}: {str(e)}")
+            continue
+    
+    # Create DataFrame for better visualization
+    df = pd.DataFrame(results)
+    
+    # Print results
+    print("\n" + "=" * 80)
+    print("PSNR Results for Each Video:")
+    print("=" * 80)
+    print(df.to_string(index=False))
+    
+    # Print summary statistics
+    print("\n" + "=" * 80)
+    print("Summary Statistics:")
+    print("=" * 80)
+    print(f"Average PSNR across all videos: {df['overall_psnr'].mean():.4f} ± {df['overall_psnr'].std():.4f}")
+    print(f"Min PSNR: {df['overall_psnr'].min():.4f} ({df.loc[df['overall_psnr'].idxmin(), 'video_name']})")
+    print(f"Max PSNR: {df['overall_psnr'].max():.4f} ({df.loc[df['overall_psnr'].idxmax(), 'video_name']})")
+    print(f"Median PSNR: {df['overall_psnr'].median():.4f}")
+    
+    # Group by video ID (video_0, video_1, etc.) and rank
+    df['video_id'] = df['video_name'].str.extract(r'(video_\d+)')[0]
+    df['rank'] = df['video_name'].str.extract(r'rank(\d+)')[0].astype(int)
+    
+    print("\n" + "=" * 80)
+    print("Average PSNR by Video ID:")
+    print("=" * 80)
+    video_group = df.groupby('video_id')['overall_psnr'].agg(['mean', 'std', 'count'])
+    print(video_group.to_string())
+    
+    print("\n" + "=" * 80)
+    print("Average PSNR by Rank:")
+    print("=" * 80)
+    rank_group = df.groupby('rank')['overall_psnr'].agg(['mean', 'std', 'count'])
+    print(rank_group.to_string())
+    
+    # Save results to CSV
+    output_csv = gt_dir.parent / "psnr_results.csv"
+    df.to_csv(output_csv, index=False)
+    print(f"\nResults saved to: {output_csv}")
+    
+    # Save summary to text file
+    output_txt = gt_dir.parent / "psnr_summary.txt"
+    with open(output_txt, 'w') as f:
+        f.write("PSNR Results Summary\n")
+        f.write("=" * 80 + "\n\n")
+        f.write(f"Average PSNR: {df['overall_psnr'].mean():.4f} ± {df['overall_psnr'].std():.4f}\n")
+        f.write(f"Min PSNR: {df['overall_psnr'].min():.4f}\n")
+        f.write(f"Max PSNR: {df['overall_psnr'].max():.4f}\n")
+        f.write(f"Median PSNR: {df['overall_psnr'].median():.4f}\n")
+        f.write("\n" + "=" * 80 + "\n")
+        f.write("Full Results:\n")
+        f.write("=" * 80 + "\n")
+        f.write(df.to_string(index=False))
+    print(f"Summary saved to: {output_txt}")
+
+
+if __name__ == "__main__":
+    main()
+

configurations/algorithm/df_video_worldmemminecraft.yaml

@@ -35,9 +35,4 @@ diffusion:
     use_linear_attn: True
     time_emb_type: rotary
 
-metrics:
-  # - fvd
-  # - fid
-  # - lpips
-
 _name: df_video_worldmemminecraft

configurations/huggingface.yaml

@@ -54,7 +54,7 @@ noise_level: random_all
 causal: True
 x_shape: [3, 360, 640]
 context_frames: 1
-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
+diffusion_path: zeqixiao/worldmem_checkpoints/diffusion_only.ckpt
+vae_path: zeqixiao/worldmem_checkpoints/vae_only.ckpt
+pose_predictor_path: zeqixiao/worldmem_checkpoints/pose_prediction_model_only.ckpt
 next_frame_length: 1

data_generator.py

@@ -0,0 +1,330 @@
+"""
+MineDojo Episode Collection Script
+
+This script generates trajectory data from MineDojo environment using a simple agent
+that randomly explores the environment. It supports parallel data collection and
+saves video and action/pose data.
+"""
+
+import argparse
+import math
+import multiprocessing as mp
+import os
+import os.path as osp
+import random
+from typing import Dict, Optional, Tuple
+
+import cv2
+import minedojo
+import numpy as np
+from tqdm import tqdm
+
+# Action mappings for the agent
+# Format: [forward/back, ?, ?, pitch, yaw, ?, ?, ?]
+ACTIONS: Dict[str, np.ndarray] = {
+    'forward': np.array([1, 0, 0, 12, 12, 0, 0, 0]),
+    'back': np.array([2, 0, 0, 12, 12, 0, 0, 0]),
+    'left': np.array([0, 0, 0, 12, 11, 0, 0, 0]),
+    'right': np.array([0, 0, 0, 12, 13, 0, 0, 0]),
+    'up': np.array([0, 0, 0, 11, 12, 0, 0, 0]),
+    'down': np.array([0, 0, 0, 13, 12, 0, 0, 0]),
+    'noop': np.array([0, 0, 0, 12, 12, 0, 0, 0])
+}
+
+
+def sample_action(prob_forward: float) -> str:
+    """
+    Sample an action based on forward probability.
+    
+    Args:
+        prob_forward: Probability of moving forward or backward
+        
+    Returns:
+        Action name string
+    """
+    prob_turn = (1 - prob_forward) / 2
+    action = np.random.choice(
+        ['forward', 'back', 'left', 'right', 'up', 'down'],
+        p=[prob_forward / 2 - 0.1, prob_forward / 2 - 0.1, prob_turn, prob_turn, 0.1, 0.1]
+    )
+    return action
+
+
+class SimpleAgent:
+    """
+    Simple agent that explores the environment with random actions.
+    
+    Attributes:
+        action_repeat: Number of times to repeat the same action
+        prob_forward: Probability of moving forward/backward
+        max_consec_fwd: Maximum consecutive forward actions (currently unused)
+    """
+    
+    def __init__(self, prob_forward: float, action_repeat: int, max_consec_fwd: int):
+        """
+        Initialize the SimpleAgent.
+        
+        Args:
+            prob_forward: Probability of moving forward or backward
+            action_repeat: How many steps to repeat each sampled action
+            max_consec_fwd: Maximum consecutive forward movements
+        """
+        self.action_repeat = action_repeat
+        self.prob_forward = prob_forward
+        self.max_consec_fwd = max_consec_fwd
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset the agent's internal state."""
+        self.n_fwd = 0
+        self.counter = 0
+        self.action = None
+
+    def sample(self, pos: np.ndarray) -> np.ndarray:
+        """
+        Sample an action given the current position.
+        
+        Args:
+            pos: Current position array containing [x, y, z, pitch, yaw]
+            
+        Returns:
+            Action array
+        """
+        prob_forward = self.prob_forward
+
+        # Sample new action if needed (or if previous was up/down)
+        if (self.action is None or 
+            self.counter % self.action_repeat == 0 or 
+            self.action in ['up', 'down']):
+            self.action = sample_action(prob_forward)
+
+        self.counter += 1
+        return ACTIONS[self.action]
+
+
+def collect_episode(env, agent: SimpleAgent, traj_length: int) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
+    """
+    Collect a single episode of trajectory data.
+    
+    Args:
+        env: MineDojo environment
+        agent: Agent to collect data with
+        traj_length: Length of trajectory to collect
+        
+    Returns:
+        Tuple of (rgb observations, actions, poses) or None if collection fails
+    """
+    agent.reset()
+
+    # Retry environment reset until successful
+    success = False
+    max_retries = 10
+    retries = 0
+    while not success and retries < max_retries:
+        try:
+            obs = env.reset()
+            success = True
+        except Exception as e:
+            retries += 1
+            if retries >= max_retries:
+                print(f"Failed to reset environment after {max_retries} retries")
+                return None
+
+    observations = [obs['rgb']]
+    actions = [env.action_space.no_op()]
+    pose = [np.concatenate([
+        obs['location_stats']['pos'],
+        obs['location_stats']['pitch'],
+        obs['location_stats']['yaw']
+    ])]
+
+    for ei in range(traj_length):
+        curr_actions = agent.sample(pose[ei])
+        obs, reward, done, info = env.step(curr_actions)
+        
+        actions.append(curr_actions)
+        observations.append(obs['rgb'])
+        pose.append(np.concatenate([
+            obs['location_stats']['pos'],
+            obs['location_stats']['pitch'],
+            obs['location_stats']['yaw']
+        ]))
+    
+    rgb = np.stack(observations, axis=0)
+    actions = np.array(actions, dtype=np.int32)
+    pose = np.array(pose)
+
+    return rgb, actions, pose
+
+
+def worker(worker_id: int, args: argparse.Namespace) -> None:
+    """
+    Worker process for parallel data collection.
+    
+    Args:
+        worker_id: Unique ID for this worker process
+        args: Command-line arguments
+    """
+    # Create worker-specific output directory
+    worker_output_dir = osp.join(args.output_dir, f'{worker_id}')
+    os.makedirs(worker_output_dir, exist_ok=True)
+    
+    # Set worker-specific random seeds for reproducibility
+    # Use a large offset between workers to ensure independent random streams
+    worker_seed = args.base_seed + worker_id * 10000
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+    agent = SimpleAgent(args.prob_forward, args.action_repeat, args.max_consec_fwd)
+
+    # Calculate number of episodes for this worker
+    num_episodes = args.num_episodes // args.n_parallel
+    if worker_id < (args.num_episodes % args.n_parallel):
+        num_episodes += 1
+    
+    pbar = tqdm(total=num_episodes, position=worker_id, desc=f"Worker {worker_id}")
+    episode_count = 0
+    
+    while episode_count < num_episodes:
+        # Create environment with unique seeds for each worker and episode
+        # Ensure world_seed and seed are different for each worker and episode
+        episode_seed_base = worker_seed + episode_count * 100
+        world_seed = episode_seed_base
+        env_seed = episode_seed_base + 1
+        
+        env = minedojo.make(
+            task_id="open-ended",
+            image_size=(360, 640),
+            world_seed=world_seed,
+            seed=env_seed,
+            generate_world_type='specified_biome',
+            specified_biome=args.env_type,
+            initial_weather='rain'
+        )
+
+        # Collect episode data
+        out = collect_episode(env, agent, args.traj_length)
+        if out is None:
+            env.close()
+            continue
+
+        rgb, actions, poses = out
+        
+        # Save video
+        video_fname = osp.join(worker_output_dir, f'{episode_count:06d}.mp4')
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        writer = cv2.VideoWriter(video_fname, fourcc, 10.0, (rgb.shape[3], rgb.shape[2]))
+        
+        for t in range(rgb.shape[0]):
+            frame = rgb[t].transpose(1, 2, 0)
+            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+            writer.write(frame)
+        writer.release()
+
+        # Save actions and poses
+        action_fname = osp.join(worker_output_dir, f'{episode_count:06d}.npz')
+        np.savez_compressed(action_fname, actions=actions, poses=poses)
+        
+        episode_count += 1
+        env.close()
+        pbar.update(1)
+    
+    pbar.close()
+
+
+def main(args: argparse.Namespace) -> None:
+    """
+    Main function to orchestrate parallel data collection.
+    
+    Args:
+        args: Command-line arguments
+    """
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Create and start worker processes
+    procs = [mp.Process(target=worker, args=(i, args)) for i in range(args.n_parallel)]
+    for p in procs:
+        p.start()
+    for p in procs:
+        p.join()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Generate MineDojo trajectory data with parallel collection'
+    )
+    parser.add_argument(
+        '-o', '--output_dir',
+        type=str,
+        default='test',
+        help='Output directory for generated data'
+    )
+    parser.add_argument(
+        '--env_type',
+        type=str,
+        default='test',
+        help='Biome type for environment generation'
+    )
+    parser.add_argument(
+        '-z', '--n_parallel',
+        type=int,
+        default=1,
+        help='Number of parallel workers (default: 1)'
+    )
+    parser.add_argument(
+        '-a', '--action_repeat',
+        type=int,
+        default=5,
+        help='Number of times to repeat each action (default: 5)'
+    )
+    parser.add_argument(
+        '-p', '--prob_forward',
+        type=float,
+        default=0.7,
+        help='Probability of forward/backward actions (default: 0.7)'
+    )
+    parser.add_argument(
+        '-m', '--max_consec_fwd',
+        type=int,
+        default=50,
+        help='Maximum consecutive forward movements (default: 50)'
+    )
+    parser.add_argument(
+        '-t', '--traj_length',
+        type=int,
+        default=1500,
+        help='Length of each trajectory (default: 1500)'
+    )
+    parser.add_argument(
+        '-n', '--num_episodes',
+        type=int,
+        default=100000,
+        help='Total number of episodes to generate (default: 100000)'
+    )
+    parser.add_argument(
+        '-r', '--resolution',
+        type=int,
+        default=128,
+        help='Resolution (currently unused, default: 128)'
+    )
+    parser.add_argument(
+        '-rh', '--resolution_h',
+        type=int,
+        default=360,
+        help='Height resolution (currently unused, default: 360)'
+    )
+    parser.add_argument(
+        '-rw', '--resolution_w',
+        type=int,
+        default=640,
+        help='Width resolution (currently unused, default: 640)'
+    )
+    parser.add_argument(
+        '--base_seed',
+        type=int,
+        default=42,
+        help='Base RNG seed; worker i uses base_seed+i (default: 42)'
+    )
+    
+    args = parser.parse_args()
+    main(args)

datasets/video/minecraft_video_dataset.py

@@ -66,22 +66,19 @@ class MinecraftVideoDataset(BaseVideoDataset):
         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.n_frames = cfg.n_frames_valid if split == "validation" or split == "test" and hasattr(cfg, "n_frames_valid") else cfg.n_frames
+        self.memory_condition_length = getattr(cfg, "memory_condition_length", 8)
         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.add_timestamp_embedding = getattr(cfg, "add_timestamp_embedding", True)
         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.
@@ -99,9 +96,9 @@ class MinecraftVideoDataset(BaseVideoDataset):
             # 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:
+        if (split == "validation" or split == "test") and self.wo_updown:
             paths = [p for p in paths if "w_updown" not in str(p)]
-        elif split == "validation":
+        elif split == "validation" or split == "test":
             paths = [p for p in paths if "w_updown" in str(p)]
 
         if not paths:
@@ -129,7 +126,7 @@ class MinecraftVideoDataset(BaseVideoDataset):
             try:
                 return self.load_data(idx)
             except Exception as e:
-                print(f"Retrying due to error: {e}")
+                # print(f"Retrying due to error: {e}")
                 idx = (idx + 1) % len(self)
 
     def load_data(self, idx):
@@ -147,7 +144,8 @@ class MinecraftVideoDataset(BaseVideoDataset):
 
         # 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}"
+        # assert poses_pool[:, 1].ptp() < 2, f"Height variation too large: {poses_pool[:, 1].ptp()} - {video_path}"
+        assert poses_pool[:, 1].ptp() < 2
 
         # Pad poses if shorter than actions
         if len(poses_pool) < len(actions_pool):

evaluate.sh

@@ -0,0 +1,22 @@
+export PYTHONWARNINGS="ignore"
+wandb offline
+python -m main +name=infer \
+    experiment.tasks=[test] \
+    dataset.validation_multiplier=1 \
+    +diffusion_model_path=zeqixiao/worldmem_checkpoints/diffusion_only.ckpt \
+    +vae_path=zeqixiao/worldmem_checkpoints/vae_only.ckpt \
+    +customized_load=true \
+    +seperate_load=true \
+    dataset.n_frames=8 \
+    dataset.save_dir=data/minecraft \
+    +dataset.n_frames_valid=700 \
+    algorithm.diffusion.sampling_timesteps=20 \
+    +algorithm.memory_condition_length=8 \
+    +algorithm.lpips_batch_size=16 \
+    +algorithm.log_video=true \
+    +algorithm.save_local=true \
+    +dataset.customized_validation=true \
+    +algorithm.n_tokens=8 \
+    algorithm.context_frames=600 \
+    experiment.test.batch_size=1 \
+    experiment.test.limit_batch=10 \

experiments/exp_base.py

@@ -49,23 +49,14 @@ def load_custom_checkpoint(algo, checkpoint_path):
         checkpoint_path = Path(checkpoint_path)
 
     if is_huggingface_model(str(checkpoint_path)):
-        # Load from Hugging Face Hub if the path contains 'yslan'
+        # Load from Hugging Face Hub if the path contains 'zeqixiao'
         hf_ckpt = str(checkpoint_path).split('/')
         repo_id = '/'.join(hf_ckpt[:2])
         file_name = '/'.join(hf_ckpt[2:])
         model_path = hf_hub_download(repo_id=repo_id, filename=file_name)
         ckpt = torch.load(model_path, map_location=torch.device('cpu'))
 
-        # quick workaround
-        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=False)
+        algo.load_state_dict(ckpt['state_dict'], strict=True)
 
     elif checkpoint_path.suffix == ".pt":
         # Load from a .pt file
@@ -106,11 +97,6 @@ def load_custom_checkpoint(algo, checkpoint_path):
             if not k in ["data_mean", "data_std"]
         }
 
-        # for k, v in filtered_state_dict.items():
-        #     if "frame_timestep_embedder" in k:
-        #         new_k = k.replace("frame_timestep_embedder", "timestamp_embedding")
-        #         filtered_state_dict[new_k] = v
-
         algo.load_state_dict(filtered_state_dict, strict=False)
 
     else:
@@ -436,13 +422,37 @@ class BaseLightningExperiment(BaseExperiment):
             detect_anomaly=False,  # self.cfg.debug,
         )
 
-        # Only load the checkpoint if only testing. Otherwise, it will have been loaded
-        # and further trained during train.
-        trainer.test(
-            self.algo,
-            dataloaders=self._build_test_loader(),
-            ckpt_path=self.ckpt_path,
-        )
+        if self.customized_load:
+            if self.seperate_load:
+                if 'oasis500m' in self.diffusion_model_path:
+                    load_custom_checkpoint(algo=self.algo.diffusion_model.model,checkpoint_path=self.diffusion_model_path)
+                else:
+                    load_custom_checkpoint(algo=self.algo.diffusion_model,checkpoint_path=self.diffusion_model_path)
+                load_custom_checkpoint(algo=self.algo.vae,checkpoint_path=self.vae_path)
+            else:
+                load_custom_checkpoint(algo=self.algo,checkpoint_path=self.ckpt_path)
+
+            if self.zero_init_gate:
+                for name, para in self.algo.diffusion_model.named_parameters():
+                    if 'r_adaLN_modulation' in name:
+                        para.requires_grad_(False)
+                        para[2*1024:3*1024] = 0
+                        para[5*1024:6*1024] = 0
+                        para.requires_grad_(True)
+            
+            trainer.test(
+                self.algo,
+                dataloaders=self._build_test_loader(),
+                ckpt_path=None,
+            )
+        else:
+            trainer.test(
+                self.algo,
+                dataloaders=self._build_test_loader(),
+                ckpt_path=self.ckpt_path,
+            )
+
+
         if not self.algo:
             self.algo = self._build_algo()
         if self.cfg.validation.compile:

infer.sh

@@ -1,9 +1,10 @@
-wandb enabled
+export PYTHONWARNINGS="ignore"
+wandb offline
 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 \
+    +diffusion_model_path=zeqixiao/worldmem_checkpoints/diffusion_only.ckpt \
+    +vae_path=zeqixiao/worldmem_checkpoints/vae_only.ckpt \
     +customized_load=true \
     +seperate_load=true \
     dataset.n_frames=8 \
@@ -17,5 +18,4 @@ python -m main +name=infer \
     algorithm.context_frames=600 \
     +algorithm.relative_embedding=true \
     +algorithm.log_video=true \
-    +algorithm.add_timestamp_embedding=true \
-    algorithm.metrics=[lpips,psnr] \
+    +algorithm.add_timestamp_embedding=true

utils/logging_utils.py

@@ -2,6 +2,7 @@ from typing import Optional
 import wandb
 import numpy as np
 import torch
+import os
 
 import matplotlib.pyplot as plt
 import cv2
@@ -9,7 +10,8 @@ import matplotlib.pyplot as plt
 from tqdm import trange, tqdm
 import matplotlib.animation as animation
 from pathlib import Path
-
+import imageio
+        
 plt.set_loglevel("warning")
 
 from torchmetrics.functional import mean_squared_error, peak_signal_noise_ratio
@@ -34,6 +36,10 @@ def log_video(
     context_frames=0,
     color=(255, 0, 0),
     logger=None,
+    fps=15,
+    format="mp4",
+    save_local=True,
+    local_save_dir=None,
 ):
     """
     take in video tensors in range [-1, 1] and log into wandb
@@ -46,37 +52,139 @@ def log_video(
     :param context_frames: an int indicating how many frames in observation_hat are ground truth given as context
     :param color: a tuple of 3 numbers specifying the color of the border for ground truth frames
     :param logger: optional logger to use. use global wandb if not specified
+    :param fps: frames per second for the video (default: 15)
+    :param format: video format, either "mp4" or "gif" (default: "mp4")
+    :param save_local: whether to save videos to local disk (default: True)
+    :param local_save_dir: directory to save local videos. If None, uses hydra output dir
     """
+    import cv2
+    import hydra
+    from pathlib import Path
+    
+    # Get local rank for distributed training
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    
     if not logger:
         logger = wandb
     
-        # observation_gt = torch.zeros_like(observation_hat)
-    # observation_hat[:context_frames] = observation_gt[:context_frames]
-    # Add red border of 1 pixel width to the context frames
-    # for i, c in enumerate(color):
-    #     c = c / 255.0
-    #     observation_hat[:context_frames, :, i, [0, -1], :] = c
-    #     observation_hat[:context_frames, :, i, :, [0, -1]] = c
-
-    #     if observation_gt is not None:
-    #         observation_gt[:context_frames, :, i, [0, -1], :] = c
-    #         observation_gt[:context_frames, :, i, :, [0, -1]] = c
-    
+    # Prepare video tensors
+    observation_hat_np = observation_hat.detach().cpu().numpy()
     if observation_gt is not None:
-        video = torch.cat([observation_hat, observation_gt], -2).detach().cpu().numpy()
+        observation_gt_np = observation_gt.detach().cpu().numpy()
     else:
-        video = torch.cat([observation_hat], -1).detach().cpu().numpy()
-    video = np.transpose(np.clip(video, a_min=0.0, a_max=1.0) * 255, (1, 0, 2, 3, 4)).astype(np.uint8)
-    # video[..., 1:] = video[..., :1]  # remove framestack, only visualize current frame
-    n_samples = len(video)
-    # use wandb directly here since pytorch lightning doesn't support logging videos yet
+        observation_gt_np = None
+    
+    # Normalize to 0-255
+    observation_hat_np = np.transpose(np.clip(observation_hat_np, a_min=0.0, a_max=1.0) * 255, (1, 0, 2, 3, 4)).astype(np.uint8)
+    if observation_gt_np is not None:
+        observation_gt_np = np.transpose(np.clip(observation_gt_np, a_min=0.0, a_max=1.0) * 255, (1, 0, 2, 3, 4)).astype(np.uint8)
+    
+    n_samples = len(observation_hat_np)
+    
+    # Setup local save directory
+    if save_local:
+        if local_save_dir is None:
+            try:
+                hydra_cfg = hydra.core.hydra_config.HydraConfig.get()
+                output_dir = Path(hydra_cfg.runtime.output_dir)
+            except:
+                output_dir = Path.cwd() / "outputs"
+            local_save_dir = output_dir / "videos" / namespace
+        else:
+            local_save_dir = Path(local_save_dir)
+        
+        local_save_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Save pred videos locally
+        pred_dir = local_save_dir / "pred"
+        pred_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Save gt videos locally if available
+        if observation_gt_np is not None:
+            gt_dir = local_save_dir / "gt"
+            gt_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Save videos
     for i in range(n_samples):
-        logger.log(
-            {
-                f"{namespace}/{prefix}_{i}": wandb.Video(video[i], fps=5),
-                f"trainer/global_step": step,
-            }
-        )
+        video_pred = observation_hat_np[i]  # (T, C, H, W)
+        
+        if save_local:
+            # Save prediction video
+            if step is not None:
+                video_filename_pred = f"{prefix}_{i}_rank{local_rank}_step{step}.{format}"
+            else:
+                video_filename_pred = f"{prefix}_{i}_rank{local_rank}.{format}"
+            
+            video_path_pred = pred_dir / video_filename_pred
+            _save_video_to_file(video_pred, str(video_path_pred), fps)
+            
+            # Save ground truth video if available
+            if observation_gt_np is not None:
+                video_gt = observation_gt_np[i]
+                if step is not None:
+                    video_filename_gt = f"{prefix}_{i}_rank{local_rank}_step{step}.{format}"
+                else:
+                    video_filename_gt = f"{prefix}_{i}_rank{local_rank}.{format}"
+                
+                video_path_gt = gt_dir / video_filename_gt
+                _save_video_to_file(video_gt, str(video_path_gt), fps)
+        
+        # Log to wandb (only rank 0 to avoid duplicate logging)
+        if local_rank == 0 and logger:
+            # Concatenate pred and gt side by side for visualization
+            if observation_gt_np is not None:
+                video_combined = torch.cat([
+                    torch.from_numpy(observation_hat_np),
+                    torch.from_numpy(observation_gt_np)
+                ], -2).numpy()  # Concatenate along width
+                logger.log(
+                    {
+                        f"{namespace}/{prefix}_{i}": wandb.Video(video_combined[i], fps=fps, format=format),
+                        f"trainer/global_step": step,
+                    }
+                )
+            else:
+                logger.log(
+                    {
+                        f"{namespace}/{prefix}_{i}": wandb.Video(video_pred, fps=fps, format=format),
+                        f"trainer/global_step": step,
+                    }
+                )
+
+
+def _save_video_to_file(video_tensor, output_path, fps=15):
+    """
+    Save a video tensor to file using imageio (better compatibility than cv2).
+    
+    :param video_tensor: numpy array of shape (T, C, H, W) with values in [0, 255]
+    :param output_path: path to save the video
+    :param fps: frames per second
+    """
+
+    T, C, H, W = video_tensor.shape
+    
+    # Convert from (T, C, H, W) to (T, H, W, C)
+    video_tensor = np.transpose(video_tensor, (0, 2, 3, 1))
+    
+    # Ensure uint8
+    video_tensor = video_tensor.astype(np.uint8)
+    
+    # Save using imageio with H.264 codec (best compatibility)
+    writer = imageio.get_writer(
+        output_path, 
+        fps=fps,
+        codec='libx264',  # H.264 codec - widely supported
+        quality=8,  # Good quality (scale 0-10, 10 is best)
+        pixelformat='yuv420p',  # Standard pixel format for compatibility
+        macro_block_size=1  # Better quality
+    )
+    
+    for frame in video_tensor:
+        writer.append_data(frame)
+    
+    writer.close()
+        
+
 
 
 def get_validation_metrics_for_videos(
@@ -85,6 +193,7 @@ def get_validation_metrics_for_videos(
     lpips_model: Optional[LearnedPerceptualImagePatchSimilarity] = None,
     fid_model: Optional[FrechetInceptionDistance] = None,
     fvd_model: Optional[FrechetVideoDistance] = None,
+    lpips_batch_size: int = 100,
 ):
     """
     :param observation_hat: predicted observation tensor of shape (frame, batch, channel, height, width)
@@ -92,6 +201,7 @@ def get_validation_metrics_for_videos(
     :param lpips_model: a LearnedPerceptualImagePatchSimilarity object from algorithm.common.metrics
     :param fid_model: a FrechetInceptionDistance object  from algorithm.common.metrics
     :param fvd_model: a FrechetVideoDistance object  from algorithm.common.metrics
+    :param lpips_batch_size: batch size for LPIPS calculation to avoid OOM (default: 100)
     :return: a tuple of metrics
     """
     frame, batch, channel, height, width = observation_hat.shape
@@ -104,39 +214,51 @@ def get_validation_metrics_for_videos(
     observation_hat = observation_hat.float()
     observation_gt = observation_gt.float()
 
-    # observation_hat = observation_hat.float().to(next(lpips_model.parameters()).device)
-    # observation_gt = observation_gt.float().to(next(lpips_model.parameters()).device)
-    # if fvd_model is not None:
-    #     output_dict["fvd"] = fvd_model.compute(torch.clamp(observation_hat, -1.0, 1.0), torch.clamp(observation_gt, -1.0, 1.0))
+    # Clip to [0, 1] range before computing metrics (matching video saving behavior)
+    observation_hat_clipped = torch.clamp(observation_hat, 0.0, 1.0)
+    observation_gt_clipped = torch.clamp(observation_gt, 0.0, 1.0)
 
+    # Compute frame-wise PSNR
     frame_wise_psnr = []
-    for f in range(observation_hat.shape[0]):
-        frame_wise_psnr.append(peak_signal_noise_ratio(observation_hat[f], observation_gt[f], data_range=2.0))
+    for f in range(observation_hat_clipped.shape[0]):
+        frame_wise_psnr.append(peak_signal_noise_ratio(observation_hat_clipped[f], observation_gt_clipped[f], data_range=1.0))
     frame_wise_psnr = torch.stack(frame_wise_psnr)
 
     output_dict["frame_wise_psnr"] = frame_wise_psnr
-    observation_hat = observation_hat.view(-1, channel, height, width)
-    observation_gt = observation_gt.view(-1, channel, height, width)
-
-    output_dict["mse"] = mean_squared_error(observation_hat, observation_gt)
+    observation_hat_clipped = observation_hat_clipped.view(-1, channel, height, width)
+    observation_gt_clipped = observation_gt_clipped.view(-1, channel, height, width)
 
-    output_dict["psnr"] = peak_signal_noise_ratio(observation_hat, observation_gt, data_range=2.0)
-    # output_dict["ssim"] = structural_similarity_index_measure(observation_hat, observation_gt, data_range=2.0)
-    # output_dict["uiqi"] = universal_image_quality_index(observation_hat, observation_gt)
-    # operations for LPIPS and FID
-    observation_hat = torch.clamp(observation_hat, -1.0, 1.0)
-    observation_gt = torch.clamp(observation_gt, -1.0, 1.0)
+    # Compute MSE and PSNR on clipped data
+    output_dict["mse"] = mean_squared_error(observation_hat_clipped, observation_gt_clipped)
+    output_dict["psnr"] = peak_signal_noise_ratio(observation_hat_clipped, observation_gt_clipped, data_range=1.0)
+    # output_dict["ssim"] = structural_similarity_index_measure(observation_hat_clipped, observation_gt_clipped, data_range=1.0)
+    # output_dict["uiqi"] = universal_image_quality_index(observation_hat_clipped, observation_gt_clipped)
 
+    # LPIPS computation
     if lpips_model is not None:
-        lpips_model.update(observation_hat, observation_gt)
+        # Process LPIPS in batches to avoid OOM
+        num_frames = observation_hat_clipped.shape[0]
+        
+        for i in range(0, num_frames, lpips_batch_size):
+            batch_end = min(i + lpips_batch_size, num_frames)
+            observation_hat_batch = observation_hat_clipped[i:batch_end]
+            observation_gt_batch = observation_gt_clipped[i:batch_end]
+            
+            lpips_model.update(observation_hat_batch, observation_gt_batch)
+            
+            # Free GPU memory after each batch
+            del observation_hat_batch, observation_gt_batch
+            torch.cuda.empty_cache()
+        
         lpips = lpips_model.compute().item()
         # Reset the states of non-functional metrics
         output_dict["lpips"] = lpips
         lpips_model.reset()
 
+    # FID computation
     if fid_model is not None:
-        observation_hat_uint8 = ((observation_hat + 1.0) / 2 * 255).type(torch.uint8)
-        observation_gt_uint8 = ((observation_gt + 1.0) / 2 * 255).type(torch.uint8)
+        observation_hat_uint8 = (observation_hat_clipped * 255).type(torch.uint8)
+        observation_gt_uint8 = (observation_gt_clipped * 255).type(torch.uint8)
         fid_model.update(observation_gt_uint8, real=True)
         fid_model.update(observation_hat_uint8, real=False)
         fid = fid_model.compute()