Auto-sync: 2026-06-30 09:56:41 (part 3)

tests/test_maniskill_policy_rollout.py

@@ -226,10 +226,14 @@ class _StubModel:
         distance = ((action - target) ** 2).reshape(action.shape[0], -1).sum(dim=1)
         return {"potential": -distance}
 
-    def forward_proposals(self, observation, instruction):
+    def forward_proposals(self, observation, instruction, proposal_types=None):
         del observation, instruction
         if self._proposals is None:
             raise AssertionError("stub proposals were not configured")
+        if proposal_types:
+            # Test stubs use proposal type names p0/p1/... to request a subset.
+            indices = [int(str(proposal_type)[1:]) for proposal_type in proposal_types]
+            return self._proposals[:, indices]
         return self._proposals
 
 
@@ -318,6 +322,30 @@ def test_proposal_lattice_mode_scores_model_generated_proposals() -> None:
     assert index.tolist() == [1]
 
 
+def test_proposal_lattice_mode_can_request_type_subset() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 2)
+    proposals = torch.tensor([[[[0.1, 0.0]], [[0.4, 0.4]], [[0.0, 0.1]]]], dtype=torch.float32)
+    offset = torch.tensor([[[0.4, 0.4]]], dtype=torch.float32)
+    model = _StubModel(torch, mean, best_offset=offset, proposals=proposals)
+
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="proposal_lattice",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=0,
+        proposal_lattice_types=("p0", "p2"),
+    )
+
+    assert torch.allclose(actions, proposals[:, 0])
+    assert index.tolist() == [0]
+
+
 def test_retrieval_residual_margin_can_abstain_to_policy() -> None:
     import torch