#!/usr/bin/env python3 """test_merge_and_config.py — 依頼 45 で足した部分の確かめ (KIOXIA にも Pod にも触らない. 一時フォルダだけ). python3 pi05/test_merge_and_config.py 中身: - merge_raw_episodes: 2 つの raw_episodes を写してまとめる / 元は変えない / ._* を写さない・tar に入れない / 形が合わない・名前が重なる・出力先がある ときは止まる - ._*.jpg (exFAT の AppleDouble) があっても check_compat が画像の枚数を数え違えない - _rmtree: 本体と一緒に ._ が消えても止まらない - add_trim_corr_config: trim2 の設定を写してデータ・ステップ数・初期の重み・学習率の予定だけ替える / 二重に入れない / 違う引数での入れ直しは止まる / 前の版 (corr_v1・corr_v2) が入っていたら止まる - pod_prep_trim_corr.sh: 正規化の値を計算し直さず trim2 の値を置く (依頼 48) """ import ast import json import subprocess import sys import tarfile import tempfile import unittest from pathlib import Path from unittest import mock import cv2 import numpy as np _HERE = Path(__file__).resolve().parent sys.path.insert(0, str(_HERE)) sys.path.insert(0, str(_HERE / "openpi_addon")) import add_trim_corr_config as atc # noqa: E402 import build_correction_episodes as bce # noqa: E402 import merge_raw_episodes as mre # noqa: E402 def write_ep(d, n=5, source="teach", size=(384, 216), appledouble=False): (d / "frames").mkdir(parents=True) for i in range(n): cv2.imwrite(str(d / "frames" / f"{i:06d}.jpg"), np.full((size[1], size[0], 3), 90, np.uint8)) if appledouble: (d / "frames" / f"._{i:06d}.jpg").write_bytes(b"\x00\x05\x16\x07") st = np.zeros((n, 7), np.float32) np.savez_compressed(d / "episode.npz", t=np.arange(n) / 10.0, t_orig=np.arange(n) / 10.0, state=st, actions=st) meta = {"task": "put the red block in the white box", "fps": 10, "steps": n, "image_size": list(size), "state_dim": 7, "state_names": ["j0", "j1", "j2", "j3", "j4", "j5", "gripper"], "source": source} if source.startswith("pi05_policy"): meta.update(correct_steps=2, model_steps=n - 2) (d / "meta.json").write_text(json.dumps(meta)) if appledouble: (d / "._episode.npz").write_bytes(b"\x00\x05\x16\x07") def run_merge(root, *extra): argv = ["merge_raw_episodes.py", "--data-root", str(root), "--src", "a", "--src", "b", "--dataset", "ab", *extra] with mock.patch.object(sys, "argv", argv): return mre.main() class TestMerge(unittest.TestCase): def setUp(self): self._td = tempfile.TemporaryDirectory() self.root = Path(self._td.name) ds = self.root / "datasets" write_ep(ds / "a" / "raw_episodes" / "run_1", n=5) write_ep(ds / "a" / "raw_episodes" / "run_2", n=6) write_ep(ds / "b" / "raw_episodes" / "pi05_1", n=4, source="pi05_policy (model + correct)", appledouble=True) def tearDown(self): self._td.cleanup() def test_merge_copies_and_counts(self): before = sorted(p.relative_to(self.root) for p in (self.root / "datasets").rglob("*")) self.assertEqual(run_merge(self.root, "--tar"), 0) out = self.root / "datasets" / "ab" self.assertEqual(sorted(p.name for p in (out / "raw_episodes").iterdir()), ["pi05_1", "run_1", "run_2"]) self.assertEqual(list((out / "raw_episodes").rglob("._*")), []) # AppleDouble は写さない idx = json.loads((out / "index.json").read_text()) self.assertEqual((idx["total_episodes"], idx["total_steps"]), (3, 15)) self.assertEqual(idx["sources"]["b"]["correct_steps"], 2) self.assertEqual(idx["repo_id"], "local/ab") self.assertTrue((out / "code" / "merge_raw_episodes.py").exists()) # データと対のコード with tarfile.open(out / "ab_raw_episodes.tar") as t: names = t.getnames() self.assertFalse([n for n in names if Path(n).name.startswith("._")]) self.assertEqual(sum(n.endswith(".jpg") for n in names), 15) # 元は変わらない (増えたのは ab/ の下だけ) after = sorted(p.relative_to(self.root) for p in (self.root / "datasets").rglob("*") if not str(p.relative_to(self.root / "datasets")).startswith("ab")) self.assertEqual(before, after) def test_refuses_existing_and_duplicates(self): self.assertEqual(run_merge(self.root), 0) with self.assertRaises(SystemExit): run_merge(self.root) # もうある self.assertEqual(run_merge(self.root, "--force"), 0) write_ep(self.root / "datasets" / "b" / "raw_episodes" / "run_1", n=3) # 名前が重なる with self.assertRaises(SystemExit): run_merge(self.root, "--force") def test_refuses_mismatch_without_writing(self): write_ep(self.root / "datasets" / "b" / "raw_episodes" / "pi05_2", n=3, size=(320, 180)) with mock.patch("sys.stdout"): self.assertEqual(run_merge(self.root), 1) self.assertFalse((self.root / "datasets" / "ab").exists()) def test_check_compat_ignores_appledouble(self): ds = self.root / "datasets" self.assertEqual(bce.check_compat(ds / "b" / "raw_episodes", ds / "a" / "raw_episodes"), []) def test_rmtree_tolerates_vanished_files(self): d = self.root / "x" (d / "sub").mkdir(parents=True) (d / "sub" / "f").write_text("1") real_unlink = __import__("os").unlink def unlink(path, *a, **k): # 1 回目で ._ 相当の 2 つ目も消えたことにする real_unlink(path, *a, **k) raise FileNotFoundError(path) with mock.patch("os.unlink", unlink): bce._rmtree(d) self.assertFalse(d.exists()) # install_addon.py の TRAIN を元に add_trim_config.py (lab/Log11/pod) と同じ置き換えをした trim2 の設定 def fake_config(root): src = (_HERE / "openpi_addon" / "install_addon.py").read_text() train = src[src.index("TRAIN = f'''") + len("TRAIN = f'''"):] train = train[:train.index("'''")].replace("{MARK}", "# >>> r2r2r mecharm addon") trim = (train.replace("# >>> r2r2r mecharm addon (train config)", "# >>> r2r2r mecharm teach_v1_trim2 (train config)") .replace("# <<< r2r2r mecharm addon (train config)", "# <<< r2r2r mecharm teach_v1_trim2 (train config)") .replace('name="pi05_mecharm_lora"', 'name="pi05_mecharm_trim_lora"') .replace('repo_id="local/mecharm_redblock_box_v1"', 'repo_id="local/mecharm_redblock_box_teach_v1_trim2"') .replace("num_train_steps=5_000,", "num_train_steps=4_000,")) cfg = root / "src" / "openpi" / "training" / "config.py" cfg.parent.mkdir(parents=True) cfg.write_text("_CONFIGS = [\n" + train + trim + "]\n") return cfg def configs(cfg): """config.py の TrainConfig(name=...) ごとの キーワード引数 (文字列のまま).""" tree = ast.parse(cfg.read_text()) out = {} for node in ast.walk(tree): if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "TrainConfig": kw = {k.arg: ast.unparse(k.value) for k in node.keywords} out[ast.literal_eval(kw["name"])] = kw return out class TestConfig(unittest.TestCase): def test_base_and_resume(self): for init, loader in (("base", "gs://openpi-assets/checkpoints/pi05_base/params"), ("/workspace/ckpt/3999/params", "/workspace/ckpt/3999/params")): with tempfile.TemporaryDirectory() as td, mock.patch("sys.stdout"): cfg = fake_config(Path(td)) self.assertEqual(atc.main([td, "--steps", "1500", "--init", init]), 0) c = configs(cfg) self.assertEqual(sorted(c), ["pi05_mecharm_lora", "pi05_mecharm_trim_corr_lora", "pi05_mecharm_trim_lora"]) new, trim = c["pi05_mecharm_trim_corr_lora"], c["pi05_mecharm_trim_lora"] self.assertIn("local/mecharm_redblock_box_teach_v1_trim2_corr_v3", new["data"]) # 依頼 48 self.assertEqual(new["num_train_steps"], "1500") self.assertIn(loader, new["weight_loader"]) # 依頼 48: 学習率の予定は続きから用に弱める. 下げきる先は trim2 と同じ lr = {k.arg: ast.literal_eval(k.value) for k in ast.parse(new["lr_schedule"]).body[0].value.keywords} lr_trim = {k.arg: ast.literal_eval(k.value) for k in ast.parse(trim["lr_schedule"]).body[0].value.keywords} self.assertEqual(lr, {"warmup_steps": 100, "peak_lr": 2.5e-5, "decay_steps": 1500, "decay_lr": lr_trim["decay_lr"]}) self.assertEqual(lr_trim, {"warmup_steps": 300, "peak_lr": 5e-5, "decay_steps": 5000, "decay_lr": 5e-6}) # trim2 の設定は変えない # ほかは trim2 と同じ for k in set(trim) - {"name", "data", "num_train_steps", "weight_loader", "lr_schedule"}: self.assertEqual(new[k], trim[k], k) self.assertEqual(trim["num_train_steps"], "4000") # trim2 の設定は変えない text = cfg.read_text() self.assertEqual(atc.main([td, "--steps", "1500", "--init", init]), 0) # 二重に入れない self.assertEqual(cfg.read_text(), text) with self.assertRaises(SystemExit): atc.main([td, "--steps", "4000", "--init", init]) # 違う引数では止まる # データ名・学習率の予定が違うものは止まる (codex r2 推奨: 5 つとも. 止まったときは書かない) for x, y in (("_trim2_corr_v3\"", "_trim2_corr_vX\""), ("warmup_steps=100,", "warmup_steps=300,"), ("peak_lr=2.5e-5,", "peak_lr=5e-5,"), ("decay_steps=1_500,", "decay_steps=5_000,"), ("decay_lr=5e-6,\n ),\n num_train_steps=1500", "decay_lr=1e-6,\n ),\n num_train_steps=1500")): bad = text.replace(x, y, 1) if x.startswith("decay_lr") else \ text[:text.index("pi05_mecharm_trim_corr_lora")] + \ text[text.index("pi05_mecharm_trim_corr_lora"):].replace(x, y, 1) self.assertNotEqual(bad, text, x) cfg.write_text(bad) with self.assertRaises(SystemExit): atc.main([td, "--steps", "1500", "--init", init]) self.assertEqual(cfg.read_text(), bad) def test_refuses_old_corr_blocks(self): """依頼 46 / 48: 前の版 (corr_v1・corr_v2) の設定が入っていたら足さずに止まる (同じ名前が 2 つにならない).""" for tag in ("teach_v1_trim2_corr_v1", "teach_v1_trim2_corr_v2"): with tempfile.TemporaryDirectory() as td, mock.patch("sys.stdout"): cfg = fake_config(Path(td)) s = cfg.read_text() cfg.write_text(s.replace("]\n", f" # >>> r2r2r mecharm {tag} (train config)\n" f" # <<< r2r2r mecharm {tag} (train config)\n]\n")) before = cfg.read_text() with self.assertRaises(SystemExit): atc.main([td, "--steps", "1500", "--init", "/workspace/ckpt/3999/params"]) self.assertEqual(cfg.read_text(), before) def test_requires_trim_first(self): with tempfile.TemporaryDirectory() as td: cfg = Path(td) / "src" / "openpi" / "training" / "config.py" cfg.parent.mkdir(parents=True) cfg.write_text("_CONFIGS = []\n") with self.assertRaises(AssertionError): atc.main([td, "--steps", "10", "--init", "base"]) def test_cli_requires_steps_and_init(self): r = subprocess.run([sys.executable, str(_HERE / "openpi_addon" / "add_trim_corr_config.py"), "/tmp"], capture_output=True, text=True) self.assertNotEqual(r.returncode, 0) class TestPodPrep(unittest.TestCase): """依頼 48: pod_prep_trim_corr.sh は正規化の値を計算し直さず, 3999 の学習の trim2 の値を置く.""" SH = _HERE / "openpi_addon" / "pod_prep_trim_corr.sh" MAC_NORM = Path("/Volumes/KIOXIA/r2r2r-data/checkpoints/pi05_mecharm_trim_lora_v1/3999/assets/local/" "mecharm_redblock_box_teach_v1_trim2/norm_stats.json") def lines(self): return [ln for ln in self.SH.read_text().splitlines() if not ln.lstrip().startswith("#")] def test_no_compute_norm_stats(self): self.assertFalse(any("compute_norm_stats" in ln for ln in self.lines())) def test_places_trim2_norm(self): code = "\n".join(self.lines()) self.assertIn("NAME=mecharm_redblock_box_teach_v1_trim2_corr_v3", code) self.assertIn("NORM_SRC=/workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json", code) self.assertIn("NORM_DST=assets/$CONFIG/local/$NAME/norm_stats.json", code) # openpi が学習で読む場所 self.assertIn('cp "$NORM_SRC" "$NORM_DST"', code) def test_sha_matches_mac_copy(self): import hashlib import re m = re.search(r"^NORM_SHA256=([0-9a-f]{64})", self.SH.read_text(), re.M) self.assertIsNotNone(m) if not self.MAC_NORM.exists(): self.skipTest("KIOXIA がない") self.assertEqual(hashlib.sha256(self.MAC_NORM.read_bytes()).hexdigest(), m.group(1)) def test_bash_syntax(self): self.assertEqual(subprocess.run(["bash", "-n", str(self.SH)]).returncode, 0) if __name__ == "__main__": unittest.main()