#!/usr/bin/env python3 """test_state_source.py — 1 コマごとの state の出どころ (state_source. 依頼 196) の確かめ. python3 pi05/test_state_source.py -v 前半はファイルを読まない部分だけ. 後半 (TestSameAsBefore) は KIOXIA の runs/pi05_policy を読み, 直す前のコード (git の HEAD) と今のコードで /tmp に作り, state / actions / t がバイト単位で同じことを見る (KIOXIA がなければ飛ばす. データセットにもアームにも触らない). """ import importlib.util import json import subprocess import sys import tempfile import unittest from pathlib import Path from types import SimpleNamespace import numpy as np _HERE = Path(__file__).resolve().parent sys.path.insert(0, str(_HERE)) import build_correction_episodes as bce # noqa: E402 def samples(still=False): """0〜1 秒は 0.1 秒おき → 1.0〜2.6 秒は抜け (j1 が 6° 動く. still なら動かない) → 2.6〜3.6 秒は 0.1 秒おき.""" a = np.round(np.arange(0.0, 1.01, 0.1), 3) b = np.round(np.arange(2.6, 3.61, 0.1), 3) at = np.concatenate([a, b]) aq = np.zeros((len(at), 6)) if not still: aq[len(a):, 1] = 6.0 return at, aq def runs_of_nonzero(src): nz = np.r_[0, (np.asarray(src) != 0).astype(int)] return int((np.diff(nz) == 1).sum()) class TestFlags(unittest.TestCase): def test_motion_gap_is_1(self): at, aq = samples() t = np.arange(0.0, 3.6, 0.1) src = bce.gap_fill_flags(t, at, aq) self.assertEqual(src.dtype, np.int8) inside = (t > 1.0) & (t < 2.6) self.assertTrue(np.all(src[inside] == bce.SRC_FILLED_MOVE)) self.assertTrue(np.all(src[~inside] == bce.SRC_MEASURED)) def test_still_gap_is_2(self): at, aq = samples(still=True) t = np.arange(0.0, 3.6, 0.1) src = bce.gap_fill_flags(t, at, aq) self.assertEqual(set(np.unique(src)), {bce.SRC_MEASURED, bce.SRC_FILLED_STILL}) def test_flags_mark_exactly_what_speed20_changes(self): """speed20 が線形補間から変えたコマは, すべて印が 0 でないコマ.""" at, aq = samples() t = np.arange(0.0, 3.6, 0.1) q_lin, _ = bce.fill_motion_gaps(t, at, aq, mode="linear") q_sp, _ = bce.fill_motion_gaps(t, at, aq, mode="speed20") src = bce.gap_fill_flags(t, at, aq) changed = np.any(np.abs(q_lin - q_sp) > 1e-9, axis=1) self.assertTrue(changed.any()) self.assertTrue(np.all(src[changed] != bce.SRC_MEASURED)) def test_no_gap_no_flag(self): at = np.round(np.arange(0.0, 3.01, 0.1), 3) aq = np.cumsum(np.ones((len(at), 6)), axis=0) self.assertEqual(int(bce.gap_fill_flags(np.arange(0.0, 3.0, 0.1), at, aq).sum()), 0) def test_compat_ignores_state_source(self): """お手本にない state_source のキーがあっても check_compat はキーの違いとして数えない.""" with tempfile.TemporaryDirectory() as d: import cv2 for name, extra in (("ours", True), ("ref", False)): ep = Path(d) / name / "e0" (ep / "frames").mkdir(parents=True) cv2.imwrite(str(ep / "frames" / "000000.jpg"), np.zeros((4, 6, 3), np.uint8)) kw = {"state_source": np.zeros(1, np.int8)} if extra else {} np.savez_compressed(ep / "episode.npz", t=np.zeros(1), t_orig=np.zeros(1), state=np.zeros((1, 7), np.float32), actions=np.zeros((1, 7), np.float32), **kw) (ep / "meta.json").write_text(json.dumps({"fps": 10, "image_size": [6, 4], "state_dim": 7, "state_names": [], "task": "x"})) self.assertEqual(bce.check_compat(Path(d) / "ours", Path(d) / "ref"), []) def _old_module(): """直す前の build_correction_episodes.py (git の HEAD) を別の名前で読み込む.""" src = subprocess.run(["git", "show", "HEAD:pi05/build_correction_episodes.py"], cwd=_HERE, capture_output=True, text=True, check=True).stdout tmp = Path(tempfile.mkdtemp()) / "bce_old.py" tmp.write_text(src) spec = importlib.util.spec_from_file_location("bce_old", tmp) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod # 補正の多い本・巻き戻しのある本など. 依頼 196 の時点で KIOXIA にあるもの SAMPLE_RUNS = ("pi05_20261031_175527", "pi05_20261031_180435") @unittest.skipUnless(Path("/Volumes/KIOXIA").is_mount() and bce.RUNS_ROOT.is_dir(), "KIOXIA がない") class TestSameAsBefore(unittest.TestCase): @classmethod def setUpClass(cls): cls.old = _old_module() def _build(self, mod, run, d, gap_fill, trim_hold): args = SimpleNamespace(gap_fill=gap_fill, trim_hold=trim_hold, no_trim_mid_idle=False, max_protrude_px=None, force=True) out = Path(d) / mod.__name__ / f"{gap_fill}_{trim_hold}" / run res = mod.build_one(bce.RUNS_ROOT / run, out, args) self.assertTrue(res.get("used"), res) z = np.load(out / "episode.npz") return {k: z[k] for k in z.files}, json.loads((out / "meta.json").read_text()) def test_bytes_and_counts(self): with tempfile.TemporaryDirectory() as d: for run in SAMPLE_RUNS: if not (bce.RUNS_ROOT / run).is_dir(): self.skipTest(f"{run} がない") for gap_fill, trim_hold in (("linear", False), ("speed20", False), ("speed20", True)): with self.subTest(run=run, gap_fill=gap_fill, trim_hold=trim_hold): a, _ = self._build(self.old, run, d, gap_fill, trim_hold) b, meta = self._build(bce, run, d, gap_fill, trim_hold) self.assertNotIn("state_source", a) for k in ("t", "t_orig", "state", "actions"): self.assertEqual(a[k].dtype, b[k].dtype, k) self.assertEqual(a[k].tobytes(), b[k].tobytes(), k) src = b["state_source"] self.assertEqual(src.dtype, np.int8) self.assertEqual(len(src), len(b["state"])) self.assertTrue(set(np.unique(src)) <= set(bce.STATE_SOURCE_CODES)) gf = meta["gap_fill"] filled = int((src != 0).sum()) self.assertEqual(meta["filled_steps"], filled) self.assertEqual(meta["filled_motion_steps"], int((src == 1).sum())) self.assertAlmostEqual(meta["filled_ratio"], filled / len(src), places=4) # 印のつながりの数は抜けの数を超えない (1 つの抜けが印を 2 つに分けることはない) self.assertLessEqual(runs_of_nonzero(src), gf["gaps"]) self.assertLessEqual(runs_of_nonzero(src == 1), gf["motion_gaps"]) if gf["gaps"] == 0: self.assertEqual(filled, 0) if __name__ == "__main__": unittest.main()