import json import tempfile import unittest from pathlib import Path import numpy as np from app.lerobot_export import (_ensure_slam_pose, _ensure_observation_state, _force_matrix_semantics, _write_stats) from app.umi_slam_action import poses_from_trajectory def _sample(t: float, x: float, z: float) -> list[float]: """一组 slam_trajectory 样本:[t, x, y, z, qx, qy, qz, qw]。""" return [float(t), float(x), 0.0, float(z), 0.0, 0.0, 0.0, 1.0] def _rows(episode: int, count: int, *, t0: float, slam_offset: float = 0.0) -> list[dict]: """造一集 rows:每帧两个 slam 样本,位姿随帧号线性推进。""" rows = [] for index in range(count): stamp = t0 + index / 30.0 rows.append({ "episode_index": episode, "timestamp": stamp, "observation.slam_trajectory": [ _sample(stamp - slam_offset, index * 0.01, index * 0.002), _sample(stamp - slam_offset + 0.01, index * 0.01 + 0.001, index * 0.002), ], }) return rows class EnsureSlamPoseTests(unittest.TestCase): def test_derives_column_and_matches_direct_reconstruction(self): """整列缺失时补出 slam_pose,且值与单独调用重建函数逐值一致。""" rows = _rows(0, 40, t0=100.0) self.assertEqual(_ensure_slam_pose(rows), "observation.slam_pose") self.assertTrue(all(len(row["observation.slam_pose"]) == 7 for row in rows)) expected = poses_from_trajectory( [row["observation.slam_trajectory"] for row in rows], np.array([row["timestamp"] for row in rows]), len(rows)) actual = np.array([row["observation.slam_pose"] for row in rows]) self.assertTrue(np.allclose(actual, expected, atol=1e-12)) def test_episodes_are_fitted_independently(self): """时间基映射必须每集单独拟合,否则跨集求差会把原点差异当运动。 两集的 slam 时钟相差一个很大的偏移量;若混在一起拟合,时间轴被拉 成两段,重建值会偏离该集单独重建的结果。 """ first = _rows(0, 30, t0=100.0, slam_offset=0.0) second = _rows(1, 30, t0=500.0, slam_offset=400.0) rows = first + second self.assertEqual(_ensure_slam_pose(rows), "observation.slam_pose") for group, offset in ((first, 0.0), (second, 400.0)): expected = poses_from_trajectory( [row["observation.slam_trajectory"] for row in group], np.array([row["timestamp"] for row in group]), len(group), ) actual = np.array([row["observation.slam_pose"] for row in group]) self.assertTrue(np.allclose(actual, expected, atol=1e-12)) def test_existing_column_is_left_untouched(self): """旧批两列都有 —— 重导时不得覆盖采集端原始 slam_pose。""" rows = _rows(0, 10, t0=100.0) original = [9.0] * 7 for row in rows: row["observation.slam_pose"] = list(original) self.assertIsNone(_ensure_slam_pose(rows)) self.assertTrue(all(row["observation.slam_pose"] == original for row in rows)) def test_skips_batches_without_trajectory(self): """手套 / D435 批次没有 slam_trajectory —— 不得凭空写入该列。""" rows = [{"episode_index": 0, "timestamp": 0.0, "observation.tactile.left": [0.0] * 256}] self.assertIsNone(_ensure_slam_pose(rows)) self.assertNotIn("observation.slam_pose", rows[0]) def test_missing_timestamp_does_not_break_reconstruction(self): """时间戳缺帧(None)不能把 np.asarray 推断成 object 而抛异常。""" rows = _rows(0, 20, t0=200.0) rows[5]["timestamp"] = None self.assertEqual(_ensure_slam_pose(rows), "observation.slam_pose") self.assertTrue(all(np.all(np.isfinite(row["observation.slam_pose"])) for row in rows)) def test_derived_pose_feeds_observation_state(self): """串联顺序:补出的 slam_pose 要能继续供 state 派生使用。 新批里 ep72/ep73 两列都缺,只补 slam_pose 而不让 state 跟上,导出 产物仍然缺 observation.state,ACT 依旧 KeyError。 """ rows = _rows(0, 25, t0=300.0) _ensure_slam_pose(rows) self.assertEqual(_ensure_observation_state(rows), "observation.state") self.assertTrue(all(len(row["observation.state"]) == 4 for row in rows)) # 位置分量是本帧相对本集第一帧的偏移 —— 首帧必须为原点 self.assertTrue(np.allclose(rows[0]["observation.state"][:3], 0.0)) class ForceMatrixSemanticsTests(unittest.TestCase): """源 info.json → 导出 features 的力矩阵语义元数据搬运。 这些字段(scale/encoding)决定数值怎么还原:力矩阵存的是「放大 scale 倍 再取行内差分」的整数,少任何一个,产物里的力矩阵都无法解读。它们只写在 info.json 里,parquet 本身不带,所以导出必须显式搬运。 """ def _root(self, features: dict) -> Path: root = Path(tempfile.mkdtemp()) (root / "meta").mkdir(parents=True) (root / "meta" / "info.json").write_text( json.dumps({"features": features}), encoding="utf-8") return root def test_carries_scale_and_encoding(self): root = self._root({ "observation.gripper_left_force_matrix": { "dtype": "int16", "shape": [250, 250, 3], "encoding": "row_diff_quantized", "scale": 100, "names": ["fx", "fy", "fz"], "units": ["mN", "mN", "mN"], }, }) self.assertEqual(_force_matrix_semantics(root), { "observation.gripper_left_force_matrix": { "encoding": "row_diff_quantized", "scale": 100, "names": ["fx", "fy", "fz"], "units": ["mN", "mN", "mN"], }, }) def test_does_not_carry_shape_or_dtype(self): """shape/dtype 由导出自己按实际写入值声明,不能被源覆盖。 源声明 shape [250,250,3] 是逻辑形状,而 parquet 里存的是展平的 187500 长列表;照搬声明会让官方加载器按 3 维 reshape 而失败。 """ root = self._root({ "observation.gripper_left_force_matrix": { "dtype": "int16", "shape": [250, 250, 3], "scale": 100, }, }) carried = _force_matrix_semantics(root)["observation.gripper_left_force_matrix"] self.assertNotIn("shape", carried) self.assertNotIn("dtype", carried) def test_ns_columns_keep_their_encoding(self): """_ns 采集时刻列同样靠 encoding 声明时基,int64 裸列无法解读。""" root = self._root({ "observation.gripper_left_force_matrix_ns": { "dtype": "int64", "encoding": "host_monotonic_ns", }, }) self.assertEqual( _force_matrix_semantics(root)["observation.gripper_left_force_matrix_ns"], {"encoding": "host_monotonic_ns"}) def test_ignores_unrelated_features(self): root = self._root({ "observation.gripper_state": {"dtype": "float32", "scale": 100}, "action": {"dtype": "float32"}, }) self.assertEqual(_force_matrix_semantics(root), {}) def test_missing_or_broken_source_degrades_to_empty(self): """源元数据读不到时返回空字典 —— 导出照常进行,不因缺元数据中断。""" self.assertEqual(_force_matrix_semantics(Path("/nonexistent/nope")), {}) broken = Path(tempfile.mkdtemp()) (broken / "meta").mkdir() (broken / "meta" / "info.json").write_text("{not json", encoding="utf-8") self.assertEqual(_force_matrix_semantics(broken), {}) class WriteStatsDataSourceTests(unittest.TestCase): """stats.json 的两种取数路径必须产出同一份统计量。 ``_write_stats`` 原先一律把整个 data/ 目录从磁盘重读一遍;调用方手里 其实就有刚写完的 arrow table,传进去可省掉这次全量往返。但两条路径的 DataFrame 必须逐列同 dtype、同值 —— 否则归一化统计会悄悄漂移,而训练 侧不会报错,只会训出偏差。 """ def _build(self, root: Path): """造一份含标量列/定长向量列/字符串列的 data/,返回写入的 table。""" import pyarrow as pa import pyarrow.parquet as pq rows = [ {"frame_index": i, "timestamp": i / 30.0, # 标量浮点 "observation.state": float(i % 3), # 标量浮点 "action": [float(i), float(i) * 2, -1.0], # 定长向量(走逐维 stats) "gesture": "open" if i % 2 else ""} # 字符串(应被跳过) for i in range(20) ] table = pa.Table.from_pylist(rows) data_dir = root / "data" / "chunk-000" data_dir.mkdir(parents=True, exist_ok=True) pq.write_table(table, data_dir / "file-000.parquet") return table def test_table_path_matches_disk_path(self): disk_root = Path(tempfile.mkdtemp()) table_root = Path(tempfile.mkdtemp()) table = self._build(disk_root) self._build(table_root) _write_stats(disk_root) # 旧路径:从磁盘重读 _write_stats(table_root, None, table) # 新路径:直接用写入的表 from_disk = json.loads((disk_root / "meta" / "stats.json").read_text()) from_table = json.loads((table_root / "meta" / "stats.json").read_text()) self.assertEqual(from_disk, from_table) # 别退化成"两边都空"的假通过 self.assertIn("observation.state", from_disk) self.assertIn("action", from_disk) def test_string_columns_are_not_given_numeric_stats(self): """官方加载器会把 stats 每一项转张量,字符串列必须跳过。""" root = Path(tempfile.mkdtemp()) table = self._build(root) _write_stats(root, None, table) stats = json.loads((root / "meta" / "stats.json").read_text()) self.assertNotIn("gesture", stats) def test_missing_data_dir_degrades_quietly(self): """没有 data/ 时直接返回,不抛异常(旧路径原有的容错)。""" root = Path(tempfile.mkdtemp()) _write_stats(root) self.assertFalse((root / "meta" / "stats.json").exists()) if __name__ == "__main__": unittest.main()