Repository: Genesis-Embodied-AI/Genesis
File: tests/test_utils.py (Python)

Code before:


@pytest.mark.parametrize("batch_size", [1, 10, 100])
def test_torch_round_trip(batch_size):
    print(f"Testing torch implementation with batch size {batch_size}...")
    rot = R.random(batch_size)
    R_np = rot.as_matrix()
    if batch_size == 1:
        R_np = R_np[0]
    R_torch = torch.tensor(R_np, dtype=torch.float32)

    d6 = R_to_rot6d(R_torch)
    R_recon = rot6d_to_R(d6)

    assert_allclose(R_torch.cpu(), R_recon.cpu(), tol=TOL)


@pytest.mark.parametrize("batch_size", [1, 10, 100])
def test_numpy_round_trip(batch_size):
    print(f"Testing numpy implementation with batch size {batch_size}...")
    rot = R.random(batch_size)
    R_np = rot.as_matrix()
    if batch_size == 1:
        R_np = R_np[0]

    d6 = R_to_rot6d(R_np)
    R_recon = rot6d_to_R(d6)

    assert_allclose(R_np, R_recon, tol=TOL)


@pytest.mark.parametrize("batch_size", [0, 1, 100])
def test_torch_identity_transform(batch_size):
    if batch_size == 0:
        pos = torch.randn(3)
        b_pos = torch.randn(10, 3)
    else:
        pos = torch.randn(batch_size, 3)
        b_pos = torch.randn(batch_size, 10, 3)

    T_identity = torch.eye(4)

    transformed = transform_by_T(pos, T_identity).cpu()
    assert_allclose(pos.cpu(), transformed, tol=TOL)
    transformed = transform_by_T(b_pos, T_identity).cpu()
    assert_allclose(b_pos.cpu(), transformed, tol=TOL)

    if batch_size > 0:
        T_batched_identity = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)
        transformed = transform_by_T(pos, T_batched_identity).cpu()
        assert_allclose(pos.cpu(), transformed, tol=TOL)

Review comment (bug):
A bit of a shame to loose the debugging error messages. You could have extended `assert_allclose` to support optional argument `err_msg`, but it is up to you. 

EDIT: Although I have to admit in this case the message was not informative it is good to remove it.

Code after:

TOL = 1e-7

pytestmark = [pytest.mark.required]


@pytest.fixture
def clear_seen_fixture():
    warnings_mod._seen.clear()
    yield
    warnings_mod._seen.clear()


@pytest.mark.parametrize("batch_size", [1, 10, 100])
def test_torch_round_trip(batch_size):
    print(f"Testing torch implementation with batch size {batch_size}...")
    rot = R.random(batch_size)
    R_np = rot.as_matrix()
    if batch_size == 1:
        R_np = R_np[0]
    R_torch = torch.tensor(R_np, dtype=torch.float32)

    d6 = R_to_rot6d(R_torch)
    R_recon = rot6d_to_R(d6)

    assert_allclose(R_torch.cpu(), R_recon.cpu(), tol=TOL)


@pytest.mark.parametrize("batch_size", [1, 10, 100])
def test_numpy_round_trip(batch_size):
    print(f"Testing numpy implementation with batch size {batch_size}...")
    rot = R.random(batch_size)
    R_np = rot.as_matrix()
    if batch_size == 1:
        R_np = R_np[0]

    d6 = R_to_rot6d(R_np)
    R_recon = rot6d_to_R(d6)

    assert_allclose(R_np, R_recon, tol=TOL)


@pytest.mark.parametrize("batch_size", [0, 1, 100])
def test_torch_identity_transform(batch_size):
    if batch_size == 0:
        pos = torch.randn(3)
        b_pos = torch.randn(10, 3)
    else:
        pos = torch.randn(batch_size, 3)
        b_pos = torch.randn(batch_size, 10, 3)