| import pytest |
| import torch |
| from dememwm_import_helper import install_dememwm_namespace |
|
|
| install_dememwm_namespace() |
| from algorithms.worldmem.dememwm.labels import RevisitCandidateLabel, plucker_overlap |
| from algorithms.worldmem.dememwm.retrieval import _select_greedy_coverage, deterministic_revisit_retrieval |
| from algorithms.worldmem.dememwm.types import MemoryRecord, MemorySourceType |
|
|
|
|
| def rec(frame, value, generated=False, pose=None, action=None, video_id="v0", chunk_id=None): |
| metadata = {"video_id": video_id} |
| if action is not None: |
| metadata["action"] = torch.tensor(action, dtype=torch.float32) |
| return MemoryRecord( |
| tokens=torch.full((2, 4), float(value)), |
| mask=torch.ones(2, dtype=torch.bool), |
| source_start=frame, |
| source_end=frame + 1, |
| frame_indices=torch.tensor([frame]), |
| pose=None if pose is None else torch.tensor(pose, dtype=torch.float32), |
| source_type=MemorySourceType.REVISIT, |
| is_generated=generated, |
| chunk_id=chunk_id or f"c{frame}", |
| metadata=metadata, |
| ) |
|
|
|
|
| def candidate_label(chunk_id, frame, fov, plucker, coverage_mask): |
| return RevisitCandidateLabel( |
| record=rec(frame, frame, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id=chunk_id), |
| valid=True, |
| gap_valid=True, |
| gap_to_target=10 - int(frame), |
| fov_overlap=float(fov), |
| plucker_overlap=float(plucker), |
| primary_overlap=float(fov), |
| coverage_mask=torch.tensor(coverage_mask, dtype=torch.bool), |
| reject_reasons=(), |
| ) |
|
|
|
|
| def test_plucker_cannot_outrank_higher_incremental_fov_gain(): |
| low_fov_high_plucker = candidate_label( |
| "low_fov_high_plucker", 0, 0.11, 1.0, [True] * 11 + [False] * 89 |
| ) |
| high_fov_low_plucker = candidate_label( |
| "high_fov_low_plucker", 1, 0.20, 0.0, [True] * 20 + [False] * 80 |
| ) |
|
|
| selected, scores, gains = _select_greedy_coverage( |
| [low_fov_high_plucker, high_fov_low_plucker], topk=1, plucker_weight=0.10 |
| ) |
|
|
| assert selected[0].record.chunk_id == "high_fov_low_plucker" |
| assert abs(scores[0] - gains[0]) < 1e-6 |
| assert abs(gains[0] - 0.20) < 1e-6 |
|
|
|
|
| def test_plucker_breaks_ties_after_fov_gain_and_overlap(): |
| low_plucker = candidate_label("low_plucker", 0, 0.20, 0.1, [True] * 20 + [False] * 80) |
| high_plucker = candidate_label("high_plucker", 0, 0.20, 0.9, [True] * 20 + [False] * 80) |
|
|
| selected, _, _ = _select_greedy_coverage([low_plucker, high_plucker], topk=1, plucker_weight=0.10) |
|
|
| assert selected[0].record.chunk_id == "high_plucker" |
|
|
|
|
| def test_plucker_overlap_handles_cuda_autocast_mixed_precision(): |
| if not torch.cuda.is_available(): |
| return |
| source_pose = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0]], device="cuda", dtype=torch.float32) |
| target_pose = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0], device="cuda", dtype=torch.float32) |
| with torch.autocast(device_type="cuda", dtype=torch.float16): |
| overlap = plucker_overlap(source_pose, target_pose) |
| assert overlap is not None |
| assert overlap > 0.0 |
|
|
|
|
| def test_revisit_candidates_require_causal_c_short_gap(): |
| records = [ |
| rec(1, 1, pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| rec(2, 2, pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| rec(9, 9, pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| ] |
| result = deterministic_revisit_retrieval( |
| records, |
| target_frame=6, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| topk=5, |
| exclude_local_context_frames=4, |
| ) |
| assert [r.max_source_frame for r in result.records] == [1] |
| assert result.diagnostics["revisit_candidate_frame_count"] == 2 |
| assert result.diagnostics["revisit_candidate_count"] == 2 |
| assert result.diagnostics["valid_revisit_frame_count"] == 1 |
| assert result.diagnostics["valid_revisit_count"] == 1 |
| assert result.diagnostics["valid_candidate_label_count"] == 1 |
| assert result.diagnostics["revisit_min_gap_to_target"] == 5 |
| assert result.diagnostics["revisit_vectorized_frame_scorer_used"] == 1 |
|
|
|
|
| def test_revisit_abstains_when_no_valid_candidate(): |
| result = deterministic_revisit_retrieval([rec(2, 2), rec(3, 3)], target_frame=6, topk=2, exclude_local_context_frames=4) |
| assert result.records == [] |
| assert result.diagnostics["abstained"] is True |
| assert result.diagnostics["valid_revisit_mask"] == 0 |
| assert result.diagnostics["no_valid_revisit_count"] == 1 |
|
|
|
|
| def test_revisit_retrieval_rejects_non_vectorized_inputs(): |
| with pytest.raises(ValueError, match="target_pose"): |
| deterministic_revisit_retrieval( |
| [rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 0.0])], |
| target_frame=10, |
| exclude_local_context_frames=4, |
| ) |
|
|
| chunk_record = MemoryRecord( |
| tokens=torch.zeros((2, 4)), |
| mask=torch.ones(2, dtype=torch.bool), |
| source_start=0, |
| source_end=2, |
| frame_indices=torch.tensor([0, 1]), |
| pose=torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]), |
| source_type=MemorySourceType.REVISIT, |
| is_generated=False, |
| chunk_id="chunk", |
| ) |
| with pytest.raises(ValueError, match="frame-level records"): |
| deterministic_revisit_retrieval( |
| [chunk_record], |
| target_frame=10, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| exclude_local_context_frames=4, |
| ) |
|
|
|
|
| def test_fov_threshold_filters_candidates_without_action(): |
| records = [ |
| rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| rec(1, 1, pose=[0.0, 0.0, 0.0, 0.0, 180.0]), |
| rec(2, 2, pose=[100.0, 0.0, 0.0, 0.0, 0.0]), |
| ] |
| result = deterministic_revisit_retrieval( |
| records, |
| target_frame=10, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| fov_overlap_threshold=0.5, |
| exclude_local_context_frames=4, |
| topk=4, |
| ) |
| assert result.diagnostics["selected_frame_record_ids"] == ["c0"] |
| assert result.diagnostics["valid_revisit_frame_count"] == 1 |
| assert result.diagnostics["valid_revisit_count"] == 1 |
| assert result.diagnostics["best_selected_fov_overlap"] == 1.0 |
| assert result.diagnostics["revisit_best_selected_fov_overlap_max"] == 1.0 |
| assert result.diagnostics["best_selected_gap_frames"] == 10 |
| assert result.diagnostics["revisit_fov_overlap_max"] == 1.0 |
| assert result.diagnostics["revisit_plucker_overlap_max"] > 0.0 |
|
|
|
|
| def test_pose_preselect_uses_local_position_and_view_direction_before_fov(): |
| records = [ |
| rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 180.0], chunk_id="opposite_same_position"), |
| rec(1, 1, pose=[90.0, 0.0, 0.0, 0.0, 0.0], chunk_id="far_same_direction"), |
| rec(2, 2, pose=[1.0, 0.0, 0.0, 0.0, 0.0], chunk_id="near_same_direction"), |
| ] |
| result = deterministic_revisit_retrieval( |
| records, |
| target_frame=10, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| fov_overlap_threshold=0.0, |
| fov_radius=30.0, |
| exclude_local_context_frames=4, |
| topk=1, |
| pose_preselect_topk=1, |
| ) |
|
|
| assert result.diagnostics["selected_frame_record_ids"] == ["near_same_direction"] |
| assert result.diagnostics["revisit_pose_preselect_input_count"] == 3 |
| assert result.diagnostics["revisit_pose_preselect_scored_count"] == 3 |
| assert result.diagnostics["revisit_pose_preselect_selected_count"] == 1 |
| assert result.diagnostics["revisit_exact_fov_candidate_count"] == 1 |
| assert result.diagnostics["revisit_vectorized_frame_scorer_used"] == 1 |
| assert abs(result.diagnostics["revisit_pose_preselect_min_distance"] - (1.0 / 30.0)) < 1e-6 |
|
|
|
|
| def test_selected_frame_carries_frame_metadata_for_projection(): |
| result = deterministic_revisit_retrieval( |
| [rec(1, 1, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id="frame_1")], |
| target_frame=8, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| fov_overlap_threshold=0.30, |
| high_quality_fov_threshold=0.70, |
| exclude_local_context_frames=4, |
| topk=1, |
| ) |
|
|
| assert result.diagnostics["selected_frame_record_ids"] == ["frame_1"] |
| assert result.selected_frame_ids == [1] |
| assert result.records[0].metadata["dememwm_selected_frame_index"] == 1 |
| assert result.records[0].metadata["dememwm_selected_frame_passes_high_quality"] is True |
| assert result.diagnostics["best_selected_frame_index"] == 1 |
| assert result.diagnostics["best_selected_frame_fov_overlap"] == 1.0 |
|
|
|
|
| def test_high_quality_threshold_is_selected_target_diagnostic_only(): |
| result = deterministic_revisit_retrieval( |
| [rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 0.0])], |
| target_frame=10, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 60.0]), |
| fov_overlap_threshold=0.30, |
| high_quality_fov_threshold=0.70, |
| exclude_local_context_frames=4, |
| topk=1, |
| ) |
| assert result.diagnostics["selected_frame_record_ids"] == ["c0"] |
| assert result.diagnostics["valid_revisit_count"] == 1 |
| assert 0.30 <= result.diagnostics["best_selected_fov_overlap"] < 0.70 |
|
|
|
|
| def test_video_metadata_does_not_filter_revisit_candidates(): |
| records = [ |
| rec(0, 0, video_id="v0", pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| rec(1, 1, video_id="other", pose=[0.0, 0.0, 0.0, 0.0, 0.0]), |
| ] |
| result = deterministic_revisit_retrieval( |
| records, |
| target_frame=10, |
| target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), |
| target_video_id="v0", |
| exclude_local_context_frames=4, |
| topk=4, |
| ) |
| assert result.diagnostics["selected_frame_record_ids"] == ["c1", "c0"] |
| assert result.diagnostics["valid_revisit_count"] == 2 |
|
|
|
|
| def test_tie_breaking_is_overlap_then_age_then_source_then_record_id(): |
| records = [ |
| rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id="b"), |
| rec(1, 1, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id="a"), |
| rec(2, 2, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id="c"), |
| ] |
| result = deterministic_revisit_retrieval(records, target_frame=10, target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), exclude_local_context_frames=4, topk=3) |
| assert result.diagnostics["selected_frame_record_ids"] == ["c", "a", "b"] |
|
|
