Auto-sync: 2026-06-27 10:45:23 (part 2)

scripts/slurm/train_dovla_h16_policy_ckpt.sbatch

@@ -25,6 +25,7 @@ CONTAINER_PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
 DATASET="${DATASET:-$SCRATCH_ROOT/experiments/h16_merged_dataset}"
 RUN_ROOT="${RUN_ROOT:-$SCRATCH_ROOT/experiments/dovla_h16_policy_ckpt_runs}"
 OBJECTIVE="${OBJECTIVE:-base}"
+POLICY_TARGET_TYPES="${POLICY_TARGET_TYPES:-}"
 SEED=$SLURM_ARRAY_TASK_ID
 OUT_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
 
@@ -45,11 +46,22 @@ PYTHON_CMD=(
   "$SIF" "$CONTAINER_PYTHON"
 )
 
+TRAIN_EXTRA_ARGS=()
+if [[ -n "$POLICY_TARGET_TYPES" ]]; then
+  TRAIN_EXTRA_ARGS+=(--policy-target-types "$POLICY_TARGET_TYPES")
+fi
+if [[ -n "${EXTRA_TRAIN_ARGS:-}" ]]; then
+  # shellcheck disable=SC2206
+  EXTRA_SPLIT=($EXTRA_TRAIN_ARGS)
+  TRAIN_EXTRA_ARGS+=("${EXTRA_SPLIT[@]}")
+fi
+
 echo "=================================================="
 echo "Train DoVLAModel h=16 with best_policy checkpoint"
 echo "Seed: $SEED"
 echo "Dataset: $DATASET"
 echo "Output: $OUT_DIR"
+echo "Policy target types: ${POLICY_TARGET_TYPES:-<best-any>}"
 echo "=================================================="
 
 "${PYTHON_CMD[@]}" -c "
@@ -77,7 +89,8 @@ assert len(ds.group_ids) == 2873
   --lr 0.001 \
   --device cuda \
   --seed "$SEED" \
-  --objective lattice_field
+  --objective lattice_field \
+  "${TRAIN_EXTRA_ARGS[@]}"
 
 test -f "$OUT_DIR/best.pt"
 test -f "$OUT_DIR/best_policy.pt"

scripts/train_dovla.py

@@ -79,6 +79,12 @@ def main(argv: list[str] | None = None) -> int:
         default="same_state",
         help="Choose whether legacy ranking pairs share the exact simulator state.",
     )
+    parser.add_argument(
+        "--policy-target-types",
+        default="",
+        help="Comma-separated candidate_type filter for policy BC targets. "
+        "Empty means best action from every group.",
+    )
     parser.add_argument(
         "--loss-weight",
         action="append",
@@ -120,6 +126,9 @@ def main(argv: list[str] | None = None) -> int:
         objective=args.objective,
         lattice_neighbors=args.lattice_neighbors,
         pair_scope=args.pair_scope,
+        policy_target_types=tuple(
+            item.strip() for item in args.policy_target_types.split(",") if item.strip()
+        ),
         losses=loss_weights,
     )
     result = DoVLATrainer(config).train()

tests/test_trainer.py

@@ -9,6 +9,7 @@ from dovla_cil.tasks.library import built_in_toy_tasks
 from dovla_cil.training.trainer import (
     DoVLATrainer,
     TrainerConfig,
+    _best_records_by_group,
     _cross_state_pair_indices,
     _reward_utility_values,
 )
@@ -107,3 +108,36 @@ def test_cross_state_scope_rejects_lattice_field_objective(tmp_path: Path) -> No
         raise AssertionError(
             "cross-state pairs cannot silently leave same-state field edges active"
         )
+
+
+def test_policy_target_type_filter_selects_best_allowed_candidate() -> None:
+    records = [
+        SimpleNamespace(
+            group_id="g0",
+            candidate_type="expert",
+            reward=SimpleNamespace(score=2.0),
+            rank_within_group=0,
+            record_id="expert",
+        ),
+        SimpleNamespace(
+            group_id="g0",
+            candidate_type="near_miss",
+            reward=SimpleNamespace(score=1.5),
+            rank_within_group=1,
+            record_id="near",
+        ),
+        SimpleNamespace(
+            group_id="g1",
+            candidate_type="expert",
+            reward=SimpleNamespace(score=1.0),
+            rank_within_group=0,
+            record_id="fallback",
+        ),
+    ]
+
+    selected = _best_records_by_group(records, candidate_types=("near_miss",))
+
+    assert {record.group_id: record.record_id for record in selected} == {
+        "g0": "near",
+        "g1": "fallback",
+    }