BonanDing commited on
Commit
dd214aa
·
1 Parent(s): a66df9a

Add DeMemWM multiview dynamic policy

Browse files
.exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md ADDED
@@ -0,0 +1,1091 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # DeMemWM Dynamic Multiview Memory Selection Plan
2
+
3
+ Repo: `/share_1/users/bonan_ding/WorldMem`
4
+
5
+ Branch: `dememwm`
6
+
7
+ Purpose: replace the current DeMemWM `dynamic` memory stream with target-relevant
8
+ multi-view memory selection while keeping the stream name `dynamic`, the packed
9
+ frame order, and the existing DeMemWM model architecture unchanged.
10
+
11
+ This is the current-branch task that should land before the unified-memory-
12
+ attention branch. It is a dataset/online-selection update only: `dynamic` still
13
+ appears in the code, configs, masks, segments, logs, and launchers, but
14
+ semantically it means multi-view memory selected from target-relevant candidate
15
+ frames.
16
+
17
+ The design intentionally stays simple and measurable:
18
+
19
+ - Add exactly one new dynamic policy name: `dynamic.selection_policy=multiview`.
20
+ - Do not add policy aliases such as `fov_multiview` or `fov_fps`.
21
+ - Do not add a hand-tuned score-weight family.
22
+ - Do not add any extra future-query-window config.
23
+ - Implement exact deterministic FOV/coverage selection for training and
24
+ validation/test.
25
+ - Also keep a pose/Plucker/FPS selector as a benchmarkable fallback under the
26
+ same `multiview` policy, so speed can be measured before choosing the final
27
+ training selector for long runs.
28
+ - Use local benchmarking to measure selector speed; do not assume FOV training
29
+ is fast enough without measurement.
30
+
31
+ ## Required Orchestration Flow
32
+
33
+ Use this plan with a new session prompt such as:
34
+
35
+ ```text
36
+ /goal Complete /share_1/users/bonan_ding/WorldMem/.exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md one substep at a time. For every substep, use a fresh implementation subagent, a fresh review subagent, and if needed fresh fix subagents until review is clean. Validate, inspect diff, commit the substep, and update this plan before moving to the next substep.
37
+ ```
38
+
39
+ The main session is the orchestra leader. It owns sequencing, diff inspection,
40
+ validation, commits, and plan-status updates. Subagents do bounded
41
+ implementation, review, or fix work only for the current substep.
42
+
43
+ For every substep below:
44
+
45
+ 1. Start a fresh implementation subagent.
46
+ 2. The implementation subagent must read:
47
+ - `AGENTS.md`
48
+ - `.codex/AGENT.md`
49
+ - `.codex/memory.md`
50
+ - `/share_1/users/bonan_ding/.codex/skills/codex-ml-cv-clean-implementation/SKILL.md`
51
+ 3. The implementation subagent must use:
52
+ - `[$codex-ml-cv-clean-implementation](/share_1/users/bonan_ding/.codex/skills/codex-ml-cv-clean-implementation/SKILL.md)`
53
+ 4. The implementation subagent must implement only the current substep.
54
+ 5. After implementation, the main session inspects the diff before review.
55
+ 6. Start a fresh review subagent.
56
+ 7. The review subagent must read the same files and use the same skill.
57
+ 8. The review subagent must check:
58
+ - correctness
59
+ - bugs or regressions
60
+ - improper or unnecessary changes
61
+ - code cleanliness
62
+ - adherence to this plan
63
+ - adherence to repository conventions
64
+ - targeted test adequacy
65
+ 9. If review finds issues, start a fresh fix subagent.
66
+ 10. The fix subagent must read the same files, use the same skill, and fix only
67
+ the review findings.
68
+ 11. Repeat review/fix/re-review until the review subagent reports no remaining
69
+ issues.
70
+ 12. Run the validation commands listed for the substep.
71
+ 13. Run `git diff --check`.
72
+ 14. Inspect `git diff`.
73
+ 15. Commit only the completed substep with a clear commit message.
74
+ 16. Mark the substep complete in this plan file and include that plan-status
75
+ update in the same substep commit if possible.
76
+ 17. Do not proceed to the next substep until the current one is reviewed clean,
77
+ validated, committed, and marked complete.
78
+
79
+ Global rules:
80
+
81
+ - Do not skip substeps.
82
+ - Do not combine unrelated substeps into one commit.
83
+ - Keep changes surgical.
84
+ - Preserve unrelated user changes.
85
+ - Do not train long jobs or submit Slurm jobs from this plan.
86
+ - Do not add or edit Berzelius/Slurm launchers in this plan.
87
+ - Do not modify `algorithms/worldmem`; use it only as a reference.
88
+ - Keep the packed frame-memory order unchanged:
89
+
90
+ ```text
91
+ [target][anchor][dynamic][revisit]
92
+ ```
93
+
94
+ - Keep the stream name `dynamic`; do not rename it to `multiview` in tensor
95
+ keys, config keys, segment names, masks, logs, or file names.
96
+ - Add exactly one new dynamic policy value: `multiview`.
97
+ - Add at most one selector-backend knob under the `dynamic` config:
98
+ `multiview_selector`. This is not a policy alias; it exists only to compare
99
+ exact FOV selection against a fast pose/Plucker/FPS fallback.
100
+ - Valid `multiview_selector` values are exactly:
101
+ - `fov_greedy`
102
+ - `pose_plucker_fps`
103
+ - Default the new local launchers to `dynamic.multiview_selector=fov_greedy`.
104
+ - Do not add policy aliases such as `fov_multiview` or `fov_fps`.
105
+ - Do not add score-weight knobs such as `fps_position_weight`,
106
+ `score_diversity_weight`, or `relevance_fov_weight`.
107
+ - Do not change the DeMemWM model architecture in this plan.
108
+ - Do not change `FrameMemoryReferenceAttention` in this plan.
109
+ - Do not change local target temporal-attention causality in this plan.
110
+ - Do not change anchor selection semantics: keep the same anchor candidate
111
+ window and `_select_anchor` scoring/diversity logic.
112
+ - Keep `dynamic.selection_policy=recent` and
113
+ `dynamic.selection_policy=event_triggered` working for backward
114
+ compatibility.
115
+ - The new `dynamic.selection_policy=multiview` should select target-relevant
116
+ co-visible support views, not event anchors.
117
+ - Dynamic multiview should avoid selecting the same raw frames as `revisit`
118
+ when enough alternatives exist. This intentionally overrides the older
119
+ duplicate-friendly behavior only for `multiview` dynamic versus revisit; keep
120
+ `recent` and `event_triggered` duplicate semantics unchanged.
121
+ - Do not compute the exact point-FOV candidate pool twice for the same query.
122
+ When `revisit` and `dynamic=multiview` both need FOV information, share one
123
+ candidate pool or one local cache entry.
124
+ - Do not add new future-memory behavior in this plan. If the current branch
125
+ already contains `_memory_candidate_frames()` and
126
+ `dataset.memory_selection.causal` from previous work, call that helper and
127
+ preserve its existing semantics exactly. If the helper is absent, use the
128
+ original past-only candidate semantics.
129
+ - Do not couple `dataset.memory_selection.causal` to `algorithm.causal`.
130
+ - Do not introduce self-generated-memory training in this plan.
131
+ - Do not add source/origin embeddings in this plan.
132
+ - If validation fails, fix through the same fresh fix-subagent and
133
+ review-subagent loop, then rerun validation.
134
+
135
+ Status legend:
136
+
137
+ - `[ ]` not started
138
+ - `[~]` in progress
139
+ - `[x]` complete and committed
140
+
141
+ ## Current Code Facts
142
+
143
+ - DeMemWM training samples are built by
144
+ `datasets/video/minecraft_video_dememwm_latent_dataset.py`, which calls
145
+ `select_memory_indices()` in `datasets/video/memory_selection.py`.
146
+ - The packed DeMemWM sequence order is currently:
147
+
148
+ ```text
149
+ [target][anchor][dynamic][revisit]
150
+ ```
151
+
152
+ - The current stream keys are `anchor`, `dynamic`, and `revisit`.
153
+ - The current dataset config exposes `max_anchor_frames`, `max_dynamic_frames`,
154
+ and `max_revisit_frames`.
155
+ - `dynamic.selection_policy` currently supports `recent` and `event_triggered`.
156
+ - The latent dataset builds `dynamic_stream` only for `event_triggered`; the new
157
+ `multiview` policy should not require a prebuilt event stream.
158
+ - `_select_revisit()` currently has different training and validation/test
159
+ behavior in the original code: training uses pose similarity, while
160
+ validation/test uses point-union/FOV selection.
161
+ - `_select_by_point_union()` already contains the validation/test FOV/Plucker/
162
+ coverage retrieval machinery needed for exact multiview selection.
163
+ - Exact point-FOV construction evaluates candidate visibility over sampled 3D
164
+ points. Selector speed must therefore be measured before choosing the final
165
+ long-run training backend.
166
+ - `_sample_points_in_sphere()` currently uses random point sampling; this plan
167
+ makes the FOV path deterministic for reproducible training, validation, and
168
+ online selection.
169
+ - Online validation/test rollout in `algorithms/dememwm/df_video.py` separately
170
+ selects online `anchor`, `dynamic`, and `revisit` memory from the generated or
171
+ observed prefix.
172
+
173
+ ## Selection Design
174
+
175
+ The strongest selector to implement first is coverage-driven FOV greedy
176
+ selection:
177
+
178
+ 1. Build a causal candidate pool for the current target positions.
179
+ 2. Pose-preselect candidates to the existing `pose_preselect_topk` budget before
180
+ exact FOV visibility is evaluated.
181
+ 3. Build deterministic target-visible 3D sample points.
182
+ 4. Compute candidate visibility masks over those points.
183
+ 5. Select direct `revisit` frames by greedy maximum coverage.
184
+ 6. Select `dynamic` multiview frames from the remaining candidates by greedy
185
+ residual coverage, initialized with the coverage already provided by selected
186
+ `revisit` frames.
187
+ 7. Use deterministic lexicographic tie-breakers only:
188
+ - larger new target-FOV coverage gain;
189
+ - larger total FOV overlap;
190
+ - larger Plucker score when available;
191
+ - larger pose distance from already selected/reference views;
192
+ - smaller temporal gap to the first target;
193
+ - smaller frame id.
194
+
195
+ The benchmarkable fallback selector is `pose_plucker_fps`:
196
+
197
+ 1. Use the current pose-similarity and optional Plucker relevance logic.
198
+ 2. Cap to `pose_preselect_topk` candidates.
199
+ 3. Select diverse views by deterministic farthest-point sampling in normalized
200
+ camera-position plus yaw/pitch pose space.
201
+ 4. Use this only as a fallback or speed-comparison backend; the new local
202
+ launchers should default to `fov_greedy`.
203
+
204
+ ## Non-Goals
205
+
206
+ - Do not create the unified dense memory attention module in this plan.
207
+ - Do not remove or rename `dynamic`.
208
+ - Do not remove the legacy `event_triggered` or `recent` dynamic policies.
209
+ - Do not submit cluster jobs.
210
+ - Do not write or edit Berzelius/Slurm scripts.
211
+ - Do not change Open-Oasis checkpoints, VAE behavior, DiT dimensions, or model
212
+ checkpoint loading.
213
+ - Do not introduce generated-memory replay or self-rollout training.
214
+ - Do not change `algorithms/worldmem`.
215
+ - Do not add a new external retrieval backend.
216
+ - Do not change the memory segment order.
217
+ - Do not add any extra future-query-window config.
218
+
219
+ ## Substep 1: Add Minimal Multiview Dynamic Policy Config
220
+
221
+ Status: `[x]`
222
+
223
+ Goal:
224
+
225
+ Add configuration and policy plumbing for `dynamic.selection_policy=multiview`
226
+ while keeping existing `recent` and `event_triggered` behavior working. Keep the
227
+ config minimal and avoid unvalidated score-weight knobs.
228
+
229
+ Primary files:
230
+
231
+ - `configurations/dataset/video_minecraft_dememwm_latent.yaml`
232
+ - `datasets/video/memory_selection.py`
233
+ - `tests/test_dememwm_latent_dataset.py`
234
+
235
+ Required behavior:
236
+
237
+ - Add or preserve this memory budget:
238
+
239
+ ```yaml
240
+ memory_selection:
241
+ max_anchor_frames: 2
242
+ max_dynamic_frames: 4
243
+ max_revisit_frames: 2
244
+ ```
245
+
246
+ - Do not add any extra future-query-window config.
247
+ - Do not switch the dataset default policy to `multiview` in this substep. The
248
+ selector is not implemented until Substep 4, so the default should remain the
249
+ current working policy here. The default can be switched after integration in
250
+ Substep 5.
251
+ - Add comments or tests that document the intended new policy value:
252
+
253
+ ```yaml
254
+ memory_selection:
255
+ dynamic:
256
+ selection_policy: multiview
257
+ multiview_selector: fov_greedy
258
+ ```
259
+
260
+ - Keep the existing global retrieval-speed controls instead of adding new
261
+ dynamic-specific score weights:
262
+
263
+ ```yaml
264
+ memory_selection:
265
+ pose_preselect_topk: 64
266
+ candidate_chunk_size: 64
267
+ ```
268
+
269
+ - Add a single selector-backend knob:
270
+
271
+ ```yaml
272
+ memory_selection:
273
+ dynamic:
274
+ multiview_selector: fov_greedy
275
+ ```
276
+
277
+ - Valid backend values are exactly:
278
+
279
+ ```python
280
+ {"fov_greedy", "pose_plucker_fps"}
281
+ ```
282
+
283
+ - In `memory_selection.py`, allow exactly these dynamic policy names:
284
+
285
+ ```python
286
+ {"recent", "event_triggered", "multiview"}
287
+ ```
288
+
289
+ - Unknown policies should raise a clear error mentioning exactly those valid
290
+ policies.
291
+ - Unknown `multiview_selector` values should raise a clear error mentioning
292
+ exactly `fov_greedy` and `pose_plucker_fps`.
293
+ - Do not add policy aliases such as `fov_multiview` or `fov_fps`.
294
+ - Do not implement the actual multiview selector yet in this substep. If
295
+ top-level selection is called with `multiview` before Substep 4, it should fail
296
+ with a clear `NotImplementedError` or equivalent explicit error rather than
297
+ silently returning an unrelated fallback.
298
+ - Do not build `dynamic_stream` for the `multiview` policy.
299
+
300
+ Implementation constraints:
301
+
302
+ - Keep changes to config/policy validation only.
303
+ - Keep tests small and deterministic.
304
+ - Do not touch model code.
305
+
306
+ Required tests:
307
+
308
+ - `dynamic.selection_policy=multiview` is accepted by policy parsing/validation.
309
+ - `dynamic.selection_policy=recent` still works.
310
+ - `dynamic.selection_policy=event_triggered` still works.
311
+ - `dynamic.selection_policy=fov_multiview` and `dynamic.selection_policy=fov_fps`
312
+ are rejected.
313
+ - `dynamic.multiview_selector=fov_greedy` is accepted.
314
+ - `dynamic.multiview_selector=pose_plucker_fps` is accepted.
315
+ - Unknown policy and backend values raise clear errors containing valid values.
316
+
317
+ Validation for this substep:
318
+
319
+ ```bash
320
+ python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
321
+ python -B -m unittest tests.test_dememwm_latent_dataset
322
+ git diff --check
323
+ ```
324
+
325
+ Review must specifically confirm:
326
+
327
+ - The stream is still named `dynamic` everywhere.
328
+ - Only one new policy value, `multiview`, was added.
329
+ - Only one backend selector knob was added.
330
+ - No new hand-tuned weight config family was added.
331
+ - Existing dynamic policies remain valid.
332
+ - The default dataset config was not flipped to an unimplemented selector in this
333
+ substep.
334
+ - No model architecture code was touched.
335
+ - No training or Slurm job was run.
336
+
337
+ Suggested commit message:
338
+
339
+ ```text
340
+ Add DeMemWM multiview dynamic policy
341
+ ```
342
+
343
+ ## Substep 2: Make FOV Candidate Construction Deterministic And Reusable
344
+
345
+ Status: `[ ]`
346
+
347
+ Goal:
348
+
349
+ Refactor the existing FOV/Plucker candidate logic so training, validation, and
350
+ online `revisit`/`dynamic=multiview` can share one deterministic candidate-pool
351
+ path.
352
+
353
+ Primary files:
354
+
355
+ - `datasets/video/memory_selection.py`
356
+ - `tests/test_dememwm_latent_dataset.py`
357
+
358
+ Required behavior:
359
+
360
+ - Replace random FOV point sampling with deterministic sampling.
361
+ - Add:
362
+
363
+ ```python
364
+ def _deterministic_points_in_unit_ball(num_points: int, device, dtype) -> torch.Tensor:
365
+ """Return deterministic points in the unit ball."""
366
+ ```
367
+
368
+ - Update `_sample_points_in_sphere(center)` to use deterministic unit-ball
369
+ points and no `torch.rand`.
370
+ - If the code currently uses a module constant such as `_FOV_NUM_POINTS`, keep
371
+ the default value unchanged unless a config already exists. Do not add a large
372
+ new sampling hyperparameter family in this plan.
373
+ - Add a reusable helper:
374
+
375
+ ```python
376
+ def _build_fov_candidate_pool(
377
+ poses: np.ndarray,
378
+ candidates: np.ndarray,
379
+ target_positions: np.ndarray,
380
+ cfg,
381
+ *,
382
+ use_plucker: bool,
383
+ ) -> dict[str, torch.Tensor]:
384
+ """Build deterministic FOV/Plucker/coverage data for candidate frames."""
385
+ ```
386
+
387
+ - The helper should:
388
+ - call `_pose_preselect()` so exact FOV runs on at most
389
+ `pose_preselect_topk` candidates when that config is set;
390
+ - reuse `_candidate_fov_masks()` and `candidate_chunk_size`;
391
+ - optionally compute Plucker scores;
392
+ - return enough candidate pose/coverage data for revisit and dynamic to reuse
393
+ the same pool for the same target query;
394
+ - return empty tensors cleanly when no candidates survive.
395
+ - The helper should return at least:
396
+
397
+ ```text
398
+ candidates_t: LongTensor [N] # frame ids after pose preselect
399
+ candidate_poses: FloatTensor [N, 5] # poses for candidates_t
400
+ inside: BoolTensor [N, P] # candidate-visible target points
401
+ fov_values: FloatTensor [N]
402
+ plucker: FloatTensor [N]
403
+ gaps: LongTensor [N]
404
+ positive_fov: BoolTensor [N]
405
+ ```
406
+
407
+ - `positive_fov` should mean `inside.any(dim=1)` or equivalently positive
408
+ visible-point coverage. It is not a new configurable threshold.
409
+ - Refactor `_select_by_point_union()` to use the helper without changing its
410
+ high-level greedy coverage behavior.
411
+ - Add one small internal pool-selection helper if needed, for example
412
+ `_select_by_point_union_from_pool(pool, count, ...)`, so multiview selection
413
+ can reuse the same FOV pool instead of recomputing point visibility.
414
+ - Add one small reuse mechanism for the helper output. Acceptable options are:
415
+ - an explicit optional `fov_pool` argument passed between selectors; or
416
+ - a local per-call cache keyed by target positions, candidate ids, and
417
+ Plucker/config values.
418
+ - The reuse mechanism must be local and simple. Do not add a new global backend
419
+ or persistent disk cache in this plan.
420
+
421
+ Implementation constraints:
422
+
423
+ - Do not implement dynamic multiview selection in this substep.
424
+ - Avoid broad refactors; keep helper local to `memory_selection.py`.
425
+ - Cache deterministic base points in a small module-level cache if needed, keyed
426
+ by `(num_points, device, dtype)`, to avoid rebuilding the same base point set
427
+ repeatedly.
428
+ - If caching GPU tensors, keep the cache bounded to this tiny base-point set or
429
+ per-query local variables. Do not introduce an unbounded candidate-pool cache.
430
+
431
+ Required tests:
432
+
433
+ - Repeated calls to the FOV selector with the same inputs return identical
434
+ selected frames.
435
+ - Empty candidate pools return empty selected arrays, not errors.
436
+ - Existing FOV revisit selection still returns causal frames when the current
437
+ branch is causal.
438
+ - Unit tests confirm `torch.rand` is no longer used in `_sample_points_in_sphere`.
439
+ - Revisit selection still passes existing deterministic tests.
440
+
441
+ Validation for this substep:
442
+
443
+ ```bash
444
+ python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
445
+ python -B -m unittest tests.test_dememwm_latent_dataset
446
+ git diff --check
447
+ ```
448
+
449
+ Review must specifically confirm:
450
+
451
+ - `torch.rand` is no longer used for FOV point sampling.
452
+ - `_select_by_point_union()` still has the same greedy coverage behavior.
453
+ - The helper does not introduce device or dtype mismatches.
454
+ - The helper output can be reused so revisit and dynamic multiview do not build
455
+ identical FOV masks twice for the same query.
456
+ - No model code was touched.
457
+
458
+ Suggested commit message:
459
+
460
+ ```text
461
+ Make DeMemWM FOV candidates deterministic
462
+ ```
463
+
464
+ ## Substep 3: Add Deterministic Multiview Selector Backends
465
+
466
+ Status: `[ ]`
467
+
468
+ Goal:
469
+
470
+ Implement the two selector backends used by `dynamic.selection_policy=multiview`:
471
+ exact FOV greedy selection and pose/Plucker/FPS fallback selection.
472
+
473
+ Primary files:
474
+
475
+ - `datasets/video/memory_selection.py`
476
+ - `tests/test_dememwm_latent_dataset.py`
477
+
478
+ Required behavior:
479
+
480
+ - Add or factor a helper for existing pose-similarity + optional Plucker
481
+ relevance, for example:
482
+
483
+ ```python
484
+ def _rank_pose_plucker_candidates(
485
+ poses: np.ndarray,
486
+ candidates: np.ndarray,
487
+ target_positions: np.ndarray,
488
+ cfg,
489
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
490
+ """Return candidate ids, candidate poses, and relevance scores."""
491
+ ```
492
+
493
+ - The helper should preserve the existing scoring semantics used by
494
+ `_select_by_pose_similarity()`:
495
+ - `pose_similarity_threshold`
496
+ - `training_use_plucker`
497
+ - `training_plucker_weight`
498
+ - deterministic ranking by score and frame id
499
+ - `_select_by_pose_similarity()` should call this helper so the old revisit
500
+ behavior remains unchanged when that path is configured.
501
+ - Add a deterministic pose-FPS helper:
502
+
503
+ ```python
504
+ def _select_pose_fps(
505
+ candidate_ids: torch.Tensor,
506
+ candidate_poses: torch.Tensor,
507
+ count: int,
508
+ *,
509
+ seed_ids: np.ndarray | None = None,
510
+ all_poses: np.ndarray | torch.Tensor | None = None,
511
+ ) -> np.ndarray:
512
+ """Select diverse camera poses from a relevance-filtered candidate set."""
513
+ ```
514
+
515
+ - Pose-FPS distance should use existing pose scales only and no tunable weights:
516
+ - translation distance normalized by `_POSE_DISTANCE_SCALE`;
517
+ - yaw/pitch angular distance normalized by `_ANGLE_DISTANCE_SCALE`;
518
+ - final distance is the unweighted sum of normalized translation and normalized
519
+ angular distance.
520
+ - Before FPS, cap the relevance-filtered candidate set to the existing
521
+ `pose_preselect_topk` budget when configured, using deterministic relevance
522
+ ranking.
523
+ - If `seed_ids` are provided, the first dynamic frame should be far from the seed
524
+ views when possible. Here the seed views will usually be selected `revisit`
525
+ frames.
526
+ - Add an FOV-greedy dynamic selector helper that consumes the FOV pool from
527
+ Substep 2:
528
+
529
+ ```python
530
+ def _select_dynamic_multiview_from_fov_pool(
531
+ pool: dict[str, torch.Tensor],
532
+ count: int,
533
+ *,
534
+ excluded: np.ndarray | None = None,
535
+ reference_frames: np.ndarray | None = None,
536
+ all_poses: np.ndarray | torch.Tensor | None = None,
537
+ ) -> np.ndarray:
538
+ """Select dynamic multiview frames by residual FOV coverage."""
539
+ ```
540
+
541
+ - FOV-greedy selection should:
542
+ 1. Remove `excluded` frames.
543
+ 2. Initialize covered target points from `reference_frames` when those frames
544
+ are available in the same pool.
545
+ 3. Consider only candidates with `positive_fov=True` when available.
546
+ 4. Repeatedly choose the next frame by the deterministic priority:
547
+ - larger new target-FOV coverage gain;
548
+ - larger total FOV overlap;
549
+ - larger Plucker score;
550
+ - larger pose distance from selected/reference views;
551
+ - smaller temporal gap to the first target;
552
+ - smaller frame id.
553
+ 5. If all remaining positive-FOV candidates have zero new coverage, continue
554
+ selecting by the same priority with zero coverage gain, so dynamic can still
555
+ add co-visible alternative views.
556
+ 6. If no positive-FOV candidate exists, fall back to `_select_pose_fps()` over
557
+ the pose-preselected candidates.
558
+ - Do not use arbitrary weighted sums.
559
+ - Tie-breaking must be deterministic.
560
+
561
+ Implementation constraints:
562
+
563
+ - Do not change anchor selection.
564
+ - Do not change the legacy event-triggered or recent dynamic paths.
565
+ - Do not add new score-weight config fields.
566
+ - Keep all helper functions local to `memory_selection.py`.
567
+
568
+ Required tests:
569
+
570
+ - The factored `_select_by_pose_similarity()` returns the same results as before
571
+ on deterministic cases.
572
+ - Pose-FPS returns at most `count` frames.
573
+ - Pose-FPS is deterministic.
574
+ - Pose-FPS caps the candidate set with `pose_preselect_topk` when configured.
575
+ - Pose-FPS prefers diverse candidate poses when relevance filtering leaves
576
+ multiple valid candidates.
577
+ - FOV-greedy dynamic selection avoids excluded revisit frames when alternatives
578
+ exist.
579
+ - FOV-greedy dynamic selection is deterministic.
580
+ - FOV-greedy dynamic selection prefers larger residual coverage before tie-breaks.
581
+
582
+ Validation for this substep:
583
+
584
+ ```bash
585
+ python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
586
+ python -B -m unittest tests.test_dememwm_latent_dataset
587
+ git diff --check
588
+ ```
589
+
590
+ Review must specifically confirm:
591
+
592
+ - FOV-greedy uses residual coverage, not an arbitrary weighted sum.
593
+ - Pose-FPS is available only as a selector backend/fallback, not as a new dynamic
594
+ policy name.
595
+ - Existing pose-similarity revisit behavior is preserved.
596
+ - No new arbitrary score-weight config fields were introduced.
597
+ - No model code was touched.
598
+
599
+ Suggested commit message:
600
+
601
+ ```text
602
+ Add DeMemWM multiview selector backends
603
+ ```
604
+
605
+ ## Substep 4: Implement Dynamic-As-Multiview Selection
606
+
607
+ Status: `[ ]`
608
+
609
+ Goal:
610
+
611
+ Implement `dynamic.selection_policy=multiview`. The stream name remains
612
+ `dynamic`, but it now selects target-relevant multi-view support frames.
613
+
614
+ Primary files:
615
+
616
+ - `datasets/video/memory_selection.py`
617
+ - `tests/test_dememwm_latent_dataset.py`
618
+
619
+ Required behavior:
620
+
621
+ - Add:
622
+
623
+ ```python
624
+ def _select_dynamic_multiview(
625
+ poses: np.ndarray,
626
+ target_positions: np.ndarray,
627
+ cfg,
628
+ count: int,
629
+ *,
630
+ excluded: np.ndarray | None = None,
631
+ reference_frames: np.ndarray | None = None,
632
+ split: str = "training",
633
+ min_candidate_frame: int = 0,
634
+ fov_pool: dict[str, torch.Tensor] | None = None,
635
+ ) -> np.ndarray:
636
+ """Select dynamic-as-multiview memory frames."""
637
+ ```
638
+
639
+ - Candidate construction:
640
+ - use the existing `_memory_candidate_frames()` helper if it is present;
641
+ - do not modify that helper or its causal/non-causal semantics in this plan;
642
+ - if the helper is absent, use the original past-only candidates before the
643
+ first target frame;
644
+ - remove `excluded` frames from the candidate pool.
645
+ - For top-level use, pass selected `revisit` frames as `excluded` and
646
+ `reference_frames` so dynamic complements direct revisit memory.
647
+ - If `dynamic.multiview_selector == "fov_greedy"`:
648
+ 1. Use the provided `fov_pool` if available; otherwise call
649
+ `_build_fov_candidate_pool(..., use_plucker=True)` once.
650
+ 2. Call `_select_dynamic_multiview_from_fov_pool()`.
651
+ 3. Use this path for both training and validation/test when configured.
652
+ - If `dynamic.multiview_selector == "pose_plucker_fps"`:
653
+ 1. Use `_rank_pose_plucker_candidates()`.
654
+ 2. Cap to `pose_preselect_topk` candidates after relevance ranking.
655
+ 3. Run `_select_pose_fps()` seeded by `reference_frames` when available.
656
+ 4. This backend exists for speed comparison and fallback, not as the default
657
+ local-launcher backend.
658
+ - Return sorted frame indices.
659
+
660
+ Implementation constraints:
661
+
662
+ - Keep the helper local to `memory_selection.py`.
663
+ - Do not remove `event_triggered` or `recent` code paths.
664
+ - Keep tie-breaking deterministic.
665
+ - Keep tensor operations small enough for unit tests.
666
+ - Do not add `fov_multiview` or `fov_fps` aliases.
667
+ - Do not add any extra future-query-window config.
668
+
669
+ Required tests:
670
+
671
+ - `dynamic.selection_policy=multiview` returns past-only frames when the current
672
+ branch uses the original past-only candidate semantics.
673
+ - If the current branch already contains `_memory_candidate_frames()` with a
674
+ configurable non-causal training mode, multiview dynamic follows that helper's
675
+ existing behavior without reimplementing future-memory logic.
676
+ - Dynamic multiview frames are not selected from already selected revisit frames
677
+ when enough alternatives exist, because top-level use passes
678
+ `excluded=revisit`.
679
+ - `fov_greedy` uses the FOV helper during training and validation/test.
680
+ - `pose_plucker_fps` uses the pose/Plucker/FPS path and does not invoke exact
681
+ FOV candidate masks.
682
+ - Repeated calls with identical inputs return identical selections.
683
+ - `recent` and `event_triggered` dynamic policies still work.
684
+
685
+ Validation for this substep:
686
+
687
+ ```bash
688
+ python -m py_compile datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
689
+ python -B -m unittest tests.test_dememwm_latent_dataset
690
+ git diff --check
691
+ ```
692
+
693
+ Review must specifically confirm:
694
+
695
+ - The new selector is target-relevant and multi-view/diversity-aware.
696
+ - `fov_greedy` works for training as well as validation/test.
697
+ - `pose_plucker_fps` is available for speed benchmarking/fallback.
698
+ - The stream name remains `dynamic`.
699
+ - Existing dynamic policies remain functional.
700
+ - No model code was touched.
701
+
702
+ Suggested commit message:
703
+
704
+ ```text
705
+ Select DeMemWM dynamic memory by multiview retrieval
706
+ ```
707
+
708
+ ## Substep 5: Integrate Multiview Dynamic In Dataset And Online Rollout
709
+
710
+ Status: `[ ]`
711
+
712
+ Goal:
713
+
714
+ Use the new dynamic-as-multiview selector in top-level dataset memory selection
715
+ and online validation/test rollout.
716
+
717
+ Primary files:
718
+
719
+ - `datasets/video/memory_selection.py`
720
+ - `algorithms/dememwm/df_video.py`
721
+ - `datasets/video/minecraft_video_dememwm_latent_dataset.py`
722
+ - `tests/test_dememwm_latent_dataset.py`
723
+
724
+ Required behavior:
725
+
726
+ - In `select_memory_indices()`:
727
+ - select `anchor` with existing anchor behavior;
728
+ - if `dynamic.selection_policy=multiview`, build one deterministic candidate
729
+ pool for the target query when `dynamic.multiview_selector=fov_greedy`;
730
+ - select `revisit` using the existing configured behavior, but allow it to
731
+ reuse the shared FOV pool when the FOV path is used;
732
+ - select `dynamic` by calling `_select_dynamic_multiview()` with
733
+ `excluded=revisit`, `reference_frames=revisit`, and the shared `fov_pool`
734
+ when available;
735
+ - if `dynamic.selection_policy=event_triggered`, keep the existing event path;
736
+ - if `dynamic.selection_policy=recent`, keep the existing recent path.
737
+ - After the selector and online path are implemented, switch the dataset config
738
+ default to:
739
+
740
+ ```yaml
741
+ memory_selection:
742
+ dynamic:
743
+ selection_policy: multiview
744
+ multiview_selector: fov_greedy
745
+ ```
746
+
747
+ - In the latent dataset, keep building `dynamic_stream` only when policy is
748
+ `event_triggered`.
749
+ - In online validation/test in `df_video.py`:
750
+ - add support for `dynamic.selection_policy=multiview`;
751
+ - use the existing current target positions from the rollout code; do not add
752
+ or emulate a retrieval horizon;
753
+ - build or retrieve the deterministic FOV candidate pool once for the online
754
+ query when using `fov_greedy` and share it between online `revisit` and
755
+ `dynamic` multiview;
756
+ - do not build or extend the online event cache for the multiview policy;
757
+ - keep the event-cache path unchanged for `event_triggered`.
758
+
759
+ Implementation constraints:
760
+
761
+ - Do not change the packed memory order.
762
+ - Do not change online causal history limits.
763
+ - Do not touch model code.
764
+ - Keep current masks/padding behavior.
765
+ - Do not add any extra future-query-window config.
766
+
767
+ Required tests:
768
+
769
+ - Dataset selection returns keys `anchor`, `dynamic`, and `revisit` with expected
770
+ maximum lengths.
771
+ - Online selector accepts `dynamic.selection_policy=multiview` without requiring
772
+ an event cache.
773
+ - Online selection uses the current rollout target positions and no retrieval
774
+ horizon.
775
+ - Online revisit and dynamic multiview share the same FOV candidate pool or
776
+ cache entry for a query when `fov_greedy` is configured.
777
+ - Existing online `event_triggered` behavior remains available.
778
+
779
+ Validation for this substep:
780
+
781
+ ```bash
782
+ python -m py_compile datasets/video/memory_selection.py algorithms/dememwm/df_video.py datasets/video/minecraft_video_dememwm_latent_dataset.py tests/test_dememwm_latent_dataset.py
783
+ python -B -m unittest tests.test_dememwm_latent_dataset
784
+ git diff --check
785
+ ```
786
+
787
+ Review must specifically confirm:
788
+
789
+ - Training and online validation both support multiview dynamic.
790
+ - The dataset default policy is switched to `multiview` only after the selector
791
+ and online integration exist.
792
+ - Online validation/test avoids avoidable duplicate exact-FOV candidate passes
793
+ when selecting revisit and dynamic for the same target query.
794
+ - `dynamic_stream` is not required for the multiview policy.
795
+ - The online event cache remains limited to `event_triggered`.
796
+ - Exact FOV construction is not duplicated for revisit and dynamic multiview in
797
+ the same online query.
798
+ - The frame-memory segment order is unchanged.
799
+
800
+ Suggested commit message:
801
+
802
+ ```text
803
+ Use multiview dynamic memory online
804
+ ```
805
+
806
+ ## Substep 6: Add Selector Speed Benchmark
807
+
808
+ Status: `[ ]`
809
+
810
+ Goal:
811
+
812
+ Add a local benchmark so selector speed can be measured before deciding whether
813
+ `fov_greedy` is acceptable for full training or whether the fallback
814
+ `pose_plucker_fps` should be used for long runs.
815
+
816
+ Primary files:
817
+
818
+ - `scripts/benchmark_dememwm_multiview_selection.py`
819
+ - `datasets/video/memory_selection.py` only if a small public/importable wrapper
820
+ is needed
821
+ - `tests/test_dememwm_latent_dataset.py` or a small benchmark smoke test
822
+ - this plan file
823
+
824
+ Required behavior:
825
+
826
+ - Add a benchmark script that can run without submitting jobs:
827
+
828
+ ```bash
829
+ python scripts/benchmark_dememwm_multiview_selection.py \
830
+ --num-frames 1200 \
831
+ --target-start 600 \
832
+ --target-len 8 \
833
+ --num-iters 100 \
834
+ --pose-preselect-topk 64 \
835
+ --candidate-chunk-size 64
836
+ ```
837
+
838
+ - The script should benchmark both selector backends on the same deterministic
839
+ pose sequence:
840
+ - `fov_greedy`
841
+ - `pose_plucker_fps`
842
+ - The script should report at least:
843
+ - mean milliseconds per selection;
844
+ - median milliseconds per selection;
845
+ - p90 milliseconds per selection;
846
+ - selected frame count;
847
+ - candidate count after pose preselection;
848
+ - whether FOV pool reuse is enabled;
849
+ - CPU/GPU device used if applicable.
850
+ - The script may support optional real-dataset inputs if this can be done
851
+ surgically, but it must have a deterministic synthetic-pose mode that works in
852
+ CI/local code validation without data paths.
853
+ - The script should not choose or edit the final default automatically. It only
854
+ reports measurements.
855
+ - Add a short Markdown report template or script option such as
856
+ `--write-report .exp_artifact/dememwm_multiview_selection_speed_report.md`.
857
+ If the benchmark is run only on synthetic data, the report must explicitly say
858
+ the result is not a substitute for a real dataset sampling benchmark.
859
+
860
+ Implementation constraints:
861
+
862
+ - Do not train.
863
+ - Do not require GPUs for the synthetic benchmark.
864
+ - Do not import heavy training modules when direct selector imports are enough.
865
+ - Do not add external dependencies.
866
+ - Do not use the benchmark result to change the committed default automatically.
867
+
868
+ Required tests:
869
+
870
+ - The benchmark script runs in synthetic mode with a tiny iteration count.
871
+ - The script reports both selector names.
872
+ - The script exits nonzero on unknown selector names.
873
+ - The script does not require dataset paths in synthetic mode.
874
+
875
+ Validation for this substep:
876
+
877
+ ```bash
878
+ python -m py_compile scripts/benchmark_dememwm_multiview_selection.py datasets/video/memory_selection.py tests/test_dememwm_latent_dataset.py
879
+ python scripts/benchmark_dememwm_multiview_selection.py --num-frames 128 --target-start 64 --target-len 4 --num-iters 3 --pose-preselect-topk 16 --candidate-chunk-size 16
880
+ python -B -m unittest tests.test_dememwm_latent_dataset
881
+ git diff --check
882
+ ```
883
+
884
+ Review must specifically confirm:
885
+
886
+ - The benchmark compares the exact FOV selector and the pose/Plucker/FPS fallback
887
+ on the same inputs.
888
+ - The benchmark reports timing rather than making an ungrounded default decision.
889
+ - The benchmark can run without Slurm, training, or dataset paths.
890
+ - No external dependencies were added.
891
+
892
+ Suggested commit message:
893
+
894
+ ```text
895
+ Add DeMemWM multiview selector benchmark
896
+ ```
897
+
898
+ ## Substep 7: Add Local Launchers And Final Validation
899
+
900
+ Status: `[ ]`
901
+
902
+ Goal:
903
+
904
+ Add local launchers for the current-branch dynamic-multiview experiment and run
905
+ cheap code validation. Do not add Berzelius or Slurm scripts in this plan.
906
+
907
+ Primary files:
908
+
909
+ - `train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_2h200_debug.sh`
910
+ - `train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_full.sh`
911
+ - `configurations/dataset/video_minecraft_dememwm_latent.yaml`
912
+ - `configurations/algorithm/dememwm_base.yaml` only if a config key is missing
913
+ - tests touched by previous substeps
914
+ - this plan file
915
+
916
+ Required launcher behavior:
917
+
918
+ - The debug launcher should be local-only and use two GPUs by default:
919
+
920
+ ```bash
921
+ CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1}
922
+ NPROC_PER_NODE=${NPROC_PER_NODE:-2}
923
+ torchrun --standalone --nproc_per_node=${NPROC_PER_NODE} -m main ...
924
+ ```
925
+
926
+ - The full launcher should also be local-only and parameterized by environment
927
+ variables. Do not assume eight local GPUs. Default to the same local 2-H200
928
+ setup unless the caller overrides it:
929
+
930
+ ```bash
931
+ CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1}
932
+ NPROC_PER_NODE=${NPROC_PER_NODE:-2}
933
+ torchrun --standalone --nproc_per_node=${NPROC_PER_NODE} -m main ...
934
+ ```
935
+
936
+ - Neither launcher may contain:
937
+ - `#SBATCH`
938
+ - `srun`
939
+ - `module load`
940
+ - `berzelius`
941
+ - cluster account names
942
+ - hard-coded cluster paths
943
+ - Both launchers should set:
944
+
945
+ ```bash
946
+ dataset.memory_selection.max_anchor_frames=2
947
+ dataset.memory_selection.max_dynamic_frames=4
948
+ dataset.memory_selection.max_revisit_frames=2
949
+ dataset.memory_selection.causal=true
950
+ dataset.memory_selection.dynamic.selection_policy=multiview
951
+ dataset.memory_selection.dynamic.multiview_selector=fov_greedy
952
+ algorithm.noise_route.anchor=all
953
+ algorithm.noise_route.dynamic=all
954
+ algorithm.noise_route.revisit=all
955
+ algorithm.memory_noise.enabled=false
956
+ algorithm.memory_noise.anchor_max_fraction=0.0
957
+ algorithm.memory_noise.dynamic_max_fraction=0.0
958
+ algorithm.memory_noise.revisit_max_fraction=0.0
959
+ +algorithm.memory_noise.validation_noisy_memory=false
960
+ ```
961
+
962
+ - The debug launcher should use tiny/debug settings, for example:
963
+
964
+ ```bash
965
+ experiment.training.max_steps=50
966
+ experiment.validation.val_every_n_step=25
967
+ experiment.validation.limit_batch=1
968
+ experiment.training.checkpointing.every_n_train_steps=25
969
+ wandb.mode=offline
970
+ ```
971
+
972
+ - The full local launcher should keep the intended full training schedule, but
973
+ use local environment variables for paths such as `DATA_ROOT`,
974
+ `FEATURE_ROOT`, `OUTPUT_ROOT`, `DIFFUSION_MODEL_PATH`, and `VAE_PATH`.
975
+ - Do not run the full training job automatically as part of Codex validation.
976
+ - The local 2H200 debug command may be run manually after code review if the
977
+ local machine has the GPUs and paths configured; keep it separate from the
978
+ cheap code-validation commands below.
979
+
980
+ Cheap final validation:
981
+
982
+ ```bash
983
+ python -m py_compile \
984
+ datasets/video/memory_selection.py \
985
+ datasets/video/minecraft_video_dememwm_latent_dataset.py \
986
+ algorithms/dememwm/df_video.py \
987
+ scripts/benchmark_dememwm_multiview_selection.py \
988
+ tests/test_dememwm_latent_dataset.py
989
+
990
+ python -B -m unittest tests.test_dememwm_latent_dataset
991
+
992
+ python scripts/benchmark_dememwm_multiview_selection.py \
993
+ --num-frames 128 \
994
+ --target-start 64 \
995
+ --target-len 4 \
996
+ --num-iters 3 \
997
+ --pose-preselect-topk 16 \
998
+ --candidate-chunk-size 16
999
+
1000
+ grep -R "#SBATCH\|srun\|berzelius\|module load" \
1001
+ train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_2h200_debug.sh \
1002
+ train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_full.sh && exit 1 || true
1003
+
1004
+
1005
+ git diff --check
1006
+ git status --short
1007
+ ```
1008
+
1009
+ Manual local run commands, not automatic Codex validation:
1010
+
1011
+ ```bash
1012
+ python scripts/benchmark_dememwm_multiview_selection.py \
1013
+ --num-frames 1200 \
1014
+ --target-start 600 \
1015
+ --target-len 8 \
1016
+ --num-iters 100 \
1017
+ --pose-preselect-topk 64 \
1018
+ --candidate-chunk-size 64 \
1019
+ --write-report .exp_artifact/dememwm_multiview_selection_speed_report.md
1020
+
1021
+ bash train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_2h200_debug.sh
1022
+ bash train_dememwm_curric_multiview_anchorclean_causal_cleanmem_local_full.sh
1023
+ ```
1024
+
1025
+ Review must specifically confirm:
1026
+
1027
+ - Both launchers are local-only.
1028
+ - No Berzelius/Slurm launcher was added or edited.
1029
+ - The debug launcher is suitable for a local 2H200 smoke run.
1030
+ - The full launcher is local, defaults to 2 GPUs, and is parameterized by
1031
+ environment variables for larger local runs.
1032
+ - Both launchers default to `dynamic.multiview_selector=fov_greedy`.
1033
+ - The benchmark script exists and can compare `fov_greedy` with
1034
+ `pose_plucker_fps`.
1035
+ - No model architecture changes were included in this plan.
1036
+ - The final behavior can be enabled with:
1037
+
1038
+ ```bash
1039
+ dataset.memory_selection.dynamic.selection_policy=multiview
1040
+ dataset.memory_selection.dynamic.multiview_selector=fov_greedy
1041
+ ```
1042
+
1043
+ Suggested commit message:
1044
+
1045
+ ```text
1046
+ Add local DeMemWM multiview launchers
1047
+ ```
1048
+
1049
+ ## New-Session Goal Prompt
1050
+
1051
+ Use this prompt to start the implementation session:
1052
+
1053
+ ```text
1054
+ Use $codex-ml-cv-clean-implementation.
1055
+ Repo: /share_1/users/bonan_ding/WorldMem
1056
+ /goal Complete .exp_artifact/dememwm_dynamic_multiview_memory_selection_plan.md step by step.
1057
+
1058
+ Act as the main orchestra leader. For each substep in the plan:
1059
+ 1. Use a fresh implementation subagent for only that substep.
1060
+ 2. Inspect the diff yourself.
1061
+ 3. Use a fresh review subagent to review correctness, bugs, design, cleanliness, and tests.
1062
+ 4. If review finds any issue, use a fresh fix subagent to fix only those findings.
1063
+ 5. Repeat review/fix with fresh subagents until review is clean.
1064
+ 6. Run the substep validation commands and git diff --check.
1065
+ 7. Inspect git diff.
1066
+ 8. Commit only that substep with the suggested or similarly clear commit message.
1067
+ 9. Mark the substep complete in the plan file and include the status update in the same commit if possible.
1068
+ 10. Move to the next substep only after the current substep is reviewed clean, validated, committed, and marked complete.
1069
+
1070
+ Hard constraints:
1071
+ - Do not train long jobs or submit Slurm jobs.
1072
+ - Do not add or edit Berzelius/Slurm scripts.
1073
+ - Do not modify algorithms/worldmem.
1074
+ - Keep the packed order [target][anchor][dynamic][revisit].
1075
+ - Keep the stream name dynamic; dynamic now semantically means multiview memory for the new policy.
1076
+ - Add only dynamic.selection_policy=multiview. Do not add fov_multiview or fov_fps policy aliases.
1077
+ - Keep existing recent and event_triggered dynamic policies working.
1078
+ - Do not change the DeMemWM model architecture or FrameMemoryReferenceAttention in this plan.
1079
+ - Keep anchor selection semantics unchanged.
1080
+ - Make FOV candidate sampling deterministic.
1081
+ - Implement exact FOV-greedy multiview selection for training and validation/test.
1082
+ - Add only one multiview backend knob: dynamic.multiview_selector with values fov_greedy and pose_plucker_fps.
1083
+ - Do not add hand-tuned scoring weight configs.
1084
+ - Do not build exact FOV masks twice for the same revisit+dynamic query; share a local FOV pool/cache.
1085
+ - Add a local selector-speed benchmark script comparing fov_greedy and pose_plucker_fps.
1086
+ - Use clean memory in the new local launchers.
1087
+ - Add local 2H200 debug and local full launchers; local full should default to 2 GPUs and scale through environment variables. No SBATCH, srun, module load, berzelius, or hard-coded cluster paths.
1088
+ - Preserve unrelated user changes.
1089
+
1090
+ Final report must include changed files, commits made, validation commands and results, selector benchmark command/results if run, local debug/full run commands, and remaining risks.
1091
+ ```
configurations/dataset/video_minecraft_dememwm_latent.yaml CHANGED
@@ -31,8 +31,11 @@ memory_selection:
31
  local_context_exclusion_frames: 8
32
  plucker_moment_radius: 30.0
33
  anchor_diverse_selection: true
 
 
34
  dynamic:
35
  selection_policy: event_triggered
 
36
  scene_threshold: 2.5
37
  state_threshold: 2.5
38
  stable_threshold: 1.0
 
31
  local_context_exclusion_frames: 8
32
  plucker_moment_radius: 30.0
33
  anchor_diverse_selection: true
34
+ pose_preselect_topk: 64
35
+ candidate_chunk_size: 64
36
  dynamic:
37
  selection_policy: event_triggered
38
+ multiview_selector: fov_greedy
39
  scene_threshold: 2.5
40
  state_threshold: 2.5
41
  stable_threshold: 1.0
datasets/video/memory_selection.py CHANGED
@@ -9,6 +9,8 @@ import torch
9
 
10
 
11
  SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
 
 
12
  _FOV_NUM_POINTS = 10000
13
  _FOV_RADIUS = 30.0
14
  _FOV_HALF_H = 105.0 / 2.0
@@ -384,7 +386,20 @@ def _select_dynamic(target_start: int, count: int, min_candidate_frame: int = 0)
384
 
385
  def _dynamic_policy(cfg) -> str:
386
  dynamic_cfg = cfg_get(cfg, "dynamic", {})
387
- return str(cfg_get(dynamic_cfg, "selection_policy", "recent"))
 
 
 
 
 
 
 
 
 
 
 
 
 
388
 
389
 
390
  def _latent_frame_vectors(latents, start: int, stop: int):
@@ -780,7 +795,12 @@ def _select_dynamic_by_policy(
780
  reference_frames=reference_frames,
781
  excluded=excluded,
782
  )
783
- raise ValueError(f"memory_selection.dynamic.selection_policy must be recent or event_triggered; got {policy!r}")
 
 
 
 
 
784
 
785
 
786
  def _select_revisit(
@@ -863,8 +883,11 @@ def select_memory_indices(
863
 
864
  target_start = int(target_positions[0])
865
  policy = _dynamic_policy(cfg)
866
- if policy not in {"recent", "event_triggered"}:
867
- raise ValueError(f"memory_selection.dynamic.selection_policy must be recent or event_triggered; got {policy!r}")
 
 
 
868
 
869
  if policy == "recent":
870
  dynamic = _select_dynamic_by_policy(
 
9
 
10
 
11
  SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
12
+ _VALID_DYNAMIC_POLICIES = ("recent", "event_triggered", "multiview")
13
+ _VALID_MULTIVIEW_SELECTORS = ("fov_greedy", "pose_plucker_fps")
14
  _FOV_NUM_POINTS = 10000
15
  _FOV_RADIUS = 30.0
16
  _FOV_HALF_H = 105.0 / 2.0
 
386
 
387
  def _dynamic_policy(cfg) -> str:
388
  dynamic_cfg = cfg_get(cfg, "dynamic", {})
389
+ policy = str(cfg_get(dynamic_cfg, "selection_policy", "recent"))
390
+ if policy not in _VALID_DYNAMIC_POLICIES:
391
+ valid = ", ".join(_VALID_DYNAMIC_POLICIES)
392
+ raise ValueError(f"memory_selection.dynamic.selection_policy must be one of {valid}; got {policy!r}")
393
+ return policy
394
+
395
+
396
+ def _dynamic_multiview_selector(cfg) -> str:
397
+ dynamic_cfg = cfg_get(cfg, "dynamic", {})
398
+ selector = str(cfg_get(dynamic_cfg, "multiview_selector", "fov_greedy"))
399
+ if selector not in _VALID_MULTIVIEW_SELECTORS:
400
+ valid = ", ".join(_VALID_MULTIVIEW_SELECTORS)
401
+ raise ValueError(f"memory_selection.dynamic.multiview_selector must be one of {valid}; got {selector!r}")
402
+ return selector
403
 
404
 
405
  def _latent_frame_vectors(latents, start: int, stop: int):
 
795
  reference_frames=reference_frames,
796
  excluded=excluded,
797
  )
798
+ if policy == "multiview":
799
+ selector = _dynamic_multiview_selector(cfg)
800
+ raise NotImplementedError(
801
+ f"memory_selection.dynamic.selection_policy='multiview' with multiview_selector={selector!r} is not implemented yet"
802
+ )
803
+ raise AssertionError(f"unhandled dynamic selection policy: {policy!r}")
804
 
805
 
806
  def _select_revisit(
 
883
 
884
  target_start = int(target_positions[0])
885
  policy = _dynamic_policy(cfg)
886
+ if policy == "multiview":
887
+ selector = _dynamic_multiview_selector(cfg)
888
+ raise NotImplementedError(
889
+ f"memory_selection.dynamic.selection_policy='multiview' with multiview_selector={selector!r} is not implemented yet"
890
+ )
891
 
892
  if policy == "recent":
893
  dynamic = _select_dynamic_by_policy(
tests/test_dememwm_latent_dataset.py CHANGED
@@ -9,13 +9,19 @@ import torch
9
  from torch import nn
10
  from omegaconf import OmegaConf
11
 
12
- from datasets.video.memory_selection import _memory_candidate_frames, select_memory_indices
 
 
 
 
 
13
  from datasets.video.minecraft_video_dememwm_latent_dataset import MinecraftVideoDeMemWMLatentDataset
14
 
15
 
16
  def _dynamic_cfg(selection_policy="recent", **overrides):
17
  cfg = {
18
  "selection_policy": selection_policy,
 
19
  "scene_threshold": 2.5,
20
  "state_threshold": 2.5,
21
  "stable_threshold": 1.0,
@@ -171,6 +177,10 @@ class MemorySelectionTests(unittest.TestCase):
171
  cfg = OmegaConf.load(config_path)
172
 
173
  self.assertIs(cfg.memory_selection.causal, True)
 
 
 
 
174
 
175
  def test_revisit_uses_sampled_fov_selection(self):
176
  poses = np.array(
@@ -597,17 +607,59 @@ class MemorySelectionTests(unittest.TestCase):
597
  self.assertEqual(indices["dynamic"].tolist(), [-1, -1])
598
  self.assertEqual(masks["dynamic"].tolist(), [False, False])
599
 
600
- def test_hybrid_policy_is_rejected(self):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
601
  poses = np.zeros((8, 5), dtype=np.float32)
602
  cfg = _selection_cfg(
603
  max_anchor_frames=0,
604
  max_dynamic_frames=2,
605
  max_revisit_frames=0,
606
- dynamic=_dynamic_cfg("hybrid"),
607
  )
608
 
609
- with self.assertRaisesRegex(ValueError, "recent or event_triggered"):
610
- select_memory_indices(poses, np.array([6]), cfg, latents=_event_latents(8, event_frame=99))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
611
 
612
  def test_event_dynamic_uses_nearest_anchors_per_revisit_frame(self):
613
  import datasets.video.memory_selection as memory_selection
@@ -1191,7 +1243,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
1191
  self.decoded_inputs.append(xs.detach().clone())
1192
  return xs + 10.0
1193
 
1194
- def _update_metric_accumulators(self, xs_pred, xs_gt):
1195
  self.metric_updates.append((xs_pred.detach().clone(), xs_gt.detach().clone()))
1196
 
1197
  def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
@@ -1246,12 +1298,12 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
1246
  self.assertEqual([call["kwargs"]["current_frame"] for call in calls], [2, 2, 4, 4])
1247
  self.assertTrue(torch.equal(calls[0]["kwargs"]["frame_idx"], calls[1]["kwargs"]["frame_idx"]))
1248
  self.assertTrue(torch.equal(calls[2]["kwargs"]["frame_idx"], calls[3]["kwargs"]["frame_idx"]))
1249
- self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0, 0.0])
1250
  self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [2.0, 0.0, 0.0, 2.0, 2.0])
1251
- self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0, 0.0])
1252
- self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11, 0])
1253
  self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 12])
1254
- self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[False]])
1255
  self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
1256
  self.assertAlmostEqual(float(loss), 1.0 / 3.0)
1257
  self.assertEqual([name for name, _ in harness.logged], ["test/latent_mse"])
@@ -1316,7 +1368,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
1316
  def decode(self, xs):
1317
  return xs
1318
 
1319
- def _update_metric_accumulators(self, xs_pred, xs_gt):
1320
  pass
1321
 
1322
  def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
@@ -1481,6 +1533,38 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
1481
  self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 100, 103])
1482
  self.assertEqual(sample["memory_masks"]["dynamic"].tolist(), [True, True])
1483
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1484
  def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
1485
  with tempfile.TemporaryDirectory() as tmp:
1486
  root = Path(tmp)
@@ -1504,7 +1588,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
1504
  self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
1505
  self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
1506
  revisit_frames = sample["frame_indices"][7:9].numpy()
1507
- self.assertEqual(revisit_frames.tolist(), [104, 105])
1508
  self.assertEqual(len(np.unique(revisit_frames)), 2)
1509
  self.assertTrue(sample["memory_masks"]["target"].all().item())
1510
  self.assertTrue(sample["memory_masks"]["anchor"].all().item())
 
9
  from torch import nn
10
  from omegaconf import OmegaConf
11
 
12
+ from datasets.video.memory_selection import (
13
+ _dynamic_multiview_selector,
14
+ _dynamic_policy,
15
+ _memory_candidate_frames,
16
+ select_memory_indices,
17
+ )
18
  from datasets.video.minecraft_video_dememwm_latent_dataset import MinecraftVideoDeMemWMLatentDataset
19
 
20
 
21
  def _dynamic_cfg(selection_policy="recent", **overrides):
22
  cfg = {
23
  "selection_policy": selection_policy,
24
+ "multiview_selector": "fov_greedy",
25
  "scene_threshold": 2.5,
26
  "state_threshold": 2.5,
27
  "stable_threshold": 1.0,
 
177
  cfg = OmegaConf.load(config_path)
178
 
179
  self.assertIs(cfg.memory_selection.causal, True)
180
+ self.assertEqual(cfg.memory_selection.dynamic.selection_policy, "event_triggered")
181
+ self.assertEqual(cfg.memory_selection.dynamic.multiview_selector, "fov_greedy")
182
+ self.assertEqual(cfg.memory_selection.pose_preselect_topk, 64)
183
+ self.assertEqual(cfg.memory_selection.candidate_chunk_size, 64)
184
 
185
  def test_revisit_uses_sampled_fov_selection(self):
186
  poses = np.array(
 
607
  self.assertEqual(indices["dynamic"].tolist(), [-1, -1])
608
  self.assertEqual(masks["dynamic"].tolist(), [False, False])
609
 
610
+ def test_dynamic_policy_validation_accepts_exact_supported_values(self):
611
+ for policy in ("recent", "event_triggered", "multiview"):
612
+ with self.subTest(policy=policy):
613
+ cfg = _selection_cfg(dynamic=_dynamic_cfg(policy))
614
+
615
+ self.assertEqual(_dynamic_policy(cfg), policy)
616
+
617
+ def test_unknown_dynamic_policies_are_rejected_without_aliases(self):
618
+ poses = np.zeros((8, 5), dtype=np.float32)
619
+ for policy in ("fov_multiview", "fov_fps", "hybrid"):
620
+ with self.subTest(policy=policy):
621
+ cfg = _selection_cfg(
622
+ max_anchor_frames=0,
623
+ max_dynamic_frames=2,
624
+ max_revisit_frames=0,
625
+ dynamic=_dynamic_cfg(policy),
626
+ )
627
+
628
+ with self.assertRaisesRegex(ValueError, "recent, event_triggered, multiview"):
629
+ select_memory_indices(poses, np.array([6]), cfg, latents=_event_latents(8, event_frame=99))
630
+
631
+ def test_multiview_selector_validation_accepts_exact_backends(self):
632
+ for selector in ("fov_greedy", "pose_plucker_fps"):
633
+ with self.subTest(selector=selector):
634
+ cfg = _selection_cfg(dynamic=_dynamic_cfg("multiview", multiview_selector=selector))
635
+
636
+ self.assertEqual(_dynamic_multiview_selector(cfg), selector)
637
+
638
+ def test_unknown_multiview_selector_is_rejected(self):
639
  poses = np.zeros((8, 5), dtype=np.float32)
640
  cfg = _selection_cfg(
641
  max_anchor_frames=0,
642
  max_dynamic_frames=2,
643
  max_revisit_frames=0,
644
+ dynamic=_dynamic_cfg("multiview", multiview_selector="fov_fps"),
645
  )
646
 
647
+ with self.assertRaisesRegex(ValueError, "fov_greedy, pose_plucker_fps"):
648
+ select_memory_indices(poses, np.array([6]), cfg)
649
+
650
+ def test_multiview_policy_selection_fails_explicitly_until_implemented(self):
651
+ poses = np.zeros((8, 5), dtype=np.float32)
652
+ for selector in ("fov_greedy", "pose_plucker_fps"):
653
+ with self.subTest(selector=selector):
654
+ cfg = _selection_cfg(
655
+ max_anchor_frames=0,
656
+ max_dynamic_frames=2,
657
+ max_revisit_frames=0,
658
+ dynamic=_dynamic_cfg("multiview", multiview_selector=selector),
659
+ )
660
+
661
+ with self.assertRaisesRegex(NotImplementedError, "multiview"):
662
+ select_memory_indices(poses, np.array([6]), cfg)
663
 
664
  def test_event_dynamic_uses_nearest_anchors_per_revisit_frame(self):
665
  import datasets.video.memory_selection as memory_selection
 
1243
  self.decoded_inputs.append(xs.detach().clone())
1244
  return xs + 10.0
1245
 
1246
+ def _update_metric_accumulators(self, xs_pred, xs_gt, valid_mask=None, eval_start=0):
1247
  self.metric_updates.append((xs_pred.detach().clone(), xs_gt.detach().clone()))
1248
 
1249
  def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
 
1298
  self.assertEqual([call["kwargs"]["current_frame"] for call in calls], [2, 2, 4, 4])
1299
  self.assertTrue(torch.equal(calls[0]["kwargs"]["frame_idx"], calls[1]["kwargs"]["frame_idx"]))
1300
  self.assertTrue(torch.equal(calls[2]["kwargs"]["frame_idx"], calls[3]["kwargs"]["frame_idx"]))
1301
+ self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0])
1302
  self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [2.0, 0.0, 0.0, 2.0, 2.0])
1303
+ self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0])
1304
+ self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11])
1305
  self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [13, 14, 10, 12, 12])
1306
+ self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[]])
1307
  self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
1308
  self.assertAlmostEqual(float(loss), 1.0 / 3.0)
1309
  self.assertEqual([name for name, _ in harness.logged], ["test/latent_mse"])
 
1368
  def decode(self, xs):
1369
  return xs
1370
 
1371
+ def _update_metric_accumulators(self, xs_pred, xs_gt, valid_mask=None, eval_start=0):
1372
  pass
1373
 
1374
  def _update_per_frame_metric_accumulators(self, xs_pred, xs_gt, mask, eval_start):
 
1533
  self.assertEqual(sample["frame_indices"].tolist(), [106, 107, 108, 100, 103])
1534
  self.assertEqual(sample["memory_masks"]["dynamic"].tolist(), [True, True])
1535
 
1536
+ def test_dataset_does_not_build_dynamic_stream_for_multiview_policy(self):
1537
+ import datasets.video.minecraft_video_dememwm_latent_dataset as dataset_module
1538
+
1539
+ with tempfile.TemporaryDirectory() as tmp:
1540
+ root = Path(tmp)
1541
+ num_frames = 112
1542
+ _write_vae_feature_clip(
1543
+ root,
1544
+ root / "vae_features",
1545
+ stem="sample",
1546
+ num_frames=num_frames,
1547
+ actions=np.zeros((num_frames, 25), dtype=np.float32),
1548
+ latents=_event_latents(num_frames, event_frame=102),
1549
+ )
1550
+ dataset = MinecraftVideoDeMemWMLatentDataset(
1551
+ _dataset_cfg(
1552
+ root,
1553
+ context_length=6,
1554
+ memory_selection=_selection_cfg(
1555
+ max_anchor_frames=0,
1556
+ max_dynamic_frames=2,
1557
+ max_revisit_frames=0,
1558
+ dynamic=_dynamic_cfg("multiview"),
1559
+ ),
1560
+ ),
1561
+ split="training",
1562
+ )
1563
+ with mock.patch.object(dataset_module, "_build_dynamic_stream", side_effect=AssertionError("unexpected dynamic stream build")):
1564
+ arrays = dataset._load_feature_arrays(dataset.data_paths[0])
1565
+
1566
+ self.assertEqual(arrays["dynamic_stream"].tolist(), [])
1567
+
1568
  def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
1569
  with tempfile.TemporaryDirectory() as tmp:
1570
  root = Path(tmp)
 
1588
  self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
1589
  self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
1590
  revisit_frames = sample["frame_indices"][7:9].numpy()
1591
+ self.assertEqual(revisit_frames.tolist(), [102, 103])
1592
  self.assertEqual(len(np.unique(revisit_frames)), 2)
1593
  self.assertTrue(sample["memory_masks"]["target"].all().item())
1594
  self.assertTrue(sample["memory_masks"]["anchor"].all().item())