File size: 4,807 Bytes
ec0a9aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """One-shot migration: fix bimanual/humanoid multiview Ctrl-World layout.
The bimanual/humanoid inference runs saved flat episode dirs like
`.../ctrlworld/episode_makovian_close_toolbox__000026/`, baking the
makovian/non_makovian category into the episode name. Per the
sampling-dataset-layout rule, `{category}` must be its own path segment:
`.../ctrlworld/makovian/episode_close_toolbox__000026/`.
This moves existing output + input episode dirs into the correct
`{category}/episode_<name>` layout, deletes the stale top-level
`all_summary.json`, and regenerates a per-category `all_summary.json` from the
moved output `metrics.json` files.
Usage:
python scripts/migrate_ctrlworld_multiview_layout.py # dry run
python scripts/migrate_ctrlworld_multiview_layout.py --apply # do it
"""
import argparse
import json
import shutil
from pathlib import Path
ROOT = Path("/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics")
EMBODIMENTS = ["bimanual", "humanoid"]
IOS = ["output", "input"]
def split_category(name):
if name.startswith("non_makovian_"):
return "non_makovian", name[len("non_makovian_"):]
if name.startswith("makovian_"):
return "makovian", name[len("makovian_"):]
return None, name
def migrate_dir(base: Path, apply: bool):
"""Move flat episode_<category>_<rest> dirs into <category>/episode_<rest>."""
if not base.exists():
print(f" [skip] {base} does not exist")
return 0
moved = 0
for child in sorted(base.iterdir()):
if not child.is_dir() or not child.name.startswith("episode_"):
continue
stem = child.name[len("episode_"):]
category, rest = split_category(stem)
if category is None:
print(f" [warn] no category prefix, leaving as-is: {child.name}")
continue
dst = base / category / f"episode_{rest}"
if dst.exists():
print(f" [conflict] dst exists, skipping: {dst}")
continue
print(f" {child.name} -> {category}/episode_{rest}")
if apply:
dst.parent.mkdir(parents=True, exist_ok=True)
shutil.move(str(child), str(dst))
moved += 1
return moved
def regenerate_summaries(output_base: Path, apply: bool):
"""Rebuild per-category all_summary.json from moved metrics.json files."""
if not output_base.exists():
return
for category_dir in sorted(output_base.iterdir()):
if not category_dir.is_dir() or category_dir.name not in ("makovian", "non_makovian"):
continue
metrics_list = []
for ep_dir in sorted(category_dir.iterdir()):
mpath = ep_dir / "metrics.json"
if mpath.exists():
with open(mpath) as f:
metrics_list.append(json.load(f))
if not metrics_list:
continue
summary = {
"mean_psnr": sum(m["psnr"] for m in metrics_list) / len(metrics_list),
"mean_ssim": sum(m["ssim"] for m in metrics_list) / len(metrics_list),
"mean_lpips": sum(m["lpips"] for m in metrics_list) / len(metrics_list),
"num_episodes": len(metrics_list),
}
out = category_dir / "all_summary.json"
print(f" summary {out.relative_to(output_base)}: "
f"n={summary['num_episodes']} PSNR={summary['mean_psnr']:.3f}")
if apply:
with open(out, "w") as f:
json.dump(summary, f, indent=2)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--apply", action="store_true", help="Actually move files (default: dry run).")
args = ap.parse_args()
mode = "APPLY" if args.apply else "DRY RUN"
print(f"=== Ctrl-World multiview layout migration [{mode}] ===\n")
for emb in EMBODIMENTS:
for io in IOS:
base = ROOT / "sampling_dataset" / "dense" / emb / io / "multiview" / "ctrlworld"
print(f"[{emb}/{io}] {base}")
n = migrate_dir(base, args.apply)
print(f" -> {n} dirs to move")
stale = base / "all_summary.json"
if stale.exists():
print(f" [stale] removing top-level {stale.name} (regenerated per-category)")
if args.apply:
stale.unlink()
print()
print("=== Regenerating per-category summaries (output only) ===")
for emb in EMBODIMENTS:
output_base = ROOT / "sampling_dataset" / "dense" / emb / "output" / "multiview" / "ctrlworld"
print(f"[{emb}/output]")
regenerate_summaries(output_base, args.apply)
print()
if not args.apply:
print("Dry run complete. Re-run with --apply to perform the migration.")
if __name__ == "__main__":
main()
|