Update DeMemWM validation metrics and launchers

algorithms/dememwm/df_video.py

@@ -738,6 +738,7 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         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)
+        self.metric_report_segment = max(0, int(getattr(cfg, "metric_report_segment", 0) or 0))
         self.log_per_frame_metrics = bool(getattr(cfg, "log_per_frame_metrics", False))
         self.log_revisit_metrics = bool(getattr(cfg, "log_revisit_metrics", False))
         self.log_memory_selection_sheet = bool(getattr(cfg, "log_memory_selection_sheet", False))
@@ -948,11 +949,18 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         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)
+        self._frame_metrics_synced = False
+        self._segment_metrics_synced = False
 
         if self.log_per_frame_metrics:
             self._frame_metric_eval_start = None
             for attr in ("mse_sum", "mse_count", "psnr_sum", "psnr_count", "lpips_sum", "lpips_count"):
                 setattr(self, f"_frame_{attr}", torch.empty(0, device=self._metric_device))
+        if self.metric_report_segment > 0:
+            self._segment_metric_eval_start = None
+            self._segment_metric_frames = 0
+            for attr in ("mse_sum", "mse_count", "psnr_sum", "psnr_count", "lpips_sum", "lpips_count"):
+                setattr(self, f"_segment_{attr}", torch.empty(0, device=self._metric_device))
         if self.log_revisit_metrics:
             self._reset_revisit_metric_accumulators()
 
@@ -960,6 +968,18 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         base = Path(self.local_save_dir) if self.local_save_dir is not None else Path.cwd() / "eval_artifacts"
         return base / "metrics"
 
+    def _wandb_run_files_dir(self) -> Path | None:
+        logger = getattr(self, "logger", None)
+        experiment = None if logger is None else getattr(logger, "experiment", None)
+        run_dir = None if experiment is None else getattr(experiment, "dir", None)
+        return None if run_dir is None else Path(run_dir)
+
+    def _segment_metric_artifact_dir(self) -> Path:
+        run_dir = self._wandb_run_files_dir()
+        if run_dir is not None:
+            return run_dir / "metrics"
+        return self._eval_artifact_dir()
+
     def _memory_sheet_artifact_dir(self) -> Path:
         base = Path(self.local_save_dir) if self.local_save_dir is not None else Path.cwd() / "eval_artifacts"
         return base / "memory_selection_sheets"
@@ -1096,13 +1116,17 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             setattr(self, name, torch.cat([getattr(self, name), torch.zeros(pad, device=self._metric_device)]))
 
     def _sync_frame_metric_accumulators(self) -> None:
+        if getattr(self, "_frame_metrics_synced", False):
+            return
         if not (dist.is_available() and dist.is_initialized()):
+            self._frame_metrics_synced = True
             return
         length = torch.tensor(int(self._frame_mse_sum.numel()), device=self._metric_device, dtype=torch.long)
         dist.all_reduce(length, op=dist.ReduceOp.MAX)
         self._pad_frame_metric_accumulators(int(length.item()))
         for attr in ("mse_sum", "mse_count", "psnr_sum", "psnr_count", "lpips_sum", "lpips_count"):
             dist.all_reduce(getattr(self, f"_frame_{attr}"), op=dist.ReduceOp.SUM)
+        self._frame_metrics_synced = True
 
     def _update_per_frame_metric_accumulators(
         self,
@@ -1115,6 +1139,7 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             return
         if not hasattr(self, "_frame_mse_sum"):
             self._reset_metric_accumulators()
+        self._frame_metrics_synced = False
 
         frames = int(xs_pred.shape[0])
         self._pad_frame_metric_accumulators(frames)
@@ -1149,6 +1174,100 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 self._frame_lpips_sum[frame_idx] += value * count
                 self._frame_lpips_count[frame_idx] += count
 
+    def _pad_segment_metric_accumulators(self, length: int) -> None:
+        current = int(self._segment_mse_sum.numel())
+        if current >= int(length):
+            return
+        pad = int(length) - current
+        for attr in ("mse_sum", "mse_count", "psnr_sum", "psnr_count", "lpips_sum", "lpips_count"):
+            name = f"_segment_{attr}"
+            setattr(self, name, torch.cat([getattr(self, name), torch.zeros(pad, device=self._metric_device)]))
+
+    def _sync_segment_metric_accumulators(self) -> None:
+        if getattr(self, "_segment_metrics_synced", False):
+            return
+        if not (dist.is_available() and dist.is_initialized()):
+            self._segment_metrics_synced = True
+            return
+        length = torch.tensor(int(self._segment_mse_sum.numel()), device=self._metric_device, dtype=torch.long)
+        frames = torch.tensor(int(getattr(self, "_segment_metric_frames", 0)), device=self._metric_device, dtype=torch.long)
+        dist.all_reduce(length, op=dist.ReduceOp.MAX)
+        dist.all_reduce(frames, op=dist.ReduceOp.MAX)
+        self._segment_metric_frames = int(frames.item())
+        self._pad_segment_metric_accumulators(int(length.item()))
+        for attr in ("mse_sum", "mse_count", "psnr_sum", "psnr_count", "lpips_sum", "lpips_count"):
+            dist.all_reduce(getattr(self, f"_segment_{attr}"), op=dist.ReduceOp.SUM)
+        self._segment_metrics_synced = True
+
+    def _update_segment_metric_accumulators(
+        self,
+        xs_pred: torch.Tensor,
+        xs_gt: torch.Tensor,
+        valid_mask: torch.Tensor | None,
+        eval_start: int,
+    ) -> None:
+        if self.metric_report_segment <= 0 or xs_pred.numel() == 0:
+            return
+        if not hasattr(self, "_segment_mse_sum"):
+            self._reset_metric_accumulators()
+        self._segment_metrics_synced = False
+
+        frames = int(xs_pred.shape[0])
+        segment = int(self.metric_report_segment)
+        num_segments = (frames + segment - 1) // segment
+        self._pad_segment_metric_accumulators(num_segments)
+        self._segment_metric_frames = max(int(getattr(self, "_segment_metric_frames", 0)), frames)
+        if self._segment_metric_eval_start is None:
+            self._segment_metric_eval_start = int(eval_start)
+
+        pred = torch.clamp(xs_pred.to(self._metric_device).float(), 0.0, 1.0)
+        gt = torch.clamp(xs_gt.to(self._metric_device).float(), 0.0, 1.0)
+        mask = torch.ones(pred.shape[:2], device=self._metric_device, dtype=torch.bool)
+        if valid_mask is not None:
+            mask = valid_mask.to(device=self._metric_device, dtype=torch.bool)
+
+        frame_mse = (pred - gt).square().flatten(start_dim=2).mean(dim=2)
+        frame_psnr = 10.0 * torch.log10(1.0 / frame_mse.clamp_min(torch.finfo(frame_mse.dtype).eps))
+        lpips_batch_size = max(1, int(self.lpips_batch_size))
+
+        for segment_idx in range(num_segments):
+            start = segment_idx * segment
+            end = min(start + segment, frames)
+            segment_mask = mask[start:end]
+            count = int(segment_mask.sum().item())
+            if count == 0:
+                continue
+
+            mask_f = segment_mask.to(dtype=frame_mse.dtype)
+            mse_sum = (frame_mse[start:end] * mask_f).sum()
+            psnr_sum = (frame_psnr[start:end] * mask_f).sum()
+            count_t = torch.tensor(float(count), device=self._metric_device)
+            pred_segment = pred[start:end][segment_mask]
+            gt_segment = gt[start:end][segment_mask]
+
+            self.validation_lpips_model.reset()
+            for batch_start in range(0, count, lpips_batch_size):
+                batch_end = min(batch_start + lpips_batch_size, count)
+                self.validation_lpips_model.update(pred_segment[batch_start:batch_end], gt_segment[batch_start:batch_end])
+            lpips = self.validation_lpips_model.compute().detach().to(self._metric_device)
+            self.validation_lpips_model.reset()
+
+            self._segment_mse_sum[segment_idx] += mse_sum
+            self._segment_mse_count[segment_idx] += count_t
+            self._segment_psnr_sum[segment_idx] += psnr_sum
+            self._segment_psnr_count[segment_idx] += count_t
+            self._segment_lpips_sum[segment_idx] += lpips * count_t
+            self._segment_lpips_count[segment_idx] += count_t
+
+            # Segment mode derives global metrics from this same pass to avoid
+            # running LPIPS once for the whole rollout and again per segment.
+            self._mse_sum += mse_sum
+            self._mse_count += count_t
+            self._psnr_sum += psnr_sum
+            self._psnr_count += count_t
+            self._lpips_sum += lpips * count_t
+            self._lpips_count += count_t
+
     def _log_per_frame_metrics(self) -> None:
         if not hasattr(self, "_frame_mse_sum"):
             return
@@ -1187,6 +1306,71 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 writer.writeheader()
                 writer.writerows(rows)
 
+    def _log_metric_segments(self) -> None:
+        if self.metric_report_segment <= 0 or not hasattr(self, "_segment_mse_sum"):
+            return
+        self._sync_segment_metric_accumulators()
+        valid = self._segment_mse_count > 0
+        if not bool(valid.any()):
+            return
+
+        segment = int(self.metric_report_segment)
+        eval_start = 0 if self._segment_metric_eval_start is None else int(self._segment_metric_eval_start)
+        rows = []
+        log_dict = {}
+        total_generated = int(getattr(self, "_segment_metric_frames", self._segment_mse_count.numel() * segment))
+        for segment_idx in torch.nonzero(valid, as_tuple=False).flatten().tolist():
+            start = int(segment_idx) * segment
+            end = min(start + segment, total_generated)
+            label = f"{start:04d}_{end:04d}"
+            mse = self._mean_or_nan_tensor(self._segment_mse_sum[segment_idx], self._segment_mse_count[segment_idx])
+            psnr = self._mean_or_nan_tensor(self._segment_psnr_sum[segment_idx], self._segment_psnr_count[segment_idx])
+            lpips = self._mean_or_nan_tensor(self._segment_lpips_sum[segment_idx], self._segment_lpips_count[segment_idx])
+            generation_length = end - start
+            log_dict[f"segment/mse_{label}"] = mse
+            log_dict[f"segment/psnr_{label}"] = psnr
+            log_dict[f"segment/lpips_{label}"] = lpips
+            log_dict[f"segment/generation_length_{label}"] = torch.tensor(float(generation_length), device=self._metric_device)
+            rows.append(
+                {
+                    "generation_start": start,
+                    "generation_end": end,
+                    "generation_length": generation_length,
+                    "absolute_start": eval_start + start,
+                    "absolute_end": eval_start + end,
+                    "mse": float(mse.detach().cpu().item()),
+                    "psnr": float(psnr.detach().cpu().item()),
+                    "lpips": float(lpips.detach().cpu().item()),
+                    "count": float(self._segment_mse_count[segment_idx].detach().cpu().item()),
+                }
+            )
+        self.log_dict(log_dict, sync_dist=False)
+        if self._is_global_rank_zero():
+            output_path = self._segment_metric_artifact_dir() / f"segment_metrics_step{int(getattr(self, 'global_step', 0)):08d}.csv"
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            with output_path.open("w", newline="", encoding="utf-8") as handle:
+                writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
+                writer.writeheader()
+                writer.writerows(rows)
+            logger = getattr(self, "logger", None)
+            if logger is not None:
+                import wandb
+
+                columns = list(rows[0].keys())
+                experiment = logger.experiment
+                experiment.log(
+                    {
+                        "segment_metrics/table": wandb.Table(
+                            columns=columns,
+                            data=[[row[column] for column in columns] for row in rows],
+                        ),
+                        "trainer/global_step": int(getattr(self, "global_step", 0)),
+                    }
+                )
+                save_fn = getattr(experiment, "save", None)
+                if callable(save_fn):
+                    save_fn(str(output_path), policy="now")
+
     def _revisit_gap_labels(self) -> list[tuple[str, int, int | None]]:
         edges = [self.revisit_metrics_min_gap]
         edges.extend(edge for edge in sorted(set(self.revisit_metrics_gap_bands)) if edge > self.revisit_metrics_min_gap)
@@ -1435,12 +1619,29 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 writer.writeheader()
                 writer.writerows(rows)
 
-    def _update_metric_accumulators(self, xs_pred: torch.Tensor, xs_gt: torch.Tensor) -> None:
+    def _update_metric_accumulators(
+        self,
+        xs_pred: torch.Tensor,
+        xs_gt: torch.Tensor,
+        valid_mask: torch.Tensor | None = None,
+        eval_start: int = 0,
+    ) -> None:
         if not hasattr(self, "_metric_device"):
             self._reset_metric_accumulators()
 
+        if self.metric_report_segment > 0:
+            self._update_segment_metric_accumulators(xs_pred, xs_gt, valid_mask, eval_start)
+            return
+
         xs_pred_device = xs_pred.to(self._metric_device)
         xs_gt_device = xs_gt.to(self._metric_device)
+        if valid_mask is not None:
+            valid_mask = valid_mask.to(device=self._metric_device, dtype=torch.bool)
+            if not bool(valid_mask.any()):
+                return
+            xs_pred_device = xs_pred_device[valid_mask].unsqueeze(1)
+            xs_gt_device = xs_gt_device[valid_mask].unsqueeze(1)
+
         metric_dict = get_validation_metrics_for_videos(
             xs_pred_device,
             xs_gt_device,
@@ -1488,6 +1689,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
 
         if self.log_per_frame_metrics:
             self._log_per_frame_metrics()
+        if self.metric_report_segment > 0:
+            self._log_metric_segments()
         if self.log_revisit_metrics:
             self._log_revisit_metrics()
 
@@ -1796,18 +1999,12 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                     )
 
                 metric_mask = None
-                if target_mask is None:
-                    self._update_metric_accumulators(xs_pred_decode, xs_decode)
-                else:
+                if target_mask is not None:
                     metric_mask = rearrange(
                         target_mask[:, eval_start:target_length].to(device=xs_pred_decode.device),
                         "b t -> t b",
                     )
-                    if metric_mask.any():
-                        self._update_metric_accumulators(
-                            xs_pred_decode[metric_mask].unsqueeze(1),
-                            xs_decode[metric_mask].unsqueeze(1),
-                        )
+                self._update_metric_accumulators(xs_pred_decode, xs_decode, metric_mask, eval_start)
                 self._update_per_frame_metric_accumulators(xs_pred_decode, xs_decode, metric_mask, eval_start)
                 self._update_revisit_metric_accumulators(
                     xs_pred_decode,

configurations/algorithm/dememwm_base.yaml

@@ -39,6 +39,8 @@ ref_mode: sequential
 focal_length: 0.35
 log_video: false
 save_local: true
+# 0 disables segment metrics; 100 logs generated-frame windows [0,100), [100,200), ...
+metric_report_segment: 0
 require_pose_prediction: false
 
 # training hyperparameters

configurations/experiment/base_pytorch.yaml

@@ -26,6 +26,7 @@ training:
     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``.
+    save_last_k: null # repo extension: if set, keep only the latest k step checkpoints.
     enable_version_counter: False # If this is ``False``, later checkpoint will be overwrite previous ones.
 
 validation:

experiments/exp_base.py

@@ -33,6 +33,41 @@ from huggingface_hub import model_info
 
 torch.set_float32_matmul_precision("high")
 
+
+class LastKStepModelCheckpoint(ModelCheckpoint):
+    """ModelCheckpoint variant that keeps only the latest k step checkpoints."""
+
+    def __init__(self, *args, save_last_k: Optional[int] = None, **kwargs):
+        self._save_last_k = None
+        if save_last_k is not None:
+            self._save_last_k = int(save_last_k)
+            if self._save_last_k < 1:
+                raise ValueError("save_last_k must be a positive integer")
+            # Lightning needs save-all mode here; this subclass prunes by step after each save.
+            kwargs["save_top_k"] = -1
+        super().__init__(*args, **kwargs)
+
+    def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+        super()._save_checkpoint(trainer, filepath)
+        self._prune_old_step_checkpoints(trainer)
+
+    def _prune_old_step_checkpoints(self, trainer: "pl.Trainer") -> None:
+        if self._save_last_k is None or self.dirpath is None:
+            return
+
+        checkpoint_dir = pathlib.Path(self.dirpath)
+        if not checkpoint_dir.exists():
+            return
+
+        step_checkpoints = []
+        for checkpoint_path in checkpoint_dir.glob("*.ckpt"):
+            match = re.search(r"(?:^|[_=-])step[=_]?(\d+)", checkpoint_path.stem)
+            if match:
+                step_checkpoints.append((int(match.group(1)), checkpoint_path.name, str(checkpoint_path)))
+
+        for _, _, checkpoint_path in sorted(step_checkpoints)[:-self._save_last_k]:
+            self._remove_checkpoint(trainer, checkpoint_path)
+
 def is_huggingface_model(path: str) -> bool:
     hf_ckpt = str(path).split('/')
     repo_id = '/'.join(hf_ckpt[:2])
@@ -404,12 +439,17 @@ class BaseLightningExperiment(BaseExperiment):
         if self.logger:
             callbacks.append(LearningRateMonitor("step", True))
         if "checkpointing" in self.cfg.training:
+            checkpointing_cfg = dict(self.cfg.training.checkpointing)
+            save_last_k = checkpointing_cfg.pop("save_last_k", None)
+            checkpoint_cls = LastKStepModelCheckpoint if save_last_k is not None else ModelCheckpoint
+            if save_last_k is not None:
+                checkpointing_cfg["save_last_k"] = save_last_k
             callbacks.append(
-                ModelCheckpoint(
+                checkpoint_cls(
                     pathlib.Path(hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]) / "checkpoints",
                     filename="epoch{epoch}_step{step}",
                     auto_insert_metric_name=False,
-                    **self.cfg.training.checkpointing,
+                    **checkpointing_cfg,
                 )
             )
         return callbacks

train_dememwm_curric_base_recent_8a100_berzelius.sh

@@ -1,93 +0,0 @@
-#!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-base-recent-8a100
-#SBATCH --nodes=1
-#SBATCH --ntasks-per-node=8
-#SBATCH --cpus-per-task=8
-#SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-base-recent-8a100-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-base-recent-8a100-%j.out
-#SBATCH --account=berzelius-2025-436
-#SBATCH --gres=gpu:A100-SXM4-80GB:8
-#SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
-
-set -eo pipefail
-
-module load buildenv-gcccuda/12.1.1-gcc12.3.0
-source $(conda info --base)/etc/profile.d/conda.sh
-
-export PYTHONPATH="./:${PYTHONPATH:-}"
-export HF_HOME=/proj/cvl/users/x_fahkh2/caches
-export TORCH_HOME=/proj/cvl/users/x_fahkh2/caches
-export PIP_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
-export TMPDIR=/proj/cvl/users/x_fahkh2/caches
-export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
-export CUDA_HOME=$CUDA_ROOT
-export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
-export HYDRA_FULL_ERROR=1
-
-git checkout main
-
-RUN_NAME=train_dememwm_curric_base_recent_a2d4r2_route_ntok8_8a100_bs8_global64_400k
-OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
-CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:120000,dataset:{wo_updown:true}},{name:with_updown,until_step:400000,dataset:{wo_updown:false}}]}'
-
-srun python -m main \
-    +name=${RUN_NAME} \
-    +output_dir=$OUTPUT_DIR/ \
-    auto_resume=true \
-    wandb.mode=offline \
-    "experiment.tasks=[training]" \
-    algorithm=dememwm_base \
-    +customized_load=true \
-    +seperate_load=true \
-    +diffusion_model_path=/proj/cvl/users/x_fahkh2/WorldMem_Repro/checkpoints/oasis500m.safetensors \
-    +vae_path=/proj/cvl/users/x_fahkh2/WorldMem_Repro/checkpoints/vit-l-20.safetensors \
-    dataset=video_minecraft_dememwm_latent \
-    dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
-    dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
-    dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
-    dataset.context_length=100 \
-    dataset.single_eval_clip=true \
-    dataset.memory_selection.enabled=true \
-    dataset.memory_selection.max_anchor_frames=2 \
-    dataset.memory_selection.max_dynamic_frames=4 \
-    dataset.memory_selection.max_revisit_frames=2 \
-    dataset.memory_selection.pose_similarity_threshold=0.6 \
-    dataset.memory_selection.training_use_plucker=true \
-    dataset.memory_selection.training_plucker_weight=1.0 \
-    dataset.memory_selection.fov_overlap_threshold=0.6 \
-    dataset.memory_selection.min_total_selected_coverage=0.1 \
-    dataset.memory_selection.local_context_exclusion_frames=8 \
-    dataset.memory_selection.anchor_diverse_selection=true \
-    +dataset.memory_selection.pose_preselect_topk=64 \
-    +dataset.memory_selection.candidate_chunk_size=64 \
-    dataset.memory_selection.dynamic.selection_policy=recent \
-    algorithm.n_tokens=8 \
-    algorithm.context_frames=600 \
-    algorithm.chunk_size=1 \
-    algorithm.warmup_steps=10000 \
-    algorithm.trainability.train_full_dit=true \
-    algorithm.trainability.full_dit_start_step=40000 \
-    algorithm.trainability.lr.memory_modules=4.0e-5 \
-    algorithm.trainability.lr.base_dit=1.0e-5 \
-    algorithm.log_video=true \
-    algorithm.save_local=true \
-    algorithm.diffusion.sampling_timesteps=20 \
-    algorithm.noise_route.anchor=high \
-    algorithm.noise_route.dynamic=low \
-    algorithm.noise_route.revisit=all \
-    algorithm.memory_noise.enabled=true \
-    algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
-    algorithm.memory_noise.revisit_max_fraction=0.25 \
-    "experiment.training.curriculum=${CURRICULUM}" \
-    experiment.training.batch_size=8 \
-    experiment.training.optim.accumulate_grad_batches=1 \
-    experiment.validation.batch_size=1 \
-    experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
-    experiment.training.max_steps=400000

train_dememwm_curric_base_recent_8h200_berzelius.sh

@@ -1,93 +0,0 @@
-#!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-base-recent-8h200
-#SBATCH --nodes=1
-#SBATCH --ntasks-per-node=8
-#SBATCH --cpus-per-task=8
-#SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-base-recent-8h200-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-base-recent-8h200-%j.out
-#SBATCH --account=berzelius-2025-436
-#SBATCH --gres=gpu:8
-#SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
-
-set -eo pipefail
-
-module load buildenv-gcccuda/12.9.1-gcc11
-source $(conda info --base)/etc/profile.d/conda.sh
-
-export PYTHONPATH="./:${PYTHONPATH:-}"
-export HF_HOME=/proj/cvl/users/x_fahkh2/caches
-export TORCH_HOME=/proj/cvl/users/x_fahkh2/caches
-export PIP_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
-export TMPDIR=/proj/cvl/users/x_fahkh2/caches
-export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
-export CUDA_HOME=$CUDA_ROOT
-export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
-export HYDRA_FULL_ERROR=1
-
-git checkout main
-
-RUN_NAME=train_dememwm_curric_base_recent_a2d4r2_route_ntok8_8h200_bs16_global128_200k
-OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
-CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:60000,dataset:{wo_updown:true}},{name:with_updown,until_step:200000,dataset:{wo_updown:false}}]}'
-
-srun python -m main \
-    +name=${RUN_NAME} \
-    +output_dir=$OUTPUT_DIR/ \
-    auto_resume=true \
-    wandb.mode=offline \
-    "experiment.tasks=[training]" \
-    algorithm=dememwm_base \
-    +customized_load=true \
-    +seperate_load=true \
-    +diffusion_model_path=/proj/cvl/users/x_fahkh2/WorldMem_Repro/checkpoints/oasis500m.safetensors \
-    +vae_path=/proj/cvl/users/x_fahkh2/WorldMem_Repro/checkpoints/vit-l-20.safetensors \
-    dataset=video_minecraft_dememwm_latent \
-    dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
-    dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
-    dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
-    dataset.context_length=100 \
-    dataset.single_eval_clip=true \
-    dataset.memory_selection.enabled=true \
-    dataset.memory_selection.max_anchor_frames=2 \
-    dataset.memory_selection.max_dynamic_frames=4 \
-    dataset.memory_selection.max_revisit_frames=2 \
-    dataset.memory_selection.pose_similarity_threshold=0.6 \
-    dataset.memory_selection.training_use_plucker=true \
-    dataset.memory_selection.training_plucker_weight=1.0 \
-    dataset.memory_selection.fov_overlap_threshold=0.6 \
-    dataset.memory_selection.min_total_selected_coverage=0.1 \
-    dataset.memory_selection.local_context_exclusion_frames=8 \
-    dataset.memory_selection.anchor_diverse_selection=true \
-    +dataset.memory_selection.pose_preselect_topk=64 \
-    +dataset.memory_selection.candidate_chunk_size=64 \
-    dataset.memory_selection.dynamic.selection_policy=recent \
-    algorithm.n_tokens=8 \
-    algorithm.context_frames=600 \
-    algorithm.chunk_size=1 \
-    algorithm.warmup_steps=5000 \
-    algorithm.trainability.train_full_dit=true \
-    algorithm.trainability.full_dit_start_step=20000 \
-    algorithm.trainability.lr.memory_modules=8.0e-5 \
-    algorithm.trainability.lr.base_dit=2.0e-5 \
-    algorithm.log_video=true \
-    algorithm.save_local=true \
-    algorithm.diffusion.sampling_timesteps=20 \
-    algorithm.noise_route.anchor=high \
-    algorithm.noise_route.dynamic=low \
-    algorithm.noise_route.revisit=all \
-    algorithm.memory_noise.enabled=true \
-    algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
-    algorithm.memory_noise.revisit_max_fraction=0.25 \
-    "experiment.training.curriculum=${CURRICULUM}" \
-    experiment.training.batch_size=16 \
-    experiment.training.optim.accumulate_grad_batches=1 \
-    experiment.validation.batch_size=1 \
-    experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
-    experiment.training.max_steps=200000

train_dememwm_curric_event_trigger_8a100_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-8a100
+#SBATCH --job-name=dememwm-curric-event-anchorclean-causal-8a100
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-8a100-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-8a100-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-anchorclean-causal-8a100-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-anchorclean-causal-8a100-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:A100-SXM4-80GB:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,12 +23,11 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_a2d4r2_route_ntok8_8a100_bs8_global64_400k
+RUN_NAME=train_dememwm_curric_event_trigger_anchorclean_causal_a2d4r2_route_ntok8_8a100_bs8_global64_400k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
 CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:120000,dataset:{wo_updown:true}},{name:with_updown,until_step:400000,dataset:{wo_updown:false}}]}'
 
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=true \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,6 +66,7 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
     algorithm.warmup_steps=10000 \
     algorithm.trainability.train_full_dit=true \
@@ -80,14 +81,15 @@ srun python -m main \
     algorithm.noise_route.revisit=all \
     algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=8 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
     experiment.training.max_steps=400000

train_dememwm_curric_event_trigger_8h200_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-8h200
+#SBATCH --job-name=dememwm-curric-event-anchorclean-causal-8h200
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-8h200-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-8h200-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-anchorclean-causal-8h200-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-anchorclean-causal-8h200-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,14 +23,13 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_a2d4r2_route_ntok8_8h200_bs16_global128_200k
+RUN_NAME=train_dememwm_curric_event_trigger_anchorclean_causal_a2d4r2_route_ntok8_8h200_bs16_global128_150k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
-CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:60000,dataset:{wo_updown:true}},{name:with_updown,until_step:200000,dataset:{wo_updown:false}}]}'
+CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:30000,dataset:{wo_updown:true}},{name:with_updown,until_step:150000,dataset:{wo_updown:false}}]}'
 
 srun python -m main \
     +name=${RUN_NAME} \
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=true \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,10 +66,11 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
-    algorithm.warmup_steps=5000 \
+    algorithm.warmup_steps=1000 \
     algorithm.trainability.train_full_dit=true \
-    algorithm.trainability.full_dit_start_step=20000 \
+    algorithm.trainability.full_dit_start_step=10000 \
     algorithm.trainability.lr.memory_modules=8.0e-5 \
     algorithm.trainability.lr.base_dit=2.0e-5 \
     algorithm.log_video=true \
@@ -80,14 +81,15 @@ srun python -m main \
     algorithm.noise_route.revisit=all \
     algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=16 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
-    experiment.training.max_steps=200000
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
+    experiment.training.max_steps=150000

train_dememwm_curric_event_trigger_no_memory_noise_8a100_berzelius.sh → train_dememwm_curric_event_trigger_anchorclean_noncausal_8a100_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-nomemnoise-8a100
+#SBATCH --job-name=dememwm-curric-event-anchorclean-noncausal-8a100
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-nomemnoise-8a100-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-nomemnoise-8a100-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-anchorclean-noncausal-8a100-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-anchorclean-noncausal-8a100-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:A100-SXM4-80GB:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,12 +23,11 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_no_memory_noise_a2d4r2_route_ntok8_8a100_bs8_global64_400k
+RUN_NAME=train_dememwm_curric_event_trigger_anchorclean_noncausal_a2d4r2_route_ntok8_8a100_bs8_global64_400k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
 CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:120000,dataset:{wo_updown:true}},{name:with_updown,until_step:400000,dataset:{wo_updown:false}}]}'
 
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=false \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,6 +66,7 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
     algorithm.warmup_steps=10000 \
     algorithm.trainability.train_full_dit=true \
@@ -78,16 +79,17 @@ srun python -m main \
     algorithm.noise_route.anchor=high \
     algorithm.noise_route.dynamic=low \
     algorithm.noise_route.revisit=all \
-    algorithm.memory_noise.enabled=false \
+    algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=8 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
     experiment.training.max_steps=400000

train_dememwm_curric_event_trigger_no_memory_noise_8h200_berzelius.sh → train_dememwm_curric_event_trigger_anchorclean_noncausal_8h200_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-nomemnoise-8h200
+#SBATCH --job-name=dememwm-curric-event-anchorclean-noncausal-8h200
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-nomemnoise-8h200-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-nomemnoise-8h200-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-anchorclean-noncausal-8h200-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-anchorclean-noncausal-8h200-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,14 +23,13 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_no_memory_noise_a2d4r2_route_ntok8_8h200_bs16_global128_200k
+RUN_NAME=train_dememwm_curric_event_trigger_anchorclean_noncausal_a2d4r2_route_ntok8_8h200_bs16_global128_150k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
-CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:60000,dataset:{wo_updown:true}},{name:with_updown,until_step:200000,dataset:{wo_updown:false}}]}'
+CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:30000,dataset:{wo_updown:true}},{name:with_updown,until_step:150000,dataset:{wo_updown:false}}]}'
 
 srun python -m main \
     +name=${RUN_NAME} \
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=false \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,10 +66,11 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
-    algorithm.warmup_steps=5000 \
+    algorithm.warmup_steps=1000 \
     algorithm.trainability.train_full_dit=true \
-    algorithm.trainability.full_dit_start_step=20000 \
+    algorithm.trainability.full_dit_start_step=10000 \
     algorithm.trainability.lr.memory_modules=8.0e-5 \
     algorithm.trainability.lr.base_dit=2.0e-5 \
     algorithm.log_video=true \
@@ -78,16 +79,17 @@ srun python -m main \
     algorithm.noise_route.anchor=high \
     algorithm.noise_route.dynamic=low \
     algorithm.noise_route.revisit=all \
-    algorithm.memory_noise.enabled=false \
+    algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=16 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
-    experiment.training.max_steps=200000
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
+    experiment.training.max_steps=150000

train_dememwm_curric_event_trigger_no_route_8a100_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-noroute-8a100
+#SBATCH --job-name=dememwm-curric-event-noroute-anchorclean-8a100
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-noroute-8a100-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-noroute-8a100-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-noroute-anchorclean-8a100-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-noroute-anchorclean-8a100-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:A100-SXM4-80GB:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,12 +23,11 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_no_route_a2d4r2_ntok8_8a100_bs8_global64_400k
+RUN_NAME=train_dememwm_curric_event_trigger_no_route_anchorclean_a2d4r2_ntok8_8a100_bs8_global64_400k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
 CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:120000,dataset:{wo_updown:true}},{name:with_updown,until_step:400000,dataset:{wo_updown:false}}]}'
 
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=true \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,6 +66,7 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
     algorithm.warmup_steps=10000 \
     algorithm.trainability.train_full_dit=true \
@@ -80,14 +81,15 @@ srun python -m main \
     algorithm.noise_route.revisit=all \
     algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=8 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
     experiment.training.max_steps=400000

train_dememwm_curric_event_trigger_no_route_8h200_berzelius.sh

@@ -1,11 +1,11 @@
 #!/usr/bin/env bash
-#SBATCH --job-name=dememwm-curric-event-noroute-8h200
+#SBATCH --job-name=dememwm-curric-event-noroute-anchorclean-8h200
 #SBATCH --nodes=1
 #SBATCH --ntasks-per-node=8
 #SBATCH --cpus-per-task=8
 #SBATCH --time=72:00:00
-#SBATCH --output=slurm_logs/dememwm-curric-event-noroute-8h200-%j.out
-#SBATCH --error=slurm_logs/dememwm-curric-event-noroute-8h200-%j.out
+#SBATCH --output=slurm_logs/dememwm-curric-event-noroute-anchorclean-8h200-%j.out
+#SBATCH --error=slurm_logs/dememwm-curric-event-noroute-anchorclean-8h200-%j.out
 #SBATCH --account=berzelius-2025-436
 #SBATCH --gres=gpu:8
 #SBATCH --chdir=/proj/cvl/users/x_fahkh2/DeMemWM
@@ -23,14 +23,13 @@ export TMPDIR=/proj/cvl/users/x_fahkh2/caches
 export TRITON_CACHE_DIR=/proj/cvl/users/x_fahkh2/caches
 export CUDA_HOME=$CUDA_ROOT
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export WANDB_DISABLED=true
 export HYDRA_FULL_ERROR=1
 
 git checkout main
 
-RUN_NAME=train_dememwm_curric_event_trigger_no_route_a2d4r2_ntok8_8h200_bs16_global128_200k
+RUN_NAME=train_dememwm_curric_event_trigger_no_route_anchorclean_a2d4r2_ntok8_8h200_bs16_global128_150k
 OUTPUT_DIR=/proj/cvl/users/x_fahkh2/DeMemWM/results/${RUN_NAME}
-CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:60000,dataset:{wo_updown:true}},{name:with_updown,until_step:200000,dataset:{wo_updown:false}}]}'
+CURRICULUM='{enabled:true,stages:[{name:no_updown,until_step:30000,dataset:{wo_updown:true}},{name:with_updown,until_step:150000,dataset:{wo_updown:false}}]}'
 
 srun python -m main \
     +name=${RUN_NAME} \
@@ -47,10 +46,11 @@ srun python -m main \
     dataset.save_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft \
     dataset.precomputed_feature_dir=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft/vae_features \
     dataset.n_frames=8 \
-    dataset.n_frames_valid=700 \
+    dataset.n_frames_valid=1100 \
     dataset.context_length=100 \
     dataset.single_eval_clip=true \
     dataset.memory_selection.enabled=true \
+    dataset.memory_selection.causal=true \
     dataset.memory_selection.max_anchor_frames=2 \
     dataset.memory_selection.max_dynamic_frames=4 \
     dataset.memory_selection.max_revisit_frames=2 \
@@ -66,10 +66,11 @@ srun python -m main \
     dataset.memory_selection.dynamic.selection_policy=event_triggered \
     algorithm.n_tokens=8 \
     algorithm.context_frames=600 \
+    algorithm.metric_report_segment=100 \
     algorithm.chunk_size=1 \
-    algorithm.warmup_steps=5000 \
+    algorithm.warmup_steps=1000 \
     algorithm.trainability.train_full_dit=true \
-    algorithm.trainability.full_dit_start_step=20000 \
+    algorithm.trainability.full_dit_start_step=10000 \
     algorithm.trainability.lr.memory_modules=8.0e-5 \
     algorithm.trainability.lr.base_dit=2.0e-5 \
     algorithm.log_video=true \
@@ -80,14 +81,15 @@ srun python -m main \
     algorithm.noise_route.revisit=all \
     algorithm.memory_noise.enabled=true \
     algorithm.memory_noise.mode=random_cleaner_fraction \
-    algorithm.memory_noise.anchor_max_fraction=0.25 \
-    algorithm.memory_noise.dynamic_max_fraction=0.50 \
+    algorithm.memory_noise.anchor_max_fraction=0.0 \
+    algorithm.memory_noise.dynamic_max_fraction=0.25 \
     algorithm.memory_noise.revisit_max_fraction=0.25 \
     "experiment.training.curriculum=${CURRICULUM}" \
     experiment.training.batch_size=16 \
     experiment.training.optim.accumulate_grad_batches=1 \
     experiment.validation.batch_size=1 \
     experiment.validation.limit_batch=1 \
-    experiment.training.checkpointing.every_n_train_steps=5000 \
-    experiment.validation.val_every_n_step=5000 \
-    experiment.training.max_steps=200000
+    experiment.training.checkpointing.every_n_train_steps=2000 \
+    experiment.training.checkpointing.save_last_k=10 \
+    experiment.validation.val_every_n_step=2000 \
+    experiment.training.max_steps=150000