import torch

from alpamayo_lite.action_space import RelativeDeltaActionSpace


def test_action_round_trip() -> None:
    history = torch.tensor([[[-1.0, 0.0], [0.0, 0.0]]])
    future = torch.tensor([[[1.0, 0.2], [2.1, 0.5], [3.2, 0.9]]])
    statistics = RelativeDeltaActionSpace.fit(history, future)
    action_space = RelativeDeltaActionSpace(statistics.mean, statistics.scale)
    action = action_space.trajectory_to_action(history, future)
    reconstructed = action_space.action_to_trajectory(action, history)
    torch.testing.assert_close(reconstructed, future)