tests/trex_track_force/test_dataset_loader.py

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)