tests/test_stability_constraint.py

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)