File size: 737 Bytes
4fd79a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | import numpy as np, torch
from spectral_world_models.v9_causal_generalization import make_trajectory_cf_batch,trajectory_cf_loss,INTERP,EXTRAP,STRUCTURAL
from spectral_world_models.models import build_model
def test_v9_regimes_are_distinct():
assert INTERP.acceleration < EXTRAP.acceleration
assert STRUCTURAL.restitution < INTERP.restitution
def test_trajectory_cf_batch_and_backward():
b=make_trajectory_cf_batch(n=2,horizon=3,seed=1,device='cpu')
assert b['left'].shape==(2,3,1,32,32)
m=build_model('swm_structured_cf_v9')
loss,parts=trajectory_cf_loss(m,b)
assert torch.isfinite(loss)
loss.backward()
assert any(p.grad is not None for p in m.parameters())
assert 'traj_cf_dir_loss' in parts
|