Keep reproduction train and eval fixes

.gitignore

@@ -0,0 +1,52 @@
+# Local environment and editor folders
+.*/
+!.gitignore
+!.gitattributes
+!.github/
+
+# Python/runtime caches
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+
+# Training, evaluation, and logging artifacts
+outputs/
+output/
+checkpoints/
+lightning_logs/
+logs/
+slurm_logs/
+wandb/
+ablation_results/
+latest-run
+*.log
+
+# Local model/data clones and generated assets
+open-oasis/
+proposal/
+tools/
+*.docx
+
+# Local reproduction and ablation helper scripts
+benchmark_memory_latency.sh
+evaluate_600_prefix_mem_len_*.sh
+evaluate_dense_reference*.sh
+evaluate_dit_no_memory.sh
+evaluate_oasis_baseline.sh
+evaluate_unified.sh
+infer_dit_no_memory.sh
+infer_oasis_baseline.sh
+run_memory_*.sh
+train_debug_berzulis.sh
+train_stage_1_debug.sh
+utils/replay_offline_wandb.py
+
+# Local experimental modules/backups kept out of the clean branch
+algorithms/worldmem/*.bak_*
+algorithms/worldmem/causal_memory.py
+algorithms/worldmem/retrieval.py
+algorithms/worldmem/revisit_consistency.py
+algorithms/worldmem/models/causal_memory.py

algorithms/worldmem/df_video.py

@@ -3,6 +3,7 @@ import random
 import math
 import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from torchvision.transforms import InterpolationMode
@@ -376,7 +377,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
             ref_mode=self.ref_mode
         )
 
-        self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity()
+        self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity(sync_on_compute=False)
         vae = VAE_models["vit-l-20-shallow-encoder"]()
         self.vae = vae.eval()
 
@@ -475,65 +476,78 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         return {"loss": loss}
     
-    def on_validation_epoch_end(self, namespace="validation") -> None:
-        if not self.validation_step_outputs:
-            return
-        
-        xs_pred = []
-        xs = []
-        for pred, gt in self.validation_step_outputs:
-            xs_pred.append(pred)
-            xs.append(gt)
-
-        xs_pred = torch.cat(xs_pred, 1)
-        if gt is not None:
-            xs = torch.cat(xs, 1)
-        else:
-            xs = None
+    def on_validation_epoch_start(self) -> None:
+        self._reset_metric_accumulators()
+
+    def on_test_epoch_start(self) -> None:
+        self._reset_metric_accumulators()
+
+    def on_validation_epoch_end(self) -> None:
+        self._on_eval_epoch_end()
+
+    def on_test_epoch_end(self) -> None:
+        self._on_eval_epoch_end()
+
+    def _reset_metric_accumulators(self) -> None:
+        self._metric_device = next(self.validation_lpips_model.parameters()).device
+        self._mse_sum = torch.tensor(0.0, device=self._metric_device)
+        self._mse_count = torch.tensor(0.0, device=self._metric_device)
+        self._psnr_sum = torch.tensor(0.0, device=self._metric_device)
+        self._psnr_count = torch.tensor(0.0, device=self._metric_device)
+        self._lpips_sum = torch.tensor(0.0, device=self._metric_device)
+        self._lpips_count = torch.tensor(0.0, device=self._metric_device)
+
+    def _update_metric_accumulators(self, xs_pred: torch.Tensor, xs_gt: torch.Tensor) -> None:
+        if not hasattr(self, "_metric_device"):
+            self._reset_metric_accumulators()
+
+        xs_pred_device = xs_pred.to(self._metric_device)
+        xs_gt_device = xs_gt.to(self._metric_device)
+        metric_dict = get_validation_metrics_for_videos(
+            xs_pred_device,
+            xs_gt_device,
+            lpips_model=self.validation_lpips_model,
+            lpips_batch_size=self.lpips_batch_size,
+        )
 
-        if self.logger and self.log_video:
-            log_video(
-                xs_pred,
-                xs,
-                step=None if namespace == "test" else self.global_step,
-                namespace=namespace + "_vis",
-                context_frames=self.context_frames,
-                logger=self.logger.experiment,
-                save_local=self.save_local,
-                local_save_dir=self.local_save_dir,
-            )
+        mse_count = torch.tensor(float(xs_pred_device.numel()), device=self._metric_device)
+        psnr_count = torch.tensor(float(xs_pred_device.shape[1]), device=self._metric_device)
+        lpips_count = torch.tensor(float(xs_pred_device.shape[0] * xs_pred_device.shape[1]), device=self._metric_device)
 
-        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_device, xs_device, 
-                lpips_model=self.validation_lpips_model,
-                lpips_batch_size=self.lpips_batch_size)
-            
+        self._mse_sum += metric_dict["mse"].detach() * mse_count
+        self._mse_count += mse_count
+        self._psnr_sum += metric_dict["psnr"].detach() * psnr_count
+        self._psnr_count += psnr_count
+        self._lpips_sum += torch.tensor(float(metric_dict["lpips"]), device=self._metric_device) * lpips_count
+        self._lpips_count += lpips_count
+
+        del xs_pred_device, xs_gt_device
+
+    def _on_eval_epoch_end(self) -> None:
+        if not hasattr(self, "_metric_device"):
+            return
+
+        if dist.is_available() and dist.is_initialized():
+            for tensor in (
+                self._mse_sum,
+                self._mse_count,
+                self._psnr_sum,
+                self._psnr_count,
+                self._lpips_sum,
+                self._lpips_count,
+            ):
+                dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        if self._mse_count.item() > 0:
             self.log_dict(
-                {"mse": metric_dict['mse'],
-                "psnr": metric_dict['psnr'],
-                "lpips": metric_dict['lpips']},
-                sync_dist=True
+                {
+                    "mse": self._mse_sum / self._mse_count.clamp_min(1.0),
+                    "psnr": self._psnr_sum / self._psnr_count.clamp_min(1.0),
+                    "lpips": self._lpips_sum / self._lpips_count.clamp_min(1.0),
+                },
+                sync_dist=False,
             )
 
-            if self.log_curve:
-                psnr_values = metric_dict['frame_wise_psnr'].cpu().tolist()
-                frames = list(range(len(psnr_values)))
-                line_plot = wandb.plot.line_series(
-                    xs = frames,
-                    ys = [psnr_values],
-                    keys = ["PSNR"],
-                    title = "Frame-wise PSNR",
-                    xname = "Frame index"
-                )
-
-                self.logger.experiment.log({"frame_wise_psnr_plot": line_plot})
-      
         self.validation_step_outputs.clear()
 
     def _preprocess_batch(self, batch):
@@ -783,10 +797,24 @@ 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 (move to CPU to save GPU memory)
-        self.validation_step_outputs.append((xs_pred.detach().cpu(), xs_decode.detach().cpu()))
+        if self.logger and self.log_video:
+            log_video(
+                xs_pred,
+                xs_decode,
+                step=batch_idx,
+                namespace=namespace + "_vis",
+                context_frames=self.context_frames,
+                logger=self.logger.experiment,
+                save_local=self.save_local,
+                local_save_dir=self.local_save_dir,
+            )
+
+        self._update_metric_accumulators(xs_pred, xs_decode)
         return
 
+    def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        return self.validation_step(batch, batch_idx, namespace="test")
+
     @torch.no_grad()
     def interactive(self, first_frame, new_actions, first_pose, device,
                     memory_latent_frames, memory_actions, memory_poses, memory_c2w, memory_frame_idx):

configurations/experiment/base_pytorch.yaml

@@ -35,7 +35,7 @@ validation:
   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!)
+  inference_mode: False # 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
@@ -45,6 +45,7 @@ test:
   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.
+  inference_mode: False # whether to run test in inference mode (enable_grad won't work!)
   data:
     num_workers: 4 # number of CPU threads for data preprocessing, for test.
     shuffle: False # whether test data will be shuffled

evaluate.sh

@@ -20,4 +20,4 @@ python -m main +name=infer \
     +algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
     experiment.test.batch_size=1 \
-    experiment.test.limit_batch=10 \
+    experiment.test.limit_batch=10

experiments/exp_base.py

@@ -9,6 +9,7 @@ from abc import ABC, abstractmethod
 from typing import Optional, Union, Literal, List, Dict
 import pathlib
 import os
+from datetime import timedelta
 
 import hydra
 import torch
@@ -361,7 +362,7 @@ class BaseLightningExperiment(BaseExperiment):
             logger=self.logger,
             devices="auto",
             num_nodes=self.cfg.num_nodes,
-            strategy=DDPStrategy(find_unused_parameters=False) if torch.cuda.device_count() > 1 else "auto",
+            strategy=DDPStrategy(find_unused_parameters=False, timeout=timedelta(hours=1)) if torch.cuda.device_count() > 1 else "auto",
             callbacks=callbacks,
             limit_val_batches=self.cfg.validation.limit_batch,
             precision=self.cfg.validation.precision,
@@ -415,10 +416,11 @@ class BaseLightningExperiment(BaseExperiment):
             logger=self.logger,
             devices="auto",
             num_nodes=self.cfg.num_nodes,
-            strategy=DDPStrategy(find_unused_parameters=False) if torch.cuda.device_count() > 1 else "auto",
+            strategy=DDPStrategy(find_unused_parameters=False, timeout=timedelta(hours=1)) if torch.cuda.device_count() > 1 else "auto",
             callbacks=callbacks,
             limit_test_batches=self.cfg.test.limit_batch,
             precision=self.cfg.test.precision,
+            inference_mode=self.cfg.test.inference_mode,
             detect_anomaly=False,  # self.cfg.debug,
         )
 

main.py

@@ -59,6 +59,7 @@ def run_local(cfg: DictConfig):
         OmegaConf.set_readonly(hydra_cfg, True)
         
     output_dir = Path(hydra_cfg.runtime.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
     
     if is_rank_zero:
         print(cyan(f"Outputs will be saved to:"), output_dir)

train_stage_1.sh

@@ -22,7 +22,7 @@ python -m main +name=train \
     +algorithm.relative_embedding=true \
     +algorithm.log_video=true \
     +algorithm.add_timestamp_embedding=true \
-    algorithm.metrics=[lpips,psnr] \
+    +algorithm.metrics=[lpips,psnr] \
     experiment.training.checkpointing.every_n_train_steps=2500 \
     experiment.training.max_steps=120000
 

train_stage_2.sh

@@ -17,9 +17,9 @@ python -m main +name=train \
     +algorithm.relative_embedding=true \
     +algorithm.log_video=true \
     +algorithm.add_timestamp_embedding=true \
-    algorithm.metrics=[lpips,psnr] \
+    +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 \
+    resume="${RESUME_RUN_ID:?Set RESUME_RUN_ID to the W&B run id from stage 1}" \
+    +output_dir="${OUTPUT_DIR:-outputs/worldmem_stage_2}" \
     experiment.training.max_steps=240000
 

train_stage_3.sh

@@ -17,8 +17,8 @@ python -m main +name=train \
     +algorithm.relative_embedding=true \
     +algorithm.log_video=true \
     +algorithm.add_timestamp_embedding=true \
-    algorithm.metrics=[lpips,psnr] \
+    +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
+    resume="${RESUME_RUN_ID:?Set RESUME_RUN_ID to the W&B run id from stage 2}" \
+    +output_dir="${OUTPUT_DIR:-outputs/worldmem_stage_3}" \
+    experiment.training.max_steps=700000

utils/distributed_utils.py

@@ -1,3 +1,3 @@
-import wandb
+import os
 
-is_rank_zero = wandb.run is not None
+is_rank_zero = int(os.environ.get("RANK", 0)) == 0

utils/logging_utils.py

@@ -67,79 +67,81 @@ def log_video(
     if not logger:
         logger = wandb
     
-    # Prepare video tensors
-    observation_hat_np = observation_hat.detach().cpu().numpy()
-    if observation_gt is not None:
-        observation_gt_np = observation_gt.detach().cpu().numpy()
-    else:
-        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)
-    
+    n_samples = int(observation_hat.shape[1])
+
+    def _to_uint8_video(video_tensor):
+        return (
+            video_tensor.detach()
+            .clamp(0.0, 1.0)
+            .mul(255)
+            .to(torch.uint8)
+            .cpu()
+            .numpy()
+        )
+
+    is_rank_zero = local_rank == 0
+    save_local = bool(save_local and is_rank_zero)
+
     # 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:
+            except Exception:
                 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:
+        if observation_gt 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):
-        video_pred = observation_hat_np[i]  # (T, C, H, W)
-        
+        if not save_local and not (is_rank_zero and logger):
+            continue
+
+        video_pred = _to_uint8_video(observation_hat[:, i])  # (T, C, H, W)
+        video_gt = None
+        if observation_gt is not None:
+            video_gt = _to_uint8_video(observation_gt[:, i])
+
         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 video_gt is not None:
                 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
+        if is_rank_zero and logger:
+            if video_gt is not None:
+                video_combined = np.concatenate([video_pred, video_gt], axis=-1)
                 logger.log(
                     {
-                        f"{namespace}/{prefix}_{i}": wandb.Video(video_combined[i], fps=fps, format=format),
+                        f"{namespace}/{prefix}_{i}": wandb.Video(video_combined, fps=fps, format=format),
                         f"trainer/global_step": step,
                     }
                 )