File size: 660 Bytes
4fd79a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | import torch
from spectral_world_models.models import build_model
def test_spectral_norm_constraint_is_enforced():
m=build_model('swm_spectral_norm')
tr=m.transition
ids=torch.arange(tr.action_dim)
A,_=tr.constrained_matrix(ids)
sigma=torch.linalg.matrix_norm(A,ord=2,dim=(-2,-1))
assert float(sigma.max()) <= tr.rho + 1e-5
def test_spectral_norm_variant_forward_shape():
m=build_model('swm_spectral_norm')
image=torch.rand(2,1,32,32); tokens=torch.ones(2,5,dtype=torch.long); action=torch.tensor([0,1])
out=m(image,tokens,action)
assert out['image_pred'].shape == image.shape
assert out['z_pred'].shape == (2,96)
|