Auto-sync: 2026-06-29 03:12:48 (part 3)

scripts/eval_maniskill_policy_rollout.py

@@ -150,6 +150,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-action-l2-penalty",
+        type=float,
+        default=0.0,
+        help="Penalty scale subtracted from residual candidate field scores according to "
+        "mean squared action displacement from the policy mean.",
+    )
     parser.add_argument(
         "--retrieval-residual-scale",
         type=float,
@@ -248,6 +255,7 @@ def main(argv: list[str] | None = None) -> int:
         retrieval_residual_source_score_bonus_scale=(
             args.retrieval_residual_source_score_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,
         retrieval_residual_anchor=args.retrieval_residual_anchor,

scripts/slurm/eval_maniskill_policy_rollout.sbatch

@@ -56,6 +56,7 @@ RETRIEVAL_TYPE_MIN_SUCCESS="${RETRIEVAL_TYPE_MIN_SUCCESS:-0.0}"
 RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS="${RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS:-0.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_ACTION_L2_PENALTY="${RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY:-0.0}"
 RETRIEVAL_RESIDUAL_SCALE="${RETRIEVAL_RESIDUAL_SCALE:-1.0}"
 RETRIEVAL_RESIDUAL_SCALES="${RETRIEVAL_RESIDUAL_SCALES:-}"
 if [[ -n "${RETRIEVAL_RESIDUAL_SCALES_COLON:-}" ]]; then
@@ -124,6 +125,7 @@ apptainer exec --nv \
     --retrieval-residual-min-source-progress "$RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS" \
     --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-action-l2-penalty "$RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY" \
     --retrieval-residual-scale "$RETRIEVAL_RESIDUAL_SCALE" \
     --retrieval-residual-scales "$RETRIEVAL_RESIDUAL_SCALES" \
     --retrieval-residual-anchor "$RETRIEVAL_RESIDUAL_ANCHOR" \

scripts/slurm/eval_maniskill_policy_rollout_cpu_smoke.sbatch

@@ -55,6 +55,7 @@ RETRIEVAL_TYPE_MIN_SUCCESS="${RETRIEVAL_TYPE_MIN_SUCCESS:-0.0}"
 RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS="${RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS:-0.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_ACTION_L2_PENALTY="${RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY:-0.0}"
 RETRIEVAL_RESIDUAL_SCALE="${RETRIEVAL_RESIDUAL_SCALE:-1.0}"
 RETRIEVAL_RESIDUAL_SCALES="${RETRIEVAL_RESIDUAL_SCALES:-}"
 if [[ -n "${RETRIEVAL_RESIDUAL_SCALES_COLON:-}" ]]; then
@@ -120,6 +121,7 @@ apptainer exec \
     --retrieval-residual-min-source-progress "$RETRIEVAL_RESIDUAL_MIN_SOURCE_PROGRESS" \
     --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-action-l2-penalty "$RETRIEVAL_RESIDUAL_ACTION_L2_PENALTY" \
     --retrieval-residual-scale "$RETRIEVAL_RESIDUAL_SCALE" \
     --retrieval-residual-scales "$RETRIEVAL_RESIDUAL_SCALES" \
     --retrieval-residual-anchor "$RETRIEVAL_RESIDUAL_ANCHOR" \

tests/test_maniskill_policy_rollout.py

@@ -615,6 +615,32 @@ def test_retrieval_residual_candidate_type_bonus_repeats_over_scale_grid() -> No
     assert index.tolist() == [1]
 
 
+def test_retrieval_residual_action_l2_penalty_prefers_shorter_tangent() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 1)
+    residual = torch.ones_like(mean)
+    model = _StubModel(torch, mean, best_offset=0.5 * residual)
+    residuals = torch.stack([torch.zeros_like(mean), residual], dim=1)
+    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,
+        action_candidates=residuals,
+        retrieval_residual_scale=1.0,
+        retrieval_residual_scales=(0.25, 0.5),
+        retrieval_residual_action_l2_penalty=2.0,
+    )
+
+    assert torch.allclose(actions, mean + 0.25 * residual)
+    assert index.tolist() == [1]
+
+
 def test_retrieval_residual_field_softmax_builds_field_weighted_tangent() -> None:
     import torch