| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import statistics |
| import subprocess |
| from collections import Counter |
| from dataclasses import dataclass |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| DEFAULT_OBJECTIVE = "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1" |
| BEST_CLEAN_SUMMARY = ( |
| "h16_transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_summary.json" |
| ) |
| BEST_CLEAN_SUCCESS = 0.3889855072463768 |
| SOURCE_SCORE_MAP = "manifests/source_score_bonus_pick001_stack005.json" |
| DATASET = "/scratch/{user}/dovla/experiments/six_task_h16_collection" |
| RUN_ROOT = "/scratch/{user}/dovla/experiments/dovla_h16_policy_ckpt_runs" |
| EXCLUDE_TYPES = ( |
| "residual_random_negative:" |
| "residual_wrong_direction:" |
| "residual_near_miss+residual_no_op:" |
| "residual_no_op+residual_wrong_gripper" |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class V1Config: |
| key: str |
| label: str |
| out_name: str |
| summary_tag: str |
| anchor: str |
| advantage_weight: float |
| min_source_advantage: float = -1.0e9 |
| candidate_oracle_rollouts: int = 0 |
|
|
|
|
| INITIAL_CONFIGS = ( |
| V1Config( |
| key="advw1p0", |
| label="Generator V1 utility-weighted tangents, expert anchor, advantage weight 1.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw1p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw1p0" |
| ), |
| anchor="expert", |
| advantage_weight=1.0, |
| ), |
| V1Config( |
| key="advw2p0", |
| label="Generator V1 utility-weighted tangents, expert anchor, advantage weight 2.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw2p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw2p0" |
| ), |
| anchor="expert", |
| advantage_weight=2.0, |
| ), |
| V1Config( |
| key="policyanchor_advw2p0", |
| label="Generator V1 positive tangents, policy anchor, advantage weight 2.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw2p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw2p0" |
| ), |
| anchor="policy", |
| advantage_weight=2.0, |
| ), |
| ) |
|
|
| NEXT_SWEEP_CONFIGS = ( |
| V1Config( |
| key="advw0p5", |
| label="Generator V1 utility-weighted tangents, expert anchor, advantage weight 0.5", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw0p5.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw0p5" |
| ), |
| anchor="expert", |
| advantage_weight=0.5, |
| ), |
| V1Config( |
| key="advw4p0", |
| label="Generator V1 utility-weighted tangents, expert anchor, advantage weight 4.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw4p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw4p0" |
| ), |
| anchor="expert", |
| advantage_weight=4.0, |
| ), |
| V1Config( |
| key="policyanchor_advw1p0", |
| label="Generator V1 positive tangents, policy anchor, advantage weight 1.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw1p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw1p0" |
| ), |
| anchor="policy", |
| advantage_weight=1.0, |
| ), |
| V1Config( |
| key="policyanchor_advw4p0", |
| label="Generator V1 positive tangents, policy anchor, advantage weight 4.0", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw4p0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw4p0" |
| ), |
| anchor="policy", |
| advantage_weight=4.0, |
| ), |
| V1Config( |
| key="advw2p0_gate0", |
| label="Generator V1 utility-weighted tangents, expert anchor, positive-source gate", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw2p0_gate0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_advw2p0_gate0" |
| ), |
| anchor="expert", |
| advantage_weight=2.0, |
| min_source_advantage=0.0, |
| ), |
| V1Config( |
| key="policyanchor_advw2p0_gate0", |
| label="Generator V1 positive tangents, policy anchor, positive-source gate", |
| out_name=( |
| "policy_rollout_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw2p0_gate0.json" |
| ), |
| summary_tag=( |
| "transport_field_reground_fieldonly_k6clean_dropnoopwg_b12_v1_" |
| "besttransport_margin0p00_k6_srcscore_task_pick001_stack005_policyanchor_advw2p0_gate0" |
| ), |
| anchor="policy", |
| advantage_weight=2.0, |
| min_source_advantage=0.0, |
| ), |
| ) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--project-dir", type=Path, default=Path.cwd()) |
| parser.add_argument("--run-root", type=Path, default=None) |
| parser.add_argument("--dataset", type=Path, default=None) |
| parser.add_argument("--objective", default=DEFAULT_OBJECTIVE) |
| parser.add_argument("--results-dir", type=Path, default=Path("results")) |
| parser.add_argument("--submit-next", action="store_true") |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--round", type=int, default=int(os.environ.get("ADVANCE_ROUND", "0"))) |
| args = parser.parse_args() |
|
|
| user = os.environ.get("USER", "knguy52") |
| run_root = args.run_root or Path(RUN_ROOT.format(user=user)) |
| dataset = args.dataset or Path(DATASET.format(user=user)) |
| project_dir = args.project_dir.resolve() |
| results_dir = args.results_dir |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| summaries = [ |
| _summarize_config(config, run_root=run_root, objective=args.objective, results_dir=results_dir) |
| for config in INITIAL_CONFIGS |
| ] |
| baseline_success = _baseline_success(results_dir) |
| complete = [item for item in summaries if item["num_completed"] == 3] |
| best = max(complete, key=lambda item: float(item["mean_success"]), default=None) |
| decision = _decision_payload( |
| summaries, |
| best=best, |
| baseline_success=baseline_success, |
| round_index=args.round, |
| ) |
| decision_path = results_dir / "v1_generator_decision.json" |
| decision_md_path = results_dir / "v1_generator_decision.md" |
| decision_path.write_text(json.dumps(decision, indent=2)) |
| decision_md_path.write_text(_render_decision_markdown(decision)) |
|
|
| if complete: |
| _run_build_paper_analysis(project_dir, dry_run=args.dry_run) |
|
|
| if args.submit_next and not decision["complete"]: |
| submitted_missing = _maybe_resubmit_missing( |
| summaries, |
| INITIAL_CONFIGS, |
| project_dir=project_dir, |
| run_root=run_root, |
| dataset=dataset, |
| objective=args.objective, |
| dry_run=args.dry_run, |
| ) |
| decision["submitted_missing"] = submitted_missing |
| decision_path.write_text(json.dumps(decision, indent=2)) |
| decision_md_path.write_text(_render_decision_markdown(decision)) |
| elif args.submit_next and complete: |
| submitted = _maybe_submit_next( |
| decision, |
| project_dir=project_dir, |
| run_root=run_root, |
| dataset=dataset, |
| objective=args.objective, |
| dry_run=args.dry_run, |
| ) |
| decision["submitted"] = submitted |
| decision_path.write_text(json.dumps(decision, indent=2)) |
| decision_md_path.write_text(_render_decision_markdown(decision)) |
|
|
| print(json.dumps(decision, indent=2)) |
| return 0 |
|
|
|
|
| def _load_json(path: Path) -> dict[str, Any]: |
| return json.loads(path.read_text()) |
|
|
|
|
| def _mean(values: list[float]) -> float: |
| return statistics.mean(values) if values else 0.0 |
|
|
|
|
| def _std(values: list[float]) -> float: |
| return statistics.stdev(values) if len(values) > 1 else 0.0 |
|
|
|
|
| def _baseline_success(results_dir: Path) -> float: |
| path = results_dir / BEST_CLEAN_SUMMARY |
| if path.exists(): |
| value = _load_json(path).get("mean_success") |
| if isinstance(value, (int, float)) and math.isfinite(float(value)): |
| return float(value) |
| return BEST_CLEAN_SUCCESS |
|
|
|
|
| def _summarize_config( |
| config: V1Config, |
| *, |
| run_root: Path, |
| objective: str, |
| results_dir: Path, |
| ) -> dict[str, Any]: |
| rows = [] |
| base_dir = run_root / objective |
| for result_path in sorted(base_dir.glob(f"seed_*/{config.out_name}")): |
| raw = _load_json(result_path) |
| seed = int(result_path.parent.name.split("_")[-1]) |
| selected_scale_counts = Counter( |
| str(row.get("selected_residual_scale")) |
| for row in raw.get("rows", []) |
| if row.get("selected_residual_scale") is not None |
| ) |
| rows.append( |
| { |
| "seed": seed, |
| "path": str(result_path), |
| "num_groups": raw.get("num_groups", 0), |
| "selection_mode": raw.get("selection_mode"), |
| "num_candidates": raw.get("num_candidates"), |
| "candidate_sigma": raw.get("candidate_sigma", 0.0), |
| "selection_margin": raw.get("selection_margin", 0.0), |
| "retrieval_neighbors": raw.get("retrieval_neighbors", 0), |
| "retrieval_metric": raw.get("retrieval_metric", "none"), |
| "retrieval_residual_anchor": raw.get("retrieval_residual_anchor", "none"), |
| "retrieval_residual_direction": raw.get("retrieval_residual_direction", "none"), |
| "retrieval_residual_reduce": raw.get("retrieval_residual_reduce", "none"), |
| "retrieval_residual_scales": raw.get("retrieval_residual_scales", []), |
| "retrieval_residual_source_score_bonus_by_task": raw.get( |
| "retrieval_residual_source_score_bonus_by_task", {} |
| ), |
| "retrieval_residual_source_advantage_weight_scale": raw.get( |
| "retrieval_residual_source_advantage_weight_scale", 0.0 |
| ), |
| "retrieval_residual_min_source_advantage": raw.get( |
| "retrieval_residual_min_source_advantage", -1.0e9 |
| ), |
| "lattice_exclude_types": raw.get("lattice_exclude_types", []), |
| "selected_residual_scale_counts": dict(selected_scale_counts), |
| "policy_rollout_success_rate": raw.get("policy_rollout_success_rate", 0.0), |
| "policy_rollout_progress": raw.get("policy_rollout_progress", 0.0), |
| "oracle_success_rate": raw.get("oracle_success_rate", 0.0), |
| "action_mse_to_best": raw.get("action_mse_to_best", 0.0), |
| "per_task": raw.get("per_task", {}), |
| } |
| ) |
| successes = [float(row["policy_rollout_success_rate"]) for row in rows] |
| progresses = [float(row["policy_rollout_progress"]) for row in rows] |
| mses = [float(row["action_mse_to_best"]) for row in rows] |
| summary = { |
| "run_root": str(run_root), |
| "objective": objective, |
| "out_name": config.out_name, |
| "key": config.key, |
| "label": config.label, |
| "num_completed": len(rows), |
| "mean_success": _mean(successes), |
| "std_success": _std(successes), |
| "mean_progress": _mean(progresses), |
| "mean_action_mse_to_best": _mean(mses), |
| "rows": rows, |
| } |
| json_path = results_dir / f"h16_{config.summary_tag}_summary.json" |
| md_path = results_dir / f"h16_{config.summary_tag}_summary.md" |
| if len(rows) == 3: |
| json_path.write_text(json.dumps(summary, indent=2)) |
| md_path.write_text(_render_summary_markdown(summary)) |
| return summary | {"summary_path": str(json_path)} |
|
|
|
|
| def _render_summary_markdown(summary: dict[str, Any]) -> str: |
| lines = [ |
| "# Generator V1 Rollout Summary", |
| "", |
| f"Objective: `{summary['objective']}`", |
| f"Result file: `{summary['out_name']}`", |
| f"Completed seeds: {summary['num_completed']}", |
| f"Mean success: {summary['mean_success']:.2%} +/- {summary['std_success']:.2%}", |
| f"Mean progress: {summary['mean_progress']:.2%}", |
| "", |
| "| seed | success | progress | mse | anchor | adv weight | min src adv | scales |", |
| "|---:|---:|---:|---:|---|---:|---|---|", |
| ] |
| for row in summary["rows"]: |
| scales = ",".join(str(item) for item in row.get("retrieval_residual_scales", [])) |
| lines.append( |
| "| {seed} | {success:.2%} | {progress:.2%} | {mse:.3f} | {anchor} | {weight:.2f} | {min_adv} | {scales} |".format( |
| seed=row["seed"], |
| success=row["policy_rollout_success_rate"], |
| progress=row["policy_rollout_progress"], |
| mse=row["action_mse_to_best"], |
| anchor=row.get("retrieval_residual_anchor", "none"), |
| weight=row.get("retrieval_residual_source_advantage_weight_scale", 0.0), |
| min_adv=_format_min_source_advantage( |
| row.get("retrieval_residual_min_source_advantage", -1.0e9) |
| ), |
| scales=scales or "none", |
| ) |
| ) |
| return "\n".join(lines) + "\n" |
|
|
|
|
| def _format_min_source_advantage(value: Any) -> str: |
| try: |
| numeric = float(value) |
| except (TypeError, ValueError): |
| return "none" |
| if numeric <= -1.0e8: |
| return "none" |
| return f"{numeric:.2f}" |
|
|
|
|
| def _decision_payload( |
| summaries: list[dict[str, Any]], |
| *, |
| best: dict[str, Any] | None, |
| baseline_success: float, |
| round_index: int, |
| ) -> dict[str, Any]: |
| return { |
| "generated_utc": datetime.now(timezone.utc).isoformat(), |
| "round": round_index, |
| "baseline_success": baseline_success, |
| "complete": all(item["num_completed"] == 3 for item in summaries), |
| "summaries": [ |
| { |
| "key": item["key"], |
| "label": item["label"], |
| "num_completed": item["num_completed"], |
| "mean_success": item["mean_success"], |
| "delta_vs_baseline": item["mean_success"] - baseline_success, |
| "summary_path": item["summary_path"], |
| } |
| for item in summaries |
| ], |
| "best_key": best["key"] if best else None, |
| "best_success": best["mean_success"] if best else None, |
| "best_delta_vs_baseline": ( |
| best["mean_success"] - baseline_success if best else None |
| ), |
| "recommendation": _recommendation(best, baseline_success), |
| } |
|
|
|
|
| def _recommendation(best: dict[str, Any] | None, baseline_success: float) -> str: |
| if best is None: |
| return "wait_for_v1_rollouts_or_debug_failures" |
| if float(best["mean_success"]) > baseline_success + 0.002: |
| return "submit_candidate_oracle_for_best_v1" |
| return "submit_wider_advantage_weight_support_sweep" |
|
|
|
|
| def _render_decision_markdown(decision: dict[str, Any]) -> str: |
| lines = [ |
| "# Generator V1 Decision", |
| "", |
| f"Generated: `{decision['generated_utc']}`", |
| f"Baseline clean success: {decision['baseline_success']:.2%}", |
| f"Recommendation: `{decision['recommendation']}`", |
| "", |
| "| key | completed | success | delta |", |
| "|---|---:|---:|---:|", |
| ] |
| for item in decision["summaries"]: |
| lines.append( |
| f"| {item['key']} | {item['num_completed']} | {item['mean_success']:.2%} | {item['delta_vs_baseline']:+.2%} |" |
| ) |
| if decision.get("submitted"): |
| lines.extend(["", "Submitted follow-up jobs:", ""]) |
| for item in decision["submitted"]: |
| lines.append(f"- `{item['key']}`: eval `{item.get('eval_job')}`, summary `{item.get('summary_job')}`") |
| if decision.get("submitted_missing"): |
| lines.extend(["", "Resubmitted missing seeds:", ""]) |
| for item in decision["submitted_missing"]: |
| lines.append( |
| f"- `{item['key']}` seeds `{item.get('missing_seeds')}`: " |
| f"eval `{item.get('eval_job')}`, summary `{item.get('summary_job')}`" |
| ) |
| return "\n".join(lines) + "\n" |
|
|
|
|
| def _run_build_paper_analysis(project_dir: Path, *, dry_run: bool) -> None: |
| if dry_run: |
| return |
| subprocess.run( |
| ["python3", "scripts/build_paper_analysis.py"], |
| cwd=project_dir, |
| check=True, |
| ) |
|
|
|
|
| def _maybe_submit_next( |
| decision: dict[str, Any], |
| *, |
| project_dir: Path, |
| run_root: Path, |
| dataset: Path, |
| objective: str, |
| dry_run: bool, |
| ) -> list[dict[str, str]]: |
| marker = project_dir / "results" / "v1_generator_next_submitted.json" |
| if marker.exists(): |
| return _load_json(marker).get("submitted", []) |
| best_key = decision.get("best_key") |
| recommendation = str(decision.get("recommendation")) |
| if recommendation == "submit_candidate_oracle_for_best_v1": |
| configs = [ |
| _oracle_config(config) |
| for config in INITIAL_CONFIGS |
| if config.key == best_key |
| ] |
| elif recommendation == "submit_wider_advantage_weight_support_sweep": |
| configs = list(NEXT_SWEEP_CONFIGS) |
| else: |
| configs = [] |
| submitted = [ |
| _submit_config( |
| config, |
| project_dir=project_dir, |
| run_root=run_root, |
| dataset=dataset, |
| objective=objective, |
| dry_run=dry_run, |
| ) |
| for config in configs |
| ] |
| marker.write_text( |
| json.dumps( |
| { |
| "generated_utc": datetime.now(timezone.utc).isoformat(), |
| "recommendation": recommendation, |
| "submitted": submitted, |
| }, |
| indent=2, |
| ) |
| ) |
| return submitted |
|
|
|
|
| def _maybe_resubmit_missing( |
| summaries: list[dict[str, Any]], |
| configs: tuple[V1Config, ...], |
| *, |
| project_dir: Path, |
| run_root: Path, |
| dataset: Path, |
| objective: str, |
| dry_run: bool, |
| ) -> list[dict[str, str]]: |
| marker = project_dir / "results" / "v1_generator_missing_resubmitted.json" |
| if marker.exists(): |
| return _load_json(marker).get("submitted", []) |
| by_key = {item.key: item for item in configs} |
| submitted: list[dict[str, str]] = [] |
| for summary in summaries: |
| config = by_key.get(str(summary.get("key"))) |
| if config is None: |
| continue |
| present = {int(row["seed"]) for row in summary.get("rows", [])} |
| missing = [str(seed) for seed in (0, 1, 2) if seed not in present] |
| if not missing: |
| continue |
| submitted.append( |
| _submit_config( |
| config, |
| project_dir=project_dir, |
| run_root=run_root, |
| dataset=dataset, |
| objective=objective, |
| array_spec=",".join(missing), |
| dry_run=dry_run, |
| ) |
| ) |
| submitted[-1]["missing_seeds"] = ",".join(missing) |
| marker.write_text( |
| json.dumps( |
| { |
| "generated_utc": datetime.now(timezone.utc).isoformat(), |
| "submitted": submitted, |
| }, |
| indent=2, |
| ) |
| ) |
| return submitted |
|
|
|
|
| def _oracle_config(config: V1Config) -> V1Config: |
| return V1Config( |
| key=f"{config.key}_oraclek8", |
| label=f"{config.label}, candidate oracle K8", |
| out_name=config.out_name.replace(".json", "_oraclek8.json"), |
| summary_tag=f"{config.summary_tag}_oraclek8", |
| anchor=config.anchor, |
| advantage_weight=config.advantage_weight, |
| min_source_advantage=config.min_source_advantage, |
| candidate_oracle_rollouts=8, |
| ) |
|
|
|
|
| def _submit_config( |
| config: V1Config, |
| *, |
| project_dir: Path, |
| run_root: Path, |
| dataset: Path, |
| objective: str, |
| array_spec: str = "0-2", |
| dry_run: bool, |
| ) -> dict[str, str]: |
| export = { |
| "PROJECT_DIR": str(project_dir), |
| "RUN_ROOT": str(run_root), |
| "DATASET": str(dataset), |
| "OBJECTIVE": objective, |
| "CHECKPOINT_NAME": "best_transport.pt", |
| "MAX_GROUPS": "all", |
| "GROUP_BATCH_SIZE": "8", |
| "EVAL_SPLIT": "validation", |
| "SELECTION_MODE": "retrieval_residual", |
| "NUM_CANDIDATES": "1", |
| "CANDIDATE_SIGMA": "0.2", |
| "SELECTION_MARGIN": "0.0", |
| "RETRIEVAL_NEIGHBORS": "6", |
| "RETRIEVAL_METRIC": "raw", |
| "RETRIEVAL_RESIDUAL_REDUCE": "compose_mean_by_type", |
| "RETRIEVAL_RESIDUAL_DIRECTION": "candidate_minus_anchor", |
| "RETRIEVAL_RESIDUAL_SCALES_COLON": "0.35:0.4:0.45", |
| "LATTICE_EXCLUDE_TYPES_COLON": EXCLUDE_TYPES, |
| "RETRIEVAL_RESIDUAL_SOURCE_SCORE_BONUS_MAP": str(project_dir / SOURCE_SCORE_MAP), |
| "OUT_NAME": config.out_name, |
| "RETRIEVAL_RESIDUAL_ANCHOR": config.anchor, |
| "RETRIEVAL_RESIDUAL_SOURCE_ADVANTAGE_WEIGHT_SCALE": str(config.advantage_weight), |
| "RETRIEVAL_RESIDUAL_MIN_SOURCE_ADVANTAGE": str(config.min_source_advantage), |
| "CANDIDATE_ORACLE_ROLLOUTS": str(config.candidate_oracle_rollouts), |
| } |
| eval_cmd = [ |
| "sbatch", |
| "--array=0-2", |
| "--job-name", |
| _job_name(config.key), |
| "--export", |
| "ALL," + ",".join(f"{key}={value}" for key, value in export.items()), |
| "scripts/slurm/eval_maniskill_policy_rollout.sbatch", |
| ] |
| eval_cmd[1] = f"--array={array_spec}" |
| eval_job = _submit(eval_cmd, project_dir=project_dir, dry_run=dry_run) |
| summary_export = { |
| "PROJECT_DIR": str(project_dir), |
| "RUN_ROOT": str(run_root), |
| "OBJECTIVE": objective, |
| "OUT_NAME": config.out_name, |
| "SUMMARY_TAG": config.summary_tag, |
| } |
| summary_cmd = [ |
| "sbatch", |
| "--dependency", |
| f"afterok:{eval_job}" if eval_job != "dry-run" else "afterok:0", |
| "--job-name", |
| _job_name(f"sum_{config.key}"), |
| "--export", |
| "ALL," + ",".join(f"{key}={value}" for key, value in summary_export.items()), |
| "scripts/slurm/summarize_h16_policy_ckpt.sbatch", |
| ] |
| summary_job = _submit(summary_cmd, project_dir=project_dir, dry_run=dry_run) |
| return {"key": config.key, "eval_job": eval_job, "summary_job": summary_job} |
|
|
|
|
| def _job_name(key: str) -> str: |
| return ("v1_" + key).replace("policyanchor", "pol")[:20] |
|
|
|
|
| def _submit(cmd: list[str], *, project_dir: Path, dry_run: bool) -> str: |
| if dry_run: |
| print("DRY RUN:", " ".join(cmd)) |
| return "dry-run" |
| result = subprocess.run( |
| cmd, |
| cwd=project_dir, |
| check=True, |
| text=True, |
| capture_output=True, |
| ) |
| for token in result.stdout.split(): |
| if token.isdigit(): |
| return token |
| raise RuntimeError(f"Could not parse sbatch job id from: {result.stdout}") |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|