Auto-sync: 2026-06-30 15:00:21 (part 4)

tests/test_maniskill_policy_rollout.py

@@ -26,12 +26,14 @@ from dovla_cil.eval.maniskill_policy_rollout import (
     _lattice_candidate_mask,
     _lattice_candidate_type_bonus,
     _load_state_archive,
+    _normalize_residual_candidate_type_name,
     _numeric_action_values,
     _reduce_residual_candidates_by_type,
     _select_action_chunk,
     _selected_candidate_type,
     _selected_residual_scale,
     _summarize_rows,
+    _task_limited_exclude_types,
     _task_limited_challenger_mask,
 )
 
@@ -675,6 +677,19 @@ def test_lattice_candidate_include_mask_uses_exact_candidate_types() -> None:
     assert mask.tolist() == [[False, True, False]]
 
 
+def test_residual_candidate_type_normalizer_accepts_shorthand() -> None:
+    assert _normalize_residual_candidate_type_name("near_miss") == "residual_near_miss"
+    assert _normalize_residual_candidate_type_name("policy") == "policy_residual"
+    assert (
+        _normalize_residual_candidate_type_name("retrieval_residual_residual_no_op")
+        == "residual_no_op"
+    )
+    assert (
+        _normalize_residual_candidate_type_name("near_miss+no_op")
+        == "residual_near_miss+residual_no_op"
+    )
+
+
 def test_task_limited_challenger_mask_respects_task_gate() -> None:
     import torch
 
@@ -719,6 +734,66 @@ def test_task_limited_challenger_mask_respects_task_gate() -> None:
     assert unrestricted.tolist() == [[False, True, True, False]]
 
 
+def test_task_limited_challenger_mask_can_limit_individual_types() -> None:
+    import torch
+
+    case = SimpleNamespace(
+        candidate_types=[
+            "policy_residual",
+            "residual_near_miss",
+            "residual_wrong_gripper",
+            "residual_near_miss+residual_wrong_gripper",
+        ]
+    )
+
+    pick_mask = _task_limited_challenger_mask(
+        [case],
+        task_id="PickCube-v1",
+        torch=torch,
+        device="cpu",
+        include_types=("near_miss", "wrong_gripper"),
+        include_tasks=(),
+        include_type_tasks={"wrong_gripper": ("PickCube-v1", "PullCube-v1")},
+    )
+    lift_mask = _task_limited_challenger_mask(
+        [case],
+        task_id="LiftPegUpright-v1",
+        torch=torch,
+        device="cpu",
+        include_types=("near_miss", "wrong_gripper"),
+        include_tasks=(),
+        include_type_tasks={"wrong_gripper": ("PickCube-v1", "PullCube-v1")},
+    )
+
+    assert pick_mask is not None
+    assert lift_mask is not None
+    assert pick_mask.tolist() == [[False, True, True, False]]
+    assert lift_mask.tolist() == [[False, True, False, False]]
+
+
+def test_task_limited_exclude_types_adds_task_specific_components() -> None:
+    base = ("residual_random_negative", "residual_wrong_direction")
+    mapping = {"residual_wrong_gripper": ("StackCube-v1",)}
+
+    stack_excludes = _task_limited_exclude_types(
+        base,
+        task_id="StackCube-v1",
+        exclude_type_tasks=mapping,
+    )
+    pick_excludes = _task_limited_exclude_types(
+        base,
+        task_id="PickCube-v1",
+        exclude_type_tasks=mapping,
+    )
+
+    assert stack_excludes == (
+        "residual_random_negative",
+        "residual_wrong_direction",
+        "residual_wrong_gripper",
+    )
+    assert pick_excludes == base
+
+
 def test_expand_residual_lattice_mask_can_limit_challenger_scales() -> None:
     import torch