Spaces:
Running
Running
| 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() | |