import torch

from alpamayo_lite.diffusion import FlowMatching


def test_flow_matching_target_velocity() -> None:
    flow = FlowMatching((4, 2), inference_steps=2)
    target = torch.ones(1, 4, 2)
    data = flow.construct_training_data(target)
    exact_velocity = data["target_action"] - data["noise"]
    assert float(flow.loss(data, exact_velocity)) == 0.0