Auto-sync: 2026-06-29 06:36:52 (part 2)

scripts/eval_maniskill_policy_rollout.py

@@ -136,6 +136,13 @@ def main(argv: list[str] | None = None) -> int:
         help="Minimum measured train-source progress for an individual residual candidate to "
         "be eligible in retrieval_residual mode. The policy_residual fallback is always kept.",
     )
+    parser.add_argument(
+        "--retrieval-residual-min-source-advantage",
+        type=float,
+        default=-1.0e9,
+        help="Minimum source reward-score advantage over the source anchor for an individual "
+        "residual candidate. The default disables advantage gating.",
+    )
     parser.add_argument(
         "--retrieval-residual-source-progress-bonus-scale",
         type=float,
@@ -150,6 +157,13 @@ def main(argv: list[str] | None = None) -> int:
         help="Scale for adding a train-source reward-score prior to each retrieved residual "
         "candidate before field selection. Score is progress plus terminal success.",
     )
+    parser.add_argument(
+        "--retrieval-residual-source-advantage-bonus-scale",
+        type=float,
+        default=0.0,
+        help="Scale for adding a train-source reward-score advantage prior relative to the "
+        "source residual anchor action.",
+    )
     parser.add_argument(
         "--retrieval-residual-action-l2-penalty",
         type=float,
@@ -249,12 +263,16 @@ def main(argv: list[str] | None = None) -> int:
         retrieval_metric=args.retrieval_metric,
         retrieval_type_min_success=args.retrieval_type_min_success,
         retrieval_residual_min_source_progress=args.retrieval_residual_min_source_progress,
+        retrieval_residual_min_source_advantage=args.retrieval_residual_min_source_advantage,
         retrieval_residual_source_progress_bonus_scale=(
             args.retrieval_residual_source_progress_bonus_scale
         ),
         retrieval_residual_source_score_bonus_scale=(
             args.retrieval_residual_source_score_bonus_scale
         ),
+        retrieval_residual_source_advantage_bonus_scale=(
+            args.retrieval_residual_source_advantage_bonus_scale
+        ),
         retrieval_residual_action_l2_penalty=args.retrieval_residual_action_l2_penalty,
         retrieval_residual_scale=args.retrieval_residual_scale,
         retrieval_residual_scales=retrieval_residual_scales,

scripts/slurm/eval_maniskill_policy_rollout.sbatch

@@ -54,8 +54,10 @@ RETRIEVAL_NEIGHBORS="${RETRIEVAL_NEIGHBORS:-1}"
 RETRIEVAL_METRIC="${RETRIEVAL_METRIC:-raw}"
 RETRIEVAL_TYPE_MIN_SUCCESS="${RETRIEVAL_TYPE_MIN_SUCCESS:-0.0}"
 RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS="${RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS:-0.0}"
+RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE="${RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE:--1000000000.0}"
 RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE:-0.0}"
 RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE:-0.0}"
+RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE:-0.0}"
 RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY="${RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY:-0.0}"
 RETRIEVAL_RESIDUAL_SCALE="${RETRIEVAL_RESIDUAL_SCALE:-1.0}"
 RETRIEVAL_RESIDUAL_SCALES="${RETRIEVAL_RESIDUAL_SCALES:-}"
@@ -123,8 +125,10 @@ apptainer exec --nv \
     --retrieval-metric "$RETRIEVAL_METRIC" \
     --retrieval-type-min-success "$RETRIEVAL_TYPE_MIN_SUCCESS" \
     --retrieval-residual-min-source-progress "$RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS" \
+    --retrieval-residual-min-source-advantage "$RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE" \
     --retrieval-residual-source-progress-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE" \
     --retrieval-residual-source-score-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE" \
+    --retrieval-residual-source-advantage-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE" \
     --retrieval-residual-action-l2-penalty "$RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY" \
     --retrieval-residual-scale "$RETRIEVAL_RESIDUAL_SCALE" \
     --retrieval-residual-scales "$RETRIEVAL_RESIDUAL_SCALES" \

scripts/slurm/eval_maniskill_policy_rollout_cpu_smoke.sbatch

@@ -53,8 +53,10 @@ RETRIEVAL_NEIGHBORS="${RETRIEVAL_NEIGHBORS:-1}"
 RETRIEVAL_METRIC="${RETRIEVAL_METRIC:-raw}"
 RETRIEVAL_TYPE_MIN_SUCCESS="${RETRIEVAL_TYPE_MIN_SUCCESS:-0.0}"
 RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS="${RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS:-0.0}"
+RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE="${RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE:--1000000000.0}"
 RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE:-0.0}"
 RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE:-0.0}"
+RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE="${RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE:-0.0}"
 RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY="${RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY:-0.0}"
 RETRIEVAL_RESIDUAL_SCALE="${RETRIEVAL_RESIDUAL_SCALE:-1.0}"
 RETRIEVAL_RESIDUAL_SCALES="${RETRIEVAL_RESIDUAL_SCALES:-}"
@@ -119,8 +121,10 @@ apptainer exec \
     --retrieval-metric "$RETRIEVAL_METRIC" \
     --retrieval-type-min-success "$RETRIEVAL_TYPE_MIN_SUCCESS" \
     --retrieval-residual-min-source-progress "$RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS" \
+    --retrieval-residual-min-source-advantage "$RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE" \
     --retrieval-residual-source-progress-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_PROGRESS_BONUS_SCALE" \
     --retrieval-residual-source-score-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_SCALE" \
+    --retrieval-residual-source-advantage-bonus-scale "$RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_BONUS_SCALE" \
     --retrieval-residual-action-l2-penalty "$RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY" \
     --retrieval-residual-scale "$RETRIEVAL_RESIDUAL_SCALE" \
     --retrieval-residual-scales "$RETRIEVAL_RESIDUAL_SCALES" \

tests/test_maniskill_policy_rollout.py

@@ -1188,6 +1188,75 @@ def test_retrieval_residual_source_progress_threshold_filters_individual_residua
     )
 
 
+def test_retrieval_residual_source_advantage_filters_against_anchor() -> None:
+    def record(
+        group_id: str,
+        candidate_type: str,
+        action_value: float,
+        progress: float,
+        terminal_success: bool = False,
+    ):
+        return SimpleNamespace(
+            group_id=group_id,
+            task_id="PickCube-v1",
+            candidate_type=candidate_type,
+            record_id=f"{group_id}-{candidate_type}-{action_value}",
+            observation_inline={"features": [0.0, 0.0]},
+            reward=SimpleNamespace(progress=progress, terminal_success=terminal_success),
+            action_chunk=ActionChunk(
+                representation="continuous",
+                horizon=1,
+                values=[[action_value, 0.0]],
+            ),
+        )
+
+    groups = {
+        "train_a": [
+            record("train_a", "expert", 1.0, 0.4),
+            record("train_a", "no_op", 1.2, 0.6),
+            record("train_a", "wrong_gripper", 1.4, 0.3),
+        ],
+        "heldout": [record("heldout", "expert", 9.0, 1.0, terminal_success=True)],
+    }
+    dataset = SimpleNamespace(
+        group_ids=list(groups),
+        get_group=lambda group_id: groups[group_id],
+    )
+    case = _RolloutCase(
+        group_id="heldout",
+        task_id="PickCube-v1",
+        source_dataset=Path("."),
+        state={},
+        observation={"features": [0.0, 0.0]},
+        instruction="pick",
+        oracle_score=1.0,
+        oracle_success=True,
+        expert_score=1.0,
+        expert_success=True,
+        best_action_values=[[9.9, 0.0]],
+        candidate_action_values=[],
+        candidate_types=[],
+    )
+
+    [attached] = _attach_retrieved_residual_candidates(
+        dataset,
+        [case],
+        heldout_group_ids=["heldout"],
+        obs_dim=2,
+        observation_mode="state",
+        retrieval_neighbors=1,
+        retrieval_residual_min_source_advantage=0.0,
+        retrieval_residual_source_advantage_bonus_scale=0.5,
+    )
+
+    assert attached.candidate_types == ["policy_residual", "residual_no_op"]
+    assert np.allclose(attached.candidate_score_bonuses, [0.0, 0.1])
+    assert np.allclose(
+        np.asarray(attached.candidate_action_values, dtype=np.float32),
+        np.asarray([[[0.0, 0.0]], [[0.2, 0.0]]]),
+    )
+
+
 def test_retrieval_residual_source_progress_padding_keeps_batches_rectangular() -> None:
     def record(
         group_id: str,