auto-sync 2026-07-02T15:44:03Z workspace (part 3)

workspace/scripts/build_oracle_selector_calibration.py

@@ -93,6 +93,36 @@ def _type_bonuses(
     return bonuses, means, counts, baseline
 
 
+def _scale_bonuses(
+    values_by_scale: dict[str, list[float]],
+    *,
+    scale: float,
+    min_count: int,
+    max_abs_bonus: float | None,
+) -> tuple[dict[str, float], dict[str, float | None], dict[str, int], float | None]:
+    pooled = [value for values in values_by_scale.values() for value in values]
+    baseline = _mean(pooled)
+    means = {
+        residual_scale: _mean(values)
+        for residual_scale, values in sorted(values_by_scale.items(), key=lambda item: float(item[0]))
+    }
+    counts = {
+        residual_scale: len(values)
+        for residual_scale, values in sorted(values_by_scale.items(), key=lambda item: float(item[0]))
+    }
+    if baseline is None or len(pooled) < min_count:
+        return {}, means, counts, baseline
+    bonuses: dict[str, float] = {}
+    for residual_scale, mean in means.items():
+        if mean is None or counts[residual_scale] < min_count:
+            continue
+        bonuses[residual_scale] = _clip(
+            float(scale) * (float(mean) - float(baseline)),
+            max_abs_bonus,
+        )
+    return bonuses, means, counts, baseline
+
+
 def _iter_rollout_paths(patterns: list[str]) -> list[Path]:
     paths: list[Path] = []
     for pattern in patterns:
@@ -112,9 +142,11 @@ def build_oracle_selector_calibration(
     max_rank: int,
     rank_scale: float,
     type_scale: float,
+    scale_scale: float,
     min_count: int,
     max_abs_rank_bias: float | None,
     max_abs_type_bonus: float | None,
+    max_abs_scale_bonus: float | None,
 ) -> dict[str, Any]:
     if max_rank <= 0:
         raise ValueError("max_rank must be positive")
@@ -127,8 +159,12 @@ def build_oracle_selector_calibration(
     type_by_task: dict[str, dict[str, list[float]]] = defaultdict(
         lambda: defaultdict(list)
     )
+    scale_by_task: dict[str, dict[str, list[float]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
     global_ranks: list[list[float]] = [[] for _ in range(max_rank)]
     global_types: dict[str, list[float]] = defaultdict(list)
+    global_scales: dict[str, list[float]] = defaultdict(list)
     rows_seen = 0
     rows_used = 0
     skipped_branches = 0
@@ -142,12 +178,24 @@ def build_oracle_selector_calibration(
             task_id = str(row.get("task_id") or "")
             values = _branch_values(row, objective)
             candidate_types = row.get("candidate_oracle_types")
+            residual_scales = row.get("candidate_oracle_residual_scales")
             valid_mask = row.get("candidate_oracle_valid_mask")
-            if not task_id or not values or not isinstance(candidate_types, list):
+            if (
+                not task_id
+                or not values
+                or not isinstance(candidate_types, list)
+                or not isinstance(residual_scales, list)
+            ):
                 continue
             if not isinstance(valid_mask, list):
                 valid_mask = [True] * len(values)
-            branch_count = min(max_rank, len(values), len(candidate_types), len(valid_mask))
+            branch_count = min(
+                max_rank,
+                len(values),
+                len(candidate_types),
+                len(residual_scales),
+                len(valid_mask),
+            )
             if branch_count <= 0:
                 continue
             rows_used += 1
@@ -157,10 +205,13 @@ def build_oracle_selector_calibration(
                     continue
                 value = float(values[rank])
                 candidate_type = str(candidate_types[rank])
+                residual_scale = f"{float(residual_scales[rank]):g}"
                 rank_by_task[task_id][rank].append(value)
                 global_ranks[rank].append(value)
                 type_by_task[task_id][candidate_type].append(value)
                 global_types[candidate_type].append(value)
+                scale_by_task[task_id][residual_scale].append(value)
+                global_scales[residual_scale].append(value)
 
     field_rank_biases_by_task: dict[str, list[float]] = {}
     rank_utility_means_by_task: dict[str, list[float | None]] = {}
@@ -213,6 +264,34 @@ def build_oracle_selector_calibration(
     type_counts_by_task["*"] = global_type_counts
     type_baselines_by_task["*"] = global_type_baseline
 
+    residual_scale_bonuses_by_task: dict[str, dict[str, float]] = {}
+    scale_utility_means_by_task: dict[str, dict[str, float | None]] = {}
+    scale_counts_by_task: dict[str, dict[str, int]] = {}
+    scale_baselines_by_task: dict[str, float | None] = {}
+    for task_id in sorted(scale_by_task):
+        bonuses, means, counts, baseline = _scale_bonuses(
+            scale_by_task[task_id],
+            scale=scale_scale,
+            min_count=min_count,
+            max_abs_bonus=max_abs_scale_bonus,
+        )
+        residual_scale_bonuses_by_task[task_id] = bonuses
+        scale_utility_means_by_task[task_id] = means
+        scale_counts_by_task[task_id] = counts
+        scale_baselines_by_task[task_id] = baseline
+    global_scale_bonuses, global_scale_means, global_scale_counts, global_scale_baseline = (
+        _scale_bonuses(
+            global_scales,
+            scale=scale_scale,
+            min_count=min_count,
+            max_abs_bonus=max_abs_scale_bonus,
+        )
+    )
+    residual_scale_bonuses_by_task["*"] = global_scale_bonuses
+    scale_utility_means_by_task["*"] = global_scale_means
+    scale_counts_by_task["*"] = global_scale_counts
+    scale_baselines_by_task["*"] = global_scale_baseline
+
     return {
         "source_rollouts": [str(path) for path in rollout_paths],
         "calibration_source": "candidate_oracle_rollout",
@@ -220,9 +299,11 @@ def build_oracle_selector_calibration(
         "max_rank": int(max_rank),
         "rank_scale": float(rank_scale),
         "type_scale": float(type_scale),
+        "scale_scale": float(scale_scale),
         "min_count": int(min_count),
         "max_abs_rank_bias": max_abs_rank_bias,
         "max_abs_type_bonus": max_abs_type_bonus,
+        "max_abs_scale_bonus": max_abs_scale_bonus,
         "rows_seen": rows_seen,
         "rows_used": rows_used,
         "skipped_branches": skipped_branches,
@@ -233,6 +314,10 @@ def build_oracle_selector_calibration(
         "type_utility_means_by_task": type_utility_means_by_task,
         "type_counts_by_task": type_counts_by_task,
         "type_baselines_by_task": type_baselines_by_task,
+        "residual_scale_bonuses_by_task": residual_scale_bonuses_by_task,
+        "scale_utility_means_by_task": scale_utility_means_by_task,
+        "scale_counts_by_task": scale_counts_by_task,
+        "scale_baselines_by_task": scale_baselines_by_task,
     }
 
 
@@ -248,13 +333,16 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--max-rank", type=int, default=4)
     parser.add_argument("--rank-scale", type=float, default=0.05)
     parser.add_argument("--type-scale", type=float, default=0.05)
+    parser.add_argument("--scale-scale", type=float, default=0.0)
     parser.add_argument("--min-count", type=int, default=20)
     parser.add_argument("--max-abs-rank-bias", type=float, default=0.02)
     parser.add_argument("--max-abs-type-bonus", type=float, default=0.02)
+    parser.add_argument("--max-abs-scale-bonus", type=float, default=0.02)
     args = parser.parse_args(argv)
 
     max_abs_rank_bias = args.max_abs_rank_bias if args.max_abs_rank_bias > 0 else None
     max_abs_type_bonus = args.max_abs_type_bonus if args.max_abs_type_bonus > 0 else None
+    max_abs_scale_bonus = args.max_abs_scale_bonus if args.max_abs_scale_bonus > 0 else None
     rollout_paths = _iter_rollout_paths(args.rollout)
     result = build_oracle_selector_calibration(
         rollout_paths,
@@ -262,9 +350,11 @@ def main(argv: list[str] | None = None) -> int:
         max_rank=args.max_rank,
         rank_scale=args.rank_scale,
         type_scale=args.type_scale,
+        scale_scale=args.scale_scale,
         min_count=args.min_count,
         max_abs_rank_bias=max_abs_rank_bias,
         max_abs_type_bonus=max_abs_type_bonus,
+        max_abs_scale_bonus=max_abs_scale_bonus,
     )
     args.out.parent.mkdir(parents=True, exist_ok=True)
     args.out.write_text(json.dumps(result, indent=2) + "\n")
@@ -282,6 +372,10 @@ def main(argv: list[str] | None = None) -> int:
                     "type_utility_means_by_task",
                     "type_counts_by_task",
                     "type_baselines_by_task",
+                    "residual_scale_bonuses_by_task",
+                    "scale_utility_means_by_task",
+                    "scale_counts_by_task",
+                    "scale_baselines_by_task",
                 }
             },
             indent=2,

workspace/scripts/eval_maniskill_policy_rollout.py

@@ -90,6 +90,33 @@ def _load_source_score_bonus_map(path: Path | None) -> dict[str, float]:
         ) from exc
 
 
+def _load_residual_scale_bonus_map(path: Path | None) -> dict[str, dict[str, float]]:
+    if path is None:
+        return {}
+    payload = json.loads(path.read_text())
+    if not isinstance(payload, dict):
+        raise SystemExit("--residual-scale-bonus-map must contain a JSON object")
+    raw_bonuses = payload.get("residual_scale_bonuses_by_task", payload)
+    if not isinstance(raw_bonuses, dict):
+        raise SystemExit(
+            "--residual-scale-bonus-map must contain a task->scale->bonus object "
+            "or residual_scale_bonuses_by_task"
+        )
+    bonuses: dict[str, dict[str, float]] = {}
+    try:
+        for task_id, values in raw_bonuses.items():
+            if not isinstance(values, dict):
+                raise ValueError
+            bonuses[str(task_id)] = {
+                str(scale): float(bonus) for scale, bonus in values.items()
+            }
+    except (TypeError, ValueError) as exc:
+        raise SystemExit(
+            "--residual-scale-bonus-map values must be numeric scale->bonus objects"
+        ) from exc
+    return bonuses
+
+
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(
         description="Execute a DoVLA policy checkpoint from restored ManiSkill CIL states."
@@ -408,6 +435,15 @@ def main(argv: list[str] | None = None) -> int:
         default=None,
         help="Optional JSON task->rank-bias map added by field rank before lattice selection.",
     )
+    parser.add_argument(
+        "--residual-scale-bonus-map",
+        type=Path,
+        default=None,
+        help=(
+            "Optional JSON task->residual_scale->bonus map added after residual "
+            "scale-grid expansion in retrieval_residual mode."
+        ),
+    )
     parser.add_argument(
         "--candidate-oracle-rollouts",
         type=int,
@@ -525,6 +561,9 @@ def main(argv: list[str] | None = None) -> int:
     source_score_bonus_by_task = _load_source_score_bonus_map(
         args.retrieval_residual_source_score_bonus_map
     )
+    residual_scale_bonuses_by_task = _load_residual_scale_bonus_map(
+        args.residual_scale_bonus_map
+    )
     result = evaluate_maniskill_policy_rollout(
         args.checkpoint,
         args.dataset,
@@ -591,6 +630,7 @@ def main(argv: list[str] | None = None) -> int:
         candidate_type_bonuses_by_task=candidate_type_bonuses_by_task,
         candidate_type_bonus_components=args.candidate_type_bonus_components,
         field_rank_biases_by_task=field_rank_biases_by_task,
+        residual_scale_bonuses_by_task=residual_scale_bonuses_by_task,
         candidate_oracle_rollouts=args.candidate_oracle_rollouts,
         candidate_oracle_unique_tolerance=args.candidate_oracle_unique_tolerance,
     )

workspace/scripts/slurm/build_oracle_selector_calibration.sbatch

@@ -20,9 +20,11 @@ CAL_OBJECTIVE="${CAL_OBJECTIVE:-score}"
 MAX_RANK="${MAX_RANK:-4}"
 RANK_SCALE="${RANK_SCALE:-0.05}"
 TYPE_SCALE="${TYPE_SCALE:-0.05}"
+SCALE_SCALE="${SCALE_SCALE:-0.0}"
 MIN_COUNT="${MIN_COUNT:-20}"
 MAX_ABS_RANK_BIAS="${MAX_ABS_RANK_BIAS:-0.02}"
 MAX_ABS_TYPE_BONUS="${MAX_ABS_TYPE_BONUS:-0.02}"
+MAX_ABS_SCALE_BONUS="${MAX_ABS_SCALE_BONUS:-0.02}"
 PYTHON="${PYTHON:-python3}"
 
 ROLLOUT="$RUN_ROOT/$OBJECTIVE/seed_$SEED/$ROLLOUT_NAME"
@@ -38,6 +40,8 @@ mkdir -p outputs/hpc/logs "$(dirname "$OUT")"
   --max-rank "$MAX_RANK" \
   --rank-scale "$RANK_SCALE" \
   --type-scale "$TYPE_SCALE" \
+  --scale-scale "$SCALE_SCALE" \
   --min-count "$MIN_COUNT" \
   --max-abs-rank-bias "$MAX_ABS_RANK_BIAS" \
-  --max-abs-type-bonus "$MAX_ABS_TYPE_BONUS"
+  --max-abs-type-bonus "$MAX_ABS_TYPE_BONUS" \
+  --max-abs-scale-bonus "$MAX_ABS_SCALE_BONUS"

workspace/scripts/slurm/eval_maniskill_policy_rollout.sbatch

@@ -121,6 +121,16 @@ if [[ -z "$FIELD_RANK_BIAS_MAP" && -n "$FIELD_RANK_BIAS_OBJECTIVE" ]]; then
   fi
   FIELD_RANK_BIAS_MAP="$RUN_ROOT/$FIELD_RANK_BIAS_OBJECTIVE/seed_$SEED/$FIELD_RANK_BIAS_NAME"
 fi
+RESIDUAL_SCALE_BONUS_MAP="${RESIDUAL_SCALE_BONUS_MAP:-}"
+RESIDUAL_SCALE_BONUS_OBJECTIVE="${RESIDUAL_SCALE_BONUS_OBJECTIVE:-}"
+RESIDUAL_SCALE_BONUS_NAME="${RESIDUAL_SCALE_BONUS_NAME:-residual_scale_bonuses.json}"
+if [[ -z "$RESIDUAL_SCALE_BONUS_MAP" && -n "$RESIDUAL_SCALE_BONUS_OBJECTIVE" ]]; then
+  if [[ -z "$RUN_ROOT" ]]; then
+    echo "RESIDUAL_SCALE_BONUS_OBJECTIVE requires RUN_ROOT" >&2
+    exit 1
+  fi
+  RESIDUAL_SCALE_BONUS_MAP="$RUN_ROOT/$RESIDUAL_SCALE_BONUS_OBJECTIVE/seed_$SEED/$RESIDUAL_SCALE_BONUS_NAME"
+fi
 CANDIDATE_ORACLE_ROLLOUTS="${CANDIDATE_ORACLE_ROLLOUTS:-0}"
 CANDIDATE_ORACLE_UNIQUE_TOLERANCE="${CANDIDATE_ORACLE_UNIQUE_TOLERANCE:-1e-6}"
 
@@ -157,6 +167,9 @@ fi
 if [[ -n "$FIELD_RANK_BIAS_MAP" ]]; then
   EXTRA_ARGS+=(--field-rank-bias-map "$FIELD_RANK_BIAS_MAP")
 fi
+if [[ -n "$RESIDUAL_SCALE_BONUS_MAP" ]]; then
+  EXTRA_ARGS+=(--residual-scale-bonus-map "$RESIDUAL_SCALE_BONUS_MAP")
+fi
 if [[ -n "$RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_MAP" ]]; then
   EXTRA_ARGS+=(--retrieval-residual-source-score-bonus-map "$RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_MAP")
 fi