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

Code before:
    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)

    np.testing.assert_allclose(
        R_np, R_recon, rtol=1e-5, atol=1e-6, err_msg=f"NumPy round-trip failed for batch size {batch_size}"
    )


@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().numpy()
    np.testing.assert_allclose(
        pos.cpu().numpy(),
        transformed,
        rtol=1e-5,
        atol=1e-6,
        err_msg=f"NumPy identity_transform failed for batch size {batch_size}",
    )
    transformed = transform_by_T(b_pos, T_identity).cpu().numpy()
    np.testing.assert_allclose(
        b_pos.cpu().numpy(),
        transformed,
        rtol=1e-5,
        atol=1e-6,
        err_msg=f"NumPy identity_transform failed for batch size {batch_size}",
    )

    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().numpy()
        np.testing.assert_allclose(
            pos.cpu().numpy(),
            transformed,
            rtol=1e-5,
            atol=1e-6,
            err_msg=f"NumPy identity_transform failed for batch size {batch_size}",
        )

Review comment (suggestion):
See my other review:
* No need to do `.numpy()` anywhere
* rather use our own `assert_allclose`
* `tolerance` should be a global constant after the imports.
* You need to jutify your tolerance threshold, why 1e-5 ? It is actually quite bad. And I possible, use `tol` instead of `atol` + `rtol`

Code after:
    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)
        transformed = transform_by_T(b_pos, T_batched_identity).cpu()
        assert_allclose(b_pos.cpu(), transformed, tol=TOL)


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