Auto-sync: 2026-06-27 10:26:03 (part 3)

tests/test_maniskill_policy_rollout.py

@@ -265,3 +265,26 @@ def test_lattice_mode_respects_candidate_mask() -> None:
 
     assert torch.allclose(actions, mean)
     assert index.tolist() == [1]
+
+
+def test_retrieval_lattice_mode_uses_candidate_tensor() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 3)
+    offset = torch.full_like(mean, 0.4)
+    model = _StubModel(torch, mean, best_offset=offset)
+    candidates = torch.stack([mean, mean + offset], dim=1)
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="retrieval_lattice",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=0,
+        action_candidates=candidates,
+    )
+
+    assert torch.allclose(actions, mean + offset)
+    assert index.tolist() == [1]