Add DeMemWM long-horizon eval metrics

algorithms/dememwm/df_video.py

@@ -1,7 +1,9 @@
 from collections.abc import Mapping
+import csv
 import os
 import random
 import time
+from pathlib import Path
 import numpy as np
 import torch
 import torch.distributed as dist
@@ -24,6 +26,7 @@ from datasets.video.memory_selection import (
     _select_revisit,
 )
 from utils.logging_utils import log_video, get_validation_metrics_for_videos
+from algorithms.worldmem.revisit_consistency import RevisitPair, compute_frame_metric_rows
 from .df_base import DiffusionForcingBase
 from .models.vae import VAE_models
 from .models.diffusion import Diffusion
@@ -670,6 +673,19 @@ 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.log_per_frame_metrics = bool(getattr(cfg, "log_per_frame_metrics", False))
+        self.log_revisit_metrics = bool(getattr(cfg, "log_revisit_metrics", False))
+        self.revisit_metrics_min_gap = int(getattr(cfg, "revisit_metrics_min_gap", 8))
+        self.revisit_metrics_gap_bands = tuple(int(x) for x in getattr(cfg, "revisit_metrics_gap_bands", (32, 128)))
+        self.revisit_metrics_self_crop_fraction = float(getattr(cfg, "revisit_metrics_self_crop_fraction", 0.5))
+        memory_selection_cfg = _cfg_get(cfg, "memory_selection", {})
+        self.revisit_metrics_fov_overlap_threshold = float(
+            getattr(
+                cfg,
+                "revisit_metrics_fov_overlap_threshold",
+                _cfg_get(memory_selection_cfg, "fov_overlap_threshold", 0.6),
+            )
+        )
 
         super().__init__(cfg)
             
@@ -864,6 +880,372 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         self._lpips_sum = torch.tensor(0.0, device=self._metric_device)
         self._lpips_count = torch.tensor(0.0, device=self._metric_device)
 
+        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.log_revisit_metrics:
+            self._reset_revisit_metric_accumulators()
+
+    def _eval_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 / "metrics"
+
+    def _is_global_rank_zero(self) -> bool:
+        return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0
+
+    def _mean_or_nan_tensor(self, total: torch.Tensor, count: torch.Tensor) -> torch.Tensor:
+        return torch.where(count > 0, total / count.clamp_min(1.0), torch.full_like(total, float("nan")))
+
+    def _pad_frame_metric_accumulators(self, length: int) -> None:
+        current = int(self._frame_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"_frame_{attr}"
+            setattr(self, name, torch.cat([getattr(self, name), torch.zeros(pad, device=self._metric_device)]))
+
+    def _sync_frame_metric_accumulators(self) -> None:
+        if not (dist.is_available() and dist.is_initialized()):
+            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)
+
+    def _update_per_frame_metric_accumulators(
+        self,
+        xs_pred: torch.Tensor,
+        xs_gt: torch.Tensor,
+        valid_mask: torch.Tensor | None,
+        eval_start: int,
+    ) -> None:
+        if not self.log_per_frame_metrics or xs_pred.numel() == 0:
+            return
+        if not hasattr(self, "_frame_mse_sum"):
+            self._reset_metric_accumulators()
+
+        frames = int(xs_pred.shape[0])
+        self._pad_frame_metric_accumulators(frames)
+        if self._frame_metric_eval_start is None:
+            self._frame_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))
+        counts = mask.sum(dim=1).to(dtype=frame_mse.dtype)
+        mask_f = mask.to(dtype=frame_mse.dtype)
+        self._frame_mse_sum[:frames] += (frame_mse * mask_f).sum(dim=1)
+        self._frame_mse_count[:frames] += counts
+        self._frame_psnr_sum[:frames] += (frame_psnr * mask_f).sum(dim=1)
+        self._frame_psnr_count[:frames] += counts
+
+        with torch.no_grad():
+            for frame_idx in range(frames):
+                frame_mask = mask[frame_idx]
+                count = int(frame_mask.sum().item())
+                if count == 0:
+                    continue
+                self.validation_lpips_model.reset()
+                self.validation_lpips_model.update(pred[frame_idx, frame_mask], gt[frame_idx, frame_mask])
+                value = self.validation_lpips_model.compute().detach().to(self._metric_device)
+                self.validation_lpips_model.reset()
+                self._frame_lpips_sum[frame_idx] += value * count
+                self._frame_lpips_count[frame_idx] += count
+
+    def _log_per_frame_metrics(self) -> None:
+        if not hasattr(self, "_frame_mse_sum"):
+            return
+        self._sync_frame_metric_accumulators()
+        valid = self._frame_mse_count > 0
+        if not bool(valid.any()):
+            return
+
+        eval_start = 0 if self._frame_metric_eval_start is None else int(self._frame_metric_eval_start)
+        rows = []
+        log_dict = {}
+        for frame_idx in torch.nonzero(valid, as_tuple=False).flatten().tolist():
+            mse = self._frame_mse_sum[frame_idx] / self._frame_mse_count[frame_idx].clamp_min(1.0)
+            psnr = self._frame_psnr_sum[frame_idx] / self._frame_psnr_count[frame_idx].clamp_min(1.0)
+            lpips = self._frame_lpips_sum[frame_idx] / self._frame_lpips_count[frame_idx].clamp_min(1.0)
+            generated_frame = int(frame_idx) + 1
+            log_dict[f"per_frame/mse_{generated_frame:04d}"] = mse
+            log_dict[f"per_frame/psnr_{generated_frame:04d}"] = psnr
+            log_dict[f"per_frame/lpips_{generated_frame:04d}"] = lpips
+            rows.append(
+                {
+                    "generated_frame": generated_frame,
+                    "absolute_frame": eval_start + int(frame_idx),
+                    "mse": float(mse.detach().cpu().item()),
+                    "psnr": float(psnr.detach().cpu().item()),
+                    "lpips": float(lpips.detach().cpu().item()),
+                    "count": float(self._frame_mse_count[frame_idx].detach().cpu().item()),
+                }
+            )
+        self.log_dict(log_dict, sync_dist=False)
+        if self._is_global_rank_zero():
+            output_path = self._eval_artifact_dir() / f"per_frame_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)
+
+    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)
+        edges.append(None)
+        return [
+            (f"{int(lo)}_{'inf' if hi is None else int(hi)}", int(lo), None if hi is None else int(hi))
+            for lo, hi in zip(edges[:-1], edges[1:])
+        ]
+
+    def _reset_revisit_metric_accumulators(self) -> None:
+        names = (
+            "num_pairs",
+            "num_revisit_targets",
+            "num_revisit_frames",
+            "num_non_revisit_frames",
+            "revisit_lpips_sum",
+            "revisit_lpips_count",
+            "non_revisit_lpips_sum",
+            "non_revisit_lpips_count",
+            "revisit_psnr_sum",
+            "revisit_psnr_count",
+            "non_revisit_psnr_sum",
+            "non_revisit_psnr_count",
+            "self_lpips_sum",
+            "self_lpips_count",
+            "self_psnr_sum",
+            "self_psnr_count",
+        )
+        for name in names:
+            setattr(self, f"_revisit_{name}", torch.tensor(0.0, device=self._metric_device))
+        self._revisit_gap_stats = {
+            label: {
+                "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_psnr_sum": torch.tensor(0.0, device=self._metric_device),
+                "self_psnr_count": torch.tensor(0.0, device=self._metric_device),
+                "revisit_lpips_sum": torch.tensor(0.0, device=self._metric_device),
+                "revisit_lpips_count": torch.tensor(0.0, device=self._metric_device),
+            }
+            for label, _, _ in self._revisit_gap_labels()
+        }
+
+    def _add_revisit_value(self, sum_attr: str, count_attr: str, value) -> None:
+        try:
+            value = float(value)
+        except (TypeError, ValueError):
+            return
+        if not np.isfinite(value):
+            return
+        getattr(self, sum_attr).add_(value)
+        getattr(self, count_attr).add_(1.0)
+
+    def _add_revisit_value_for_dict(self, stats: dict[str, torch.Tensor], name: str, value) -> None:
+        try:
+            value = float(value)
+        except (TypeError, ValueError):
+            return
+        if not np.isfinite(value):
+            return
+        stats[f"{name}_sum"] += value
+        stats[f"{name}_count"] += 1.0
+
+    def _mine_revisit_pairs_with_validation_selector(
+        self,
+        poses: torch.Tensor,
+        frame_indices: torch.Tensor,
+        memory_selection_cfg,
+        eval_start: int,
+        split: str,
+        clip_id: str,
+    ) -> list[RevisitPair]:
+        poses_np = poses.detach().cpu().numpy()
+        frame_indices_cpu = frame_indices.detach().cpu()
+        excluded = np.empty((0,), dtype=np.int64)
+        pairs: list[RevisitPair] = []
+        for target_index in range(int(eval_start), int(poses.shape[0])):
+            selected = _select_revisit(
+                poses_np,
+                np.asarray([target_index], dtype=np.int64),
+                memory_selection_cfg,
+                1,
+                excluded,
+                split,
+            )
+            if len(selected) == 0:
+                continue
+            source_index = int(selected[0])
+            source_frame = int(frame_indices_cpu[source_index].item())
+            target_frame = int(frame_indices_cpu[target_index].item())
+            if source_index >= target_index:
+                continue
+            source_pose = poses[source_index].detach().float().cpu()
+            target_pose = poses[target_index].detach().float().cpu()
+            position_distance = float(torch.linalg.vector_norm(target_pose[:3] - source_pose[:3]).item())
+            yaw_delta = float(abs(((target_pose[4] - source_pose[4] + 180.0) % 360.0) - 180.0).item())
+            pitch_delta = float(abs(((target_pose[3] - source_pose[3] + 180.0) % 360.0) - 180.0).item())
+            pairs.append(
+                RevisitPair(
+                    clip_id=str(clip_id),
+                    source_index=source_index,
+                    target_index=target_index,
+                    source_frame=source_frame,
+                    target_frame=target_frame,
+                    gap=int(target_frame - source_frame),
+                    fov_overlap=1.0,
+                    plucker_overlap=0.0,
+                    position_distance=position_distance,
+                    yaw_delta_deg=yaw_delta,
+                    pitch_delta_deg=pitch_delta,
+                    positional=bool(position_distance <= 2.0),
+                )
+            )
+        return pairs
+
+    def _update_revisit_metric_accumulators(
+        self,
+        xs_pred: torch.Tensor,
+        xs_gt: torch.Tensor,
+        target_poses: torch.Tensor,
+        target_frame_indices: torch.Tensor,
+        valid_mask: torch.Tensor | None,
+        eval_start: int,
+        batch_idx: int,
+        namespace: str,
+    ) -> None:
+        if not self.log_revisit_metrics or xs_pred.numel() == 0:
+            return
+        if not hasattr(self, "_revisit_num_pairs"):
+            self._reset_revisit_metric_accumulators()
+
+        valid_mask_cpu = None if valid_mask is None else valid_mask.detach().cpu()
+        batch_size = int(xs_pred.shape[1])
+        for batch_i in range(batch_size):
+            clip_id = f"batch{batch_idx:06d}_sample{batch_i:02d}"
+            poses = target_poses[:, batch_i, :5].detach()
+            frame_indices = target_frame_indices[:, batch_i].detach()
+            pairs = self._mine_revisit_pairs_with_validation_selector(
+                poses,
+                frame_indices,
+                _cfg_get(self.cfg, "memory_selection", {}),
+                eval_start,
+                namespace,
+                clip_id,
+            )
+            self._revisit_num_pairs += float(len(pairs))
+            self._revisit_num_revisit_targets += float(len({int(pair.target_index) for pair in pairs}))
+            rows = compute_frame_metric_rows(
+                branch="validation",
+                clip_id=clip_id,
+                pred=xs_pred[:, batch_i],
+                gt=xs_gt[:, batch_i],
+                poses=poses,
+                pairs=pairs,
+                context_frames=int(eval_start),
+                lpips_model=self.validation_lpips_model,
+                self_crop_fraction=self.revisit_metrics_self_crop_fraction,
+                compute_plucker_similarity=False,
+            )
+            for row in rows:
+                output_index = int(row["output_index"])
+                if valid_mask_cpu is not None and not bool(valid_mask_cpu[output_index, batch_i].item()):
+                    continue
+                if bool(row["is_revisit"]):
+                    self._revisit_num_revisit_frames += 1.0
+                    self._add_revisit_value("_revisit_revisit_lpips_sum", "_revisit_revisit_lpips_count", row["gt_lpips"])
+                    self._add_revisit_value("_revisit_revisit_psnr_sum", "_revisit_revisit_psnr_count", row["gt_psnr"])
+                    self._add_revisit_value("_revisit_self_lpips_sum", "_revisit_self_lpips_count", row["self_lpips"])
+                    self._add_revisit_value("_revisit_self_psnr_sum", "_revisit_self_psnr_count", row["self_psnr"])
+                    try:
+                        gap = int(row["gap"])
+                    except (TypeError, ValueError):
+                        gap = None
+                    if gap is not None:
+                        for label, lo, hi in self._revisit_gap_labels():
+                            if gap >= lo and (hi is None or gap < hi):
+                                stats = self._revisit_gap_stats[label]
+                                stats["count"] += 1.0
+                                self._add_revisit_value_for_dict(stats, "revisit_lpips", row["gt_lpips"])
+                                self._add_revisit_value_for_dict(stats, "self_lpips", row["self_lpips"])
+                                self._add_revisit_value_for_dict(stats, "self_psnr", row["self_psnr"])
+                                break
+                else:
+                    self._revisit_num_non_revisit_frames += 1.0
+                    self._add_revisit_value("_revisit_non_revisit_lpips_sum", "_revisit_non_revisit_lpips_count", row["gt_lpips"])
+                    self._add_revisit_value("_revisit_non_revisit_psnr_sum", "_revisit_non_revisit_psnr_count", row["gt_psnr"])
+
+    def _sync_revisit_metric_accumulators(self) -> None:
+        if not (dist.is_available() and dist.is_initialized()):
+            return
+        attrs = (
+            "num_pairs",
+            "num_revisit_targets",
+            "num_revisit_frames",
+            "num_non_revisit_frames",
+            "revisit_lpips_sum",
+            "revisit_lpips_count",
+            "non_revisit_lpips_sum",
+            "non_revisit_lpips_count",
+            "revisit_psnr_sum",
+            "revisit_psnr_count",
+            "non_revisit_psnr_sum",
+            "non_revisit_psnr_count",
+            "self_lpips_sum",
+            "self_lpips_count",
+            "self_psnr_sum",
+            "self_psnr_count",
+        )
+        for attr in attrs:
+            dist.all_reduce(getattr(self, f"_revisit_{attr}"), op=dist.ReduceOp.SUM)
+        for stats in self._revisit_gap_stats.values():
+            for value in stats.values():
+                dist.all_reduce(value, op=dist.ReduceOp.SUM)
+
+    def _log_revisit_metrics(self) -> None:
+        if not hasattr(self, "_revisit_num_pairs"):
+            return
+        self._sync_revisit_metric_accumulators()
+        revisit_lpips = self._mean_or_nan_tensor(self._revisit_revisit_lpips_sum, self._revisit_revisit_lpips_count)
+        non_revisit_lpips = self._mean_or_nan_tensor(self._revisit_non_revisit_lpips_sum, self._revisit_non_revisit_lpips_count)
+        log_dict = {
+            "revisit_metrics/num_pairs": self._revisit_num_pairs,
+            "revisit_metrics/num_revisit_targets": self._revisit_num_revisit_targets,
+            "revisit_metrics/num_revisit_frames": self._revisit_num_revisit_frames,
+            "revisit_metrics/num_non_revisit_frames": self._revisit_num_non_revisit_frames,
+            "revisit_metrics/revisit_lpips": revisit_lpips,
+            "revisit_metrics/non_revisit_lpips": non_revisit_lpips,
+            "revisit_metrics/revisit_psnr": self._mean_or_nan_tensor(self._revisit_revisit_psnr_sum, self._revisit_revisit_psnr_count),
+            "revisit_metrics/non_revisit_psnr": self._mean_or_nan_tensor(self._revisit_non_revisit_psnr_sum, self._revisit_non_revisit_psnr_count),
+            "revisit_metrics/self_consistency_lpips": self._mean_or_nan_tensor(self._revisit_self_lpips_sum, self._revisit_self_lpips_count),
+            "revisit_metrics/self_consistency_psnr": self._mean_or_nan_tensor(self._revisit_self_psnr_sum, self._revisit_self_psnr_count),
+            "revisit_metrics/revisit_minus_non_revisit_lpips": revisit_lpips - non_revisit_lpips,
+        }
+        for label, stats in self._revisit_gap_stats.items():
+            log_dict[f"revisit_metrics/num_gap_{label}"] = stats["count"]
+            log_dict[f"revisit_metrics/self_consistency_lpips_gap_{label}"] = self._mean_or_nan_tensor(stats["self_lpips_sum"], stats["self_lpips_count"])
+            log_dict[f"revisit_metrics/self_consistency_psnr_gap_{label}"] = self._mean_or_nan_tensor(stats["self_psnr_sum"], stats["self_psnr_count"])
+            log_dict[f"revisit_metrics/revisit_lpips_gap_{label}"] = self._mean_or_nan_tensor(stats["revisit_lpips_sum"], stats["revisit_lpips_count"])
+        self.log_dict(log_dict, sync_dist=False)
+        if self._is_global_rank_zero():
+            rows = [{"metric": key, "value": float(value.detach().cpu().item())} for key, value in log_dict.items()]
+            output_path = self._eval_artifact_dir() / f"revisit_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=["metric", "value"])
+                writer.writeheader()
+                writer.writerows(rows)
+
     def _update_metric_accumulators(self, xs_pred: torch.Tensor, xs_gt: torch.Tensor) -> None:
         if not hasattr(self, "_metric_device"):
             self._reset_metric_accumulators()
@@ -915,6 +1297,11 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                 sync_dist=False,
             )
 
+        if self.log_per_frame_metrics:
+            self._log_per_frame_metrics()
+        if self.log_revisit_metrics:
+            self._log_revisit_metrics()
+
         self.validation_step_outputs.clear()
 
     def _preprocess_batch(self, batch):
@@ -1185,24 +1572,37 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                         xs_decode,
                         step=getattr(self, "global_step", 0),
                         namespace=namespace + "_vis",
+                        prefix=f"batch{batch_idx:06d}",
                         context_frames=self.context_frames,
                         logger=self.logger.experiment,
                         save_local=self.save_local,
                         local_save_dir=self.local_save_dir,
                     )
 
+                metric_mask = None
                 if target_mask is None:
                     self._update_metric_accumulators(xs_pred_decode, xs_decode)
                 else:
-                    valid_mask = rearrange(
+                    metric_mask = rearrange(
                         target_mask[:, eval_start:target_length].to(device=xs_pred_decode.device),
                         "b t -> t b",
                     )
-                    if valid_mask.any():
+                    if metric_mask.any():
                         self._update_metric_accumulators(
-                            xs_pred_decode[valid_mask].unsqueeze(1),
-                            xs_decode[valid_mask].unsqueeze(1),
+                            xs_pred_decode[metric_mask].unsqueeze(1),
+                            xs_decode[metric_mask].unsqueeze(1),
                         )
+                self._update_per_frame_metric_accumulators(xs_pred_decode, xs_decode, metric_mask, eval_start)
+                self._update_revisit_metric_accumulators(
+                    xs_pred_decode,
+                    xs_decode,
+                    target_poses,
+                    target_frame_indices,
+                    metric_mask,
+                    eval_start,
+                    batch_idx,
+                    namespace,
+                )
             return latent_loss
 
     def test_step(self, batch, batch_idx) -> STEP_OUTPUT: