Auto-sync: 2026-06-30 20:06:30 (part 3)

scripts/slurm/summarize_h16_policy_ckpt.sbatch

@@ -112,6 +112,9 @@ for result_path in sorted(base_dir.glob(f"seed_*/{out_name}")):
             "retrieval_residual_challenger_margin": data.get(
                 "retrieval_residual_challenger_margin", 0.0
             ),
+            "retrieval_residual_challenger_type_margins": data.get(
+                "retrieval_residual_challenger_type_margins", {}
+            ),
             "retrieval_residual_challenger_tasks": data.get(
                 "retrieval_residual_challenger_tasks", []
             ),
@@ -406,6 +409,16 @@ for row in rows:
         else ""
     )
     challenger_type_tasks = row.get("retrieval_residual_challenger_type_tasks") or {}
+    challenger_type_margins = row.get("retrieval_residual_challenger_type_margins") or {}
+    challenger_type_margin_suffix = (
+        " margins "
+        + ",".join(
+            f"{candidate_type}={float(margin):.2f}"
+            for candidate_type, margin in sorted(challenger_type_margins.items())
+        )
+        if challenger_type_margins
+        else ""
+    )
     challenger_type_task_suffix = (
         " with "
         + ";".join(
@@ -454,6 +467,7 @@ for row in rows:
                 )
                 + challenger_task_suffix
                 + challenger_type_task_suffix
+                + challenger_type_margin_suffix
                 + f"@{row.get('retrieval_residual_challenger_margin') or 0.0:.2f}"
                 if row.get("retrieval_residual_challenger_types")
                 else ""

tests/test_maniskill_policy_rollout.py

@@ -401,6 +401,36 @@ def test_retrieval_residual_challenger_can_override_primary_anchor() -> None:
     assert index.tolist() == [2]
 
 
+def test_retrieval_residual_challenger_can_use_candidate_margins() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 1)
+    offset = torch.full_like(mean, 0.4)
+    model = _StubModel(torch, mean, best_offset=offset)
+    residuals = torch.tensor([[[[0.0]], [[0.1]], [[0.4]]]], dtype=torch.float32)
+
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="retrieval_residual",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=0,
+        selection_margin=1.0,
+        action_candidates=residuals,
+        challenger_mask=torch.tensor([[False, False, True]]),
+        challenger_margin=0.05,
+        retrieval_residual_challenger_margin_by_candidate=torch.tensor(
+            [[0.05, 0.05, 0.2]]
+        ),
+    )
+
+    assert torch.allclose(actions, mean)
+    assert index.tolist() == [0]
+
+
 def test_candidate_oracle_prefix_keeps_deployed_candidate_first() -> None:
     import torch