from types import SimpleNamespace import numpy as np import pandas as pd from groot.vla.data.dataset.lerobot_sharded import ( ShardedLeRobotSubLangSingleActionChunkDatasetDROID, ) from groot.vla.model.trex_track_force.dataset import ( AR_BLOCKS, FORCE_COLUMN, TRACK_VISIBILITY_COLUMN, TRACK_XY_COLUMN, TrexTrackForceShardedDataset, eef62_delta_base, nearest_timestamp_indices, uniform_target_times, ) from groot.vla.model.trex_track_force.runtime import ( TrexRuntimeStatistics, delta_base_to_absolute, ) def _pose9(translation: tuple[float, float, float], rotation: np.ndarray) -> np.ndarray: return np.concatenate((np.asarray(translation), rotation[:, 0], rotation[:, 1])) def test_nearest_timestamp_indices_builds_16_step_30_to_20_grid() -> None: source = np.arange(61, dtype=np.float64) / 30.0 target = uniform_target_times(0.0, range(16), 20.0) selection = nearest_timestamp_indices(source, target) expected = np.floor(np.arange(16, dtype=np.float64) * 1.5).astype(np.int64) np.testing.assert_array_equal(selection.indices, expected) assert not selection.padding_mask.any() assert target[-1] == 0.75 def test_nearest_timestamp_indices_marks_history_edge_padding() -> None: source = np.arange(31, dtype=np.float64) / 30.0 target = uniform_target_times(0.0, range(-15, 1), 5.0) selection = nearest_timestamp_indices(source, target) assert selection.padding_mask[:-1].all() assert not selection.padding_mask[-1] assert np.all(selection.indices[:-1] == 0) def test_step_filter_rejects_float_roundoff_future_padding(monkeypatch) -> None: timestamps = np.arange(1126, dtype=np.float64) / 30.0 dataset = object.__new__(TrexTrackForceShardedDataset) dataset.action_rate_hz = 20.0 dataset.video_rate_hz = 10.0 dataset.num_ar_blocks = AR_BLOCKS dataset._trajectory_ids = np.array([0], dtype=np.int64) dataset._logical_anchor_times = {} monkeypatch.setattr( ShardedLeRobotSubLangSingleActionChunkDatasetDROID, "_get_step_filter", lambda self: {0: np.arange(timestamps.size, dtype=np.int64)}, ) monkeypatch.setattr(dataset, "get_parquet_path", lambda trajectory_id: "unused") monkeypatch.setattr( pd, "read_parquet", lambda *args, **kwargs: pd.DataFrame({"timestamp": timestamps}), ) filtered = dataset._get_step_filter() assert 1027 in filtered[0] assert 1029 not in filtered[0] assert 1029 not in dataset._logical_anchor_times[0] def test_column_array_collapses_pyarrow_nested_track_objects() -> None: row = np.empty(2, dtype=object) row[:] = ( np.asarray([0.1, 0.2], dtype=np.float32), np.asarray([0.3, 0.4], dtype=np.float32), ) frame = pd.DataFrame({TRACK_XY_COLUMN: [row, row.copy()]}) values = TrexTrackForceShardedDataset._column_array(frame, TRACK_XY_COLUMN) assert values.shape == (2, 2, 2) assert values.dtype == np.float32 np.testing.assert_allclose(values[0, 1], [0.3, 0.4]) def test_loader_preserves_four_ar_blocks_with_16_actions_each() -> None: frames = 121 timestamps = np.arange(frames, dtype=np.float64) / 30.0 identity_rot6d = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) eef = np.zeros((frames, 62), dtype=np.float32) eef[:, 3:9] = identity_rot6d eef[:, 34:40] = identity_rot6d eef[:, 0] = np.arange(frames, dtype=np.float32) / 30.0 tracks = np.zeros((frames, 250, 2), dtype=np.float32) tracks[..., 0] = np.arange(frames, dtype=np.float32)[:, None] / frames frame = pd.DataFrame( { "timestamp": timestamps, "observation.state_eef62": list(eef), "action.eef62_absolute": list(eef), TRACK_XY_COLUMN: list(tracks), TRACK_VISIBILITY_COLUMN: list( np.ones((frames, 250), dtype=np.float32) ), FORCE_COLUMN: list( np.zeros((frames, 60), dtype=np.float32) ), } ) dataset = object.__new__(TrexTrackForceShardedDataset) dataset.action_rate_hz = 20.0 dataset.tactile_rate_hz = 5.0 dataset.video_rate_hz = 10.0 dataset.num_ar_blocks = AR_BLOCKS dataset._logical_anchor_times = {0: {0: 0.0}} dataset._force_low = np.zeros((10, 6), dtype=np.float32) dataset._force_high = np.ones((10, 6), dtype=np.float32) dataset._lerobot_modality_meta = SimpleNamespace( state={ "eef62": SimpleNamespace( original_key="observation.state_eef62", start=0, end=62 ) }, action={ "eef62": SimpleNamespace( original_key="action.eef62_absolute", start=0, end=62 ) }, ) dataset.cached_shard = { key: np.zeros((frames, 4, 6, 3), dtype=np.uint8) for key in ("video.head_left", "video.left_wrist", "video.right_wrist") } dataset.shard_start_indices = {0: 0} dataset._modality_keys = { "video": list(dataset.cached_shard), "language": [], } dataset.get_trajectory_data = lambda trajectory_id: frame sample = dataset.get_step_data(0, {"action.eef62": np.array([0])}) assert sample["action.eef62"].shape == (4 * 16, 62) assert sample["state.eef62"].shape == (4, 62) assert sample["track_past_xy"].shape == (4, 16, 250, 2) assert sample["track_future_xy"].shape == (4, 16, 250, 2) np.testing.assert_allclose( sample["track_future_xy"][:, 0], sample["track_past_xy"][:, -1] ) assert not sample["track_past_visibility"][0, :-1].any() assert sample["track_past_visibility"][0, -1].all() assert sample["current_force"].shape == (4, 4, 10, 6) assert sample["force_history"].shape == (4, 4, 16, 10, 6) assert np.count_nonzero(sample["force_history"][0, 0, :-1]) == 0 np.testing.assert_array_equal(sample["current_force"][0, 0], -1.0) assert sample["force_history_padding_mask"].shape == (4, 4, 16) for key in dataset.modality_keys["video"]: assert sample[key].shape == (33, 4, 6, 3) def test_training_and_runtime_force_normalization_are_identical() -> None: low = np.zeros((10, 6), dtype=np.float32) high = np.ones((10, 6), dtype=np.float32) high[0, 0] = 0 force = np.linspace(-0.5, 1.5, 120, dtype=np.float32).reshape(2, 10, 6) dataset = object.__new__(TrexTrackForceShardedDataset) dataset._force_low = low dataset._force_high = high runtime = TrexRuntimeStatistics( action_q01=np.zeros(62), action_q99=np.ones(62), state_q01=np.zeros(62), state_q99=np.ones(62), force_q01=low, force_q99=high, ) np.testing.assert_allclose( dataset._normalize_force(force), runtime.normalize_force(force) ) def test_eef62_delta_base_matches_trex_chunk_start_frame_math() -> None: identity = np.eye(3) angle = np.pi / 2 reference_rotation = np.array( [ [np.cos(angle), -np.sin(angle), 0.0], [np.sin(angle), np.cos(angle), 0.0], [0.0, 0.0, 1.0], ] ) reference = np.zeros(62, dtype=np.float32) reference[:9] = _pose9((0.0, 0.0, 0.0), reference_rotation) reference[31:40] = _pose9((0.0, 0.0, 0.0), identity) targets = np.zeros((2, 62), dtype=np.float32) targets[:, :9] = _pose9((1.0, 0.0, 0.0), reference_rotation) targets[:, 31:40] = _pose9((0.0, 2.0, 0.0), identity) targets[:, 9:31] = 0.25 targets[:, 40:62] = -0.5 delta = eef62_delta_base(reference, targets) np.testing.assert_allclose(delta[:, :3], [[0.0, -1.0, 0.0]] * 2, atol=1e-6) np.testing.assert_allclose( delta[:, 3:9], [[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]] * 2, atol=1e-6 ) np.testing.assert_allclose(delta[:, 31:34], [[0.0, 2.0, 0.0]] * 2) np.testing.assert_allclose(delta[:, 9:31], 0.25) np.testing.assert_allclose(delta[:, 40:62], -0.5) reconstructed = delta_base_to_absolute(reference, delta) np.testing.assert_allclose(reconstructed, targets, atol=1e-6)