vla / workspace /scripts /advance_v1_generator.py
anhtld's picture
auto-sync 2026-07-02T19:07:26Z workspace (part 6)
828bff4 verified
Raw
History Blame Contribute Delete
24.6 kB
#!/usr/bin/env python3
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())