seeker
66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
import torch
|
|
|
|
from pi0mini.flow import FlowMatching
|
|
|
|
|
|
def _time_view(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
|
|
|
|
|
|
def test_training_batch_obeys_linear_path() -> None:
|
|
flow = FlowMatching()
|
|
action = torch.tensor([[[1.0, -2.0], [0.5, 3.0]], [[-1.0, 0.0], [2.5, -0.5]]])
|
|
x_t, t, velocity = flow.training_batch(
|
|
action, generator=torch.Generator().manual_seed(123)
|
|
)
|
|
|
|
noise = action - velocity
|
|
expected = (1.0 - _time_view(t, action)) * noise + _time_view(t, action) * action
|
|
torch.testing.assert_close(x_t, expected)
|
|
assert x_t.shape == action.shape
|
|
assert torch.all((t >= 0.0) & (t <= 1.0))
|
|
|
|
|
|
def test_training_batch_is_reproducible() -> None:
|
|
action = torch.randn(3, 16, 2)
|
|
flow = FlowMatching()
|
|
a = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
|
|
b = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
|
|
for left, right in zip(a, b):
|
|
torch.testing.assert_close(left, right)
|
|
|
|
|
|
def test_training_times_use_low_biased_shifted_beta() -> None:
|
|
flow = FlowMatching()
|
|
times = flow.sample_training_times(
|
|
100_000,
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
generator=torch.Generator().manual_seed(17),
|
|
)
|
|
|
|
assert torch.all(times >= 0.0)
|
|
assert torch.all(times <= 0.999)
|
|
# E[0.999 * (1 - Beta(1.5, 1))] = 0.3996. This also guards
|
|
# against accidentally reverting to a uniform or high-time-biased draw.
|
|
assert abs(float(times.mean()) - 0.3996) < 0.005
|
|
|
|
|
|
def test_loss_is_velocity_mse() -> None:
|
|
flow = FlowMatching()
|
|
target = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
|
|
prediction = target + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
|
|
torch.testing.assert_close(flow.loss(prediction, target), torch.tensor(1.5))
|
|
|
|
|
|
def test_exact_vector_field_integrates_to_action() -> None:
|
|
flow = FlowMatching(inference_steps=10)
|
|
action = torch.tensor([[[0.25, -0.5], [1.0, 2.0]], [[-2.0, 0.75], [0.0, 0.5]]])
|
|
|
|
def exact_field(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
|
return (action.to(x_t) - x_t) / (1.0 - _time_view(t, x_t))
|
|
|
|
sample = flow.sample(exact_field, tuple(action.shape), action.device,
|
|
generator=torch.Generator().manual_seed(5))
|
|
torch.testing.assert_close(sample, action, atol=1e-5, rtol=1e-5)
|